"""
Copyright ©2017. The Regents of the University of California (Regents). All Rights Reserved.
Permission to use, copy, modify, and distribute this software and its documentation for educational,
research, and not-for-profit purposes, without fee and without a signed licensing agreement, is
hereby granted, provided that the above copyright notice, this paragraph and the following two
paragraphs appear in all copies, modifications, and distributions. Contact The Office of Technology
Licensing, UC Berkeley, 2150 Shattuck Avenue, Suite 510, Berkeley, CA 94720-1620, (510) 643-
7201, otl@berkeley.edu, http://ipira.berkeley.edu/industry-info for commercial licensing opportunities.

IN NO EVENT SHALL REGENTS BE LIABLE TO ANY PARTY FOR DIRECT, INDIRECT, SPECIAL,
INCIDENTAL, OR CONSEQUENTIAL DAMAGES, INCLUDING LOST PROFITS, ARISING OUT OF
THE USE OF THIS SOFTWARE AND ITS DOCUMENTATION, EVEN IF REGENTS HAS BEEN
ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

REGENTS SPECIFICALLY DISCLAIMS ANY WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
PURPOSE. THE SOFTWARE AND ACCOMPANYING DOCUMENTATION, IF ANY, PROVIDED
HEREUNDER IS PROVIDED "AS IS". REGENTS HAS NO OBLIGATION TO PROVIDE
MAINTENANCE, SUPPORT, UPDATES, ENHANCEMENTS, OR MODIFICATIONS.

Author: Mike Danielczuk
"""

import os
import time
import numpy as np
import gym
import scipy.stats as sstats
import trimesh

from autolab_core import Logger, RigidTransform

from .random_variables import CameraRandomVariable
from .states import ObjectState, HeapState, CameraState, HeapAndCameraState
from .constants import KEY_SEP_TOKEN, TRAIN_ID, TEST_ID

class CameraStateSpace(gym.Space):
    """ State space for a camera. """
    def __init__(self, config):
        self._config = config

        # read params
        self.frame = config['name']

        # random variable for pose of camera
        self.camera_rv = CameraRandomVariable(config)

    def sample(self):
        """ Sample a camera state. """
        pose, intrinsics = self.camera_rv.sample(size=1)
        return CameraState(self.frame, pose, intrinsics)

class HeapStateSpace(gym.Space):
    """ State space for object heaps. """
    def __init__(self, physics_engine, config):

        self._physics_engine = physics_engine
        self._config = config

        # set up logger
        self._logger = Logger.get_logger(self.__class__.__name__)

        # read subconfigs
        obj_config = config['objects']
        workspace_config = config['workspace']

        self.num_objs_rv = sstats.poisson(config['mean_objs']-1)
        self.max_objs = config['max_objs']
        self.min_objs = 1
        if 'min_objs' in list(config.keys()):
            self.min_objs = config['min_objs']

        self.max_obj_diam = config['max_obj_diam']
        self.drop_height = config['drop_height']
        self.max_settle_steps = config['max_settle_steps']
        self.mag_v_thresh = config['mag_v_thresh']
        self.mag_w_thresh = config['mag_w_thresh']

        # bounds of heap center in the table plane
        min_heap_center = np.array(config['center']['min'])
        max_heap_center = np.array(config['center']['max'])
        self.heap_center_space = gym.spaces.Box(min_heap_center, max_heap_center, dtype=np.float32)

        # Set up object configs
        # bounds of object drop pose in the table plane
        # organized as [tx, ty, theta] where theta is in degrees
        min_obj_pose = np.r_[obj_config['planar_translation']['min'], 0]
        max_obj_pose = np.r_[obj_config['planar_translation']['max'], 2*np.pi]
        self.obj_planar_pose_space = gym.spaces.Box(min_obj_pose, max_obj_pose, dtype=np.float32)

        # bounds of object drop orientation
        min_sph_coords = np.array([0.0, 0.0])
        max_sph_coords = np.array([2*np.pi, np.pi])
        self.obj_orientation_space = gym.spaces.Box(min_sph_coords, max_sph_coords, dtype=np.float32)

        # bounds of center of mass
        delta_com_sigma = max(1e-6, obj_config['center_of_mass']['sigma'])
        self.delta_com_rv = sstats.multivariate_normal(np.zeros(3), delta_com_sigma**2)

        self.obj_density = obj_config['density']

        # bounds of workspace (for checking out of bounds)
        min_workspace_trans = np.array(workspace_config['min'])
        max_workspace_trans = np.array(workspace_config['max'])
        self.workspace_space = gym.spaces.Box(min_workspace_trans, max_workspace_trans, dtype=np.float32)

        # Setup object keys and directories
        object_keys = []
        mesh_filenames = []
        self._train_pct = obj_config['train_pct']
        num_objects = obj_config['num_objects']
        self._mesh_dir = obj_config['mesh_dir']
        if not os.path.isabs(self._mesh_dir):
            self._mesh_dir = os.path.join(os.getcwd(), self._mesh_dir)
        for root, dirs, files in os.walk(self._mesh_dir):
            dataset_name = os.path.basename(root)
            if dataset_name in list(obj_config['object_keys'].keys()):
                for f in files:
                    filename, ext = os.path.splitext(f)
                    if ext.split('.')[1] in trimesh.exchange.load.mesh_formats() \
                        and (filename in obj_config['object_keys'][dataset_name] or obj_config['object_keys'][dataset_name] == 'all'):
                        obj_key = '{}{}{}'.format(dataset_name,KEY_SEP_TOKEN,filename)
                        object_keys.append(obj_key)
                        mesh_filenames.append(os.path.join(root,f))

        inds = np.arange(len(object_keys))
        np.random.shuffle(inds)
        self.all_object_keys = list(np.array(object_keys)[inds][:num_objects])
        mesh_filenames = list(np.array(mesh_filenames)[inds][:num_objects])
        self.train_keys = self.all_object_keys[:int(len(self.all_object_keys)*self._train_pct)]
        self.test_keys = self.all_object_keys[int(len(self.all_object_keys)*self._train_pct):]
        self.obj_ids = dict([(key, i+1) for i,key in enumerate(self.all_object_keys)])
        self.mesh_filenames = {}
        [self.mesh_filenames.update({k:v}) for k,v in zip(self.all_object_keys, mesh_filenames)]

    @property
    def obj_keys(self):
        return self.all_object_keys

    @obj_keys.setter
    def obj_keys(self, keys):
        self.all_object_keys = keys

    @property
    def num_objects(self):
        return len(self.all_object_keys)

    @property
    def obj_id_map(self):
        return self.obj_ids

    @obj_id_map.setter
    def obj_id_map(self, id_map):
        self.obj_ids = id_map

    @property
    def obj_splits(self):
        obj_splits = {}
        for key in self.all_object_keys:
            if key in self.train_keys:
                obj_splits[key] = TRAIN_ID
            else:
                obj_splits[key] = TEST_ID
        return obj_splits

    def set_splits(self, obj_splits):
        self.train_keys = []
        self.test_keys = []
        for k in list(obj_splits.keys()):
            if obj_splits[k] == TRAIN_ID:
                self.train_keys.append(k)
            else:
                self.test_keys.append(k)

    def in_workspace(self, pose):
        """ Check whether a pose is in the workspace. """
        return self.workspace_space.contains(pose.translation)

    def sample(self):
        """ Samples a state from the space
        Returns
        -------
        :obj:`HeapState`
            state of the object pile
        """

        # Start physics engine
        self._physics_engine.start()

        # setup workspace
        workspace_obj_states = []
        workspace_objs = self._config['workspace']['objects']
        for work_key, work_config in list(workspace_objs.items()):

            # make paths absolute
            mesh_filename = work_config['mesh_filename']
            pose_filename = work_config['pose_filename']

            if not os.path.isabs(mesh_filename):
                mesh_filename = os.path.join(os.path.dirname(os.path.realpath(__file__)),
                                             '..',
                                             mesh_filename)
            if not os.path.isabs(pose_filename):
                pose_filename = os.path.join(os.path.dirname(os.path.realpath(__file__)),
                                             '..',
                                             pose_filename)

            # load mesh
            mesh = trimesh.load_mesh(mesh_filename)
            mesh.density = self.obj_density
            pose = RigidTransform.load(pose_filename)
            workspace_obj = ObjectState(work_key, mesh, pose)
            self._physics_engine.add(workspace_obj, static=True)
            workspace_obj_states.append(workspace_obj)

        # sample state
        train = True
        if np.random.rand() > self._train_pct:
            train = False
            sample_keys = self.test_keys
            self._logger.info('Sampling from test')
        else:
            sample_keys = self.train_keys
            self._logger.info('Sampling from train')

        total_num_objs = len(sample_keys)

        # sample object ids
        num_objs = min(self.num_objs_rv.rvs(size=1)[0], total_num_objs-1) + 1
        num_objs = min(num_objs, self.max_objs)
        num_objs = max(num_objs, self.min_objs)
        obj_inds = np.random.permutation(np.arange(total_num_objs))

        # log
        self._logger.info('Sampling %d objects' %(num_objs))

        # sample pile center
        heap_center = self.heap_center_space.sample()
        t_heap_world = np.array([heap_center[0], heap_center[1], 0])
        self._logger.debug('Sampled pile location: %.3f %.3f' %(t_heap_world[0], t_heap_world[1]))

        # sample object, center of mass, pose
        objs_in_heap = []
        total_drops = 0
        while total_drops < total_num_objs and len(objs_in_heap) < num_objs:
            obj_key = sample_keys[total_drops]
            obj_mesh = trimesh.load_mesh(self.mesh_filenames[obj_key])
            obj_mesh.density = self.obj_density
            obj = ObjectState(obj_key, obj_mesh)
            _, radius = trimesh.nsphere.minimum_nsphere(obj.mesh)
            if 2*radius > self.max_obj_diam:
                self._logger.warning('Obj too big, skipping .....')
                total_drops += 1
                continue

            # sample center of mass
            delta_com = self.delta_com_rv.rvs(size=1)
            center_of_mass = obj.mesh.center_mass + delta_com
            obj.mesh.center_mass = center_of_mass

            # sample obj drop pose
            obj_orientation = self.obj_orientation_space.sample()
            az = obj_orientation[0]
            elev = obj_orientation[1]
            T_obj_table = RigidTransform.sph_coords_to_pose(az, elev).as_frames('obj', 'world')

            # sample object planar pose
            obj_planar_pose = self.obj_planar_pose_space.sample()
            theta = obj_planar_pose[2]
            R_table_world = RigidTransform.z_axis_rotation(theta)
            R_obj_drop_world = R_table_world.dot(T_obj_table.rotation)
            t_obj_drop_heap = np.array([obj_planar_pose[0],
                                        obj_planar_pose[1],
                                        self.drop_height])
            t_obj_drop_world = t_obj_drop_heap + t_heap_world
            obj.pose = RigidTransform(rotation=R_obj_drop_world,
                                      translation=t_obj_drop_world,
                                      from_frame='obj',
                                      to_frame='world')

            self._physics_engine.add(obj)
            try:
                v, w = self._physics_engine.get_velocity(obj.key)
            except:
                self._logger.warning('Could not get base velocity for object %s. Skipping ...' %(obj_key))
                self._physics_engine.remove(obj.key)
                total_drops += 1
                continue

            objs_in_heap.append(obj)

            # setup until approx zero velocity
            wait = time.time()
            objects_in_motion = True
            num_steps = 0
            while objects_in_motion and num_steps < self.max_settle_steps:

                # step simulation
                self._physics_engine.step()

                # check velocities
                max_mag_v = 0
                max_mag_w = 0
                objs_to_remove = set()
                for o in objs_in_heap:
                    try:
                        v, w = self._physics_engine.get_velocity(o.key)
                    except:
                        self._logger.warning('Could not get base velocity for object %s. Skipping ...' %(o.key))
                        objs_to_remove.add(o)
                        continue
                    mag_v = np.linalg.norm(v)
                    mag_w = np.linalg.norm(w)
                    if mag_v > max_mag_v:
                        max_mag_v = mag_v
                    if mag_w > max_mag_w:
                        max_mag_w = mag_w

                # Remove invalid objects
                for o in objs_to_remove:
                    self._physics_engine.remove(o.key)
                    objs_in_heap.remove(o)

                # check objs in motion
                if max_mag_v < self.mag_v_thresh and max_mag_w < self.mag_w_thresh:
                    objects_in_motion = False

                num_steps += 1

            # read out object poses
            objs_to_remove = set()
            for o in objs_in_heap:
                obj_pose = self._physics_engine.get_pose(o.key)
                o.pose = obj_pose.copy()

                # remove the object if its outside of the workspace
                if not self.in_workspace(obj_pose):
                    self._logger.warning('Object {} fell out of the workspace!'.format(o.key))
                    objs_to_remove.add(o)

            # remove invalid objects
            for o in objs_to_remove:
                self._physics_engine.remove(o.key)
                objs_in_heap.remove(o)

            total_drops += 1
            self._logger.debug('Waiting for zero velocity took %.3f sec' %(time.time()-wait))

        # Stop physics engine
        self._physics_engine.stop()

        # add metadata for heap state and return it
        metadata = {'split': TRAIN_ID}
        if not train:
            metadata['split'] = TEST_ID

        return HeapState(workspace_obj_states, objs_in_heap, metadata=metadata)

    def _add_object_to_simulation(self, obj, useFixedBase=0):

        # create URDF
        urdf_dir = os.path.join(self._urdf_cache_dir, obj.key)
        name = os.path.basename(urdf_dir)
        name = '{}.urdf'.format(name)
        urdf_filename = os.path.join(urdf_dir, name)
        self.urdf_filenames[obj.key] = urdf_filename
        obj.mesh.density = self.obj_density
        # Mess with center of mass
        geometry = obj.mesh.copy()
        geometry.apply_translation(-obj.mesh.center_mass)
        if not os.path.exists(urdf_filename):
            try:
                os.makedirs(urdf_dir)
            except:
                self._logger.warning('Failed to create dir %s. The object may have been created simultaneously by another process' %(urdf_dir))
            self._logger.info('Exporting URDF for object %s' %(obj.key))
            trimesh.exchange.export.export_urdf(geometry, urdf_dir)

        # add to simulator
        obj_id = self._physics_engine.add_urdf(self.urdf_filenames[obj.key], obj.pose, obj.mesh.center_mass, useFixedBase=useFixedBase)
        return obj_id

class HeapAndCameraStateSpace(gym.Space):
    """ State space for environments. """
    def __init__(self, physics_engine, config):

        heap_config = config['heap']
        cam_config = config['camera']

        # individual state spaces
        self.heap = HeapStateSpace(physics_engine, heap_config)
        self.camera = CameraStateSpace(cam_config)

    @property
    def obj_id_map(self):
        return self.heap.obj_id_map

    @obj_id_map.setter
    def obj_id_map(self, id_map):
        self.heap.obj_ids = id_map

    @property
    def obj_keys(self):
        return self.heap.obj_keys

    @obj_keys.setter
    def obj_keys(self, keys):
        self.heap.all_object_keys = keys

    @property
    def obj_splits(self):
        return self.heap.obj_splits

    def set_splits(self, splits):
        self.heap.set_splits(splits)

    @property
    def mesh_filenames(self):
        return self.heap.mesh_filenames

    @mesh_filenames.setter
    def mesh_filenames(self, fns):
        self.heap.mesh_filenames = fns

    def sample(self):
        """ Sample a state. """
        # sample individual states
        heap_state = self.heap.sample()
        cam_state = self.camera.sample()

        return HeapAndCameraState(heap_state, cam_state)


from __future__ import print_function, division, absolute_import

import atexit
import collections
import functools
import io
import itertools
import os
import threading
import timeit
import math
import sys
import traceback
import weakref
from types import ModuleType
import numpy as np

from .six import *
from .errors import UnsupportedError
try:
    # preferred over pure-python StringIO due to threadsafety
    # note: parallel write to StringIO could cause data to go missing
    from cStringIO import StringIO
except ImportError:
    from io import StringIO
from numba.config import PYVERSION, MACHINE_BITS, DEVELOPER_MODE

IS_PY3 = PYVERSION >= (3, 0)

if IS_PY3:
    import builtins
    INT_TYPES = (int,)
    longint = int
    get_ident = threading.get_ident
    intern = sys.intern
    file_replace = os.replace
    asbyteint = int

    def erase_traceback(exc_value):
        """
        Erase the traceback and hanging locals from the given exception instance.
        """
        if exc_value.__traceback__ is not None:
            traceback.clear_frames(exc_value.__traceback__)
        return exc_value.with_traceback(None)

else:
    import thread
    import __builtin__ as builtins
    INT_TYPES = (int, long)
    longint = long
    get_ident = thread.get_ident
    intern = intern
    def asbyteint(x):
        # convert 1-char str into int
        return ord(x)

    if sys.platform == 'win32':
        def file_replace(src, dest):
            # Best-effort emulation of os.replace()
            try:
                os.rename(src, dest)
            except OSError:
                os.unlink(dest)
                os.rename(src, dest)
    else:
        file_replace = os.rename

    def erase_traceback(exc_value):
        """
        Erase the traceback and hanging locals from the given exception instance.
        """
        return exc_value

try:
    from inspect import signature as pysignature
    from inspect import Signature as pySignature
    from inspect import Parameter as pyParameter
except ImportError:
    try:
        from funcsigs import signature as _pysignature
        from funcsigs import Signature as pySignature
        from funcsigs import Parameter as pyParameter
        from funcsigs import BoundArguments

        def pysignature(*args, **kwargs):
            try:
                return _pysignature(*args, **kwargs)
            except ValueError as e:
                msg = ("Cannot obtain a signature for: %s. The error message "
                       "from funcsigs was: '%s'." % (args,  e.message))
                raise UnsupportedError(msg)

        # monkey patch `apply_defaults` onto `BoundArguments` cf inspect in py3
        # This patch is from https://github.com/aliles/funcsigs/pull/30/files
        # with minor modifications, and thanks!
        # See LICENSES.third-party.
        def apply_defaults(self):
            arguments = self.arguments

            # Creating a new one and not modifying in-place for thread safety.
            new_arguments = []

            for name, param in self._signature.parameters.items():
                try:
                    new_arguments.append((name, arguments[name]))
                except KeyError:
                    if param.default is not param.empty:
                        val = param.default
                    elif param.kind is _VAR_POSITIONAL:
                        val = ()
                    elif param.kind is _VAR_KEYWORD:
                        val = {}
                    else:
                        # BoundArguments was likely created by bind_partial
                        continue
                    new_arguments.append((name, val))

            self.arguments = collections.OrderedDict(new_arguments)
        BoundArguments.apply_defaults = apply_defaults
    except ImportError:
        raise ImportError("please install the 'funcsigs' package "
                          "('pip install funcsigs')")

try:
    from functools import singledispatch
except ImportError:
    try:
        import singledispatch
    except ImportError:
        raise ImportError("please install the 'singledispatch' package "
                          "('pip install singledispatch')")
    else:
        # Hotfix for https://bitbucket.org/ambv/singledispatch/issues/8/inconsistent-hierarchy-with-enum
        def _c3_merge(sequences):
            """Merges MROs in *sequences* to a single MRO using the C3 algorithm.

            Adapted from http://www.python.org/download/releases/2.3/mro/.

            """
            result = []
            while True:
                sequences = [s for s in sequences if s]   # purge empty sequences
                if not sequences:
                    return result
                for s1 in sequences:   # find merge candidates among seq heads
                    candidate = s1[0]
                    for s2 in sequences:
                        if candidate in s2[1:]:
                            candidate = None
                            break      # reject the current head, it appears later
                    else:
                        break
                if candidate is None:
                    raise RuntimeError("Inconsistent hierarchy")
                result.append(candidate)
                # remove the chosen candidate
                for seq in sequences:
                    if seq[0] == candidate:
                        del seq[0]

        singledispatch._c3_merge = _c3_merge

        singledispatch = singledispatch.singledispatch

# Mapping between operator module functions and the corresponding built-in
# operators.

BINOPS_TO_OPERATORS = {
    '+': operator.add,
    '-': operator.sub,
    '*': operator.mul,
    '//': operator.floordiv,
    '/': operator.truediv,
    '%': operator.mod,
    '**': operator.pow,
    '&': operator.and_,
    '|': operator.or_,
    '^': operator.xor,
    '<<': operator.lshift,
    '>>': operator.rshift,
    '==': operator.eq,
    '!=': operator.ne,
    '<': operator.lt,
    '<=': operator.le,
    '>': operator.gt,
    '>=': operator.ge,
    'is': operator.is_,
    'is not': operator.is_not,
    # This one has its args reversed!
    'in': operator.contains
}

INPLACE_BINOPS_TO_OPERATORS = {
    '+=': operator.iadd,
    '-=': operator.isub,
    '*=': operator.imul,
    '//=': operator.ifloordiv,
    '/=': operator.itruediv,
    '%=': operator.imod,
    '**=': operator.ipow,
    '&=': operator.iand,
    '|=': operator.ior,
    '^=': operator.ixor,
    '<<=': operator.ilshift,
    '>>=': operator.irshift,
}

UNARY_BUITINS_TO_OPERATORS = {
    '+': operator.pos,
    '-': operator.neg,
    '~': operator.invert,
    'not': operator.not_,
    'is_true': operator.truth
}

OPERATORS_TO_BUILTINS = {
    operator.add: '+',
    operator.iadd: '+=',
    operator.sub: '-',
    operator.isub: '-=',
    operator.mul: '*',
    operator.imul: '*=',
    operator.floordiv: '//',
    operator.ifloordiv: '//=',
    operator.truediv: '/',
    operator.itruediv: '/=',
    operator.mod: '%',
    operator.imod: '%=',
    operator.pow: '**',
    operator.ipow: '**=',
    operator.and_: '&',
    operator.iand: '&=',
    operator.or_: '|',
    operator.ior: '|=',
    operator.xor: '^',
    operator.ixor: '^=',
    operator.lshift: '<<',
    operator.ilshift: '<<=',
    operator.rshift: '>>',
    operator.irshift: '>>=',
    operator.eq: '==',
    operator.ne: '!=',
    operator.lt: '<',
    operator.le: '<=',
    operator.gt: '>',
    operator.ge: '>=',
    operator.is_: 'is',
    operator.is_not: 'is not',
    # This one has its args reversed!
    operator.contains: 'in',
    # Unary
    operator.pos: '+',
    operator.neg: '-',
    operator.invert: '~',
    operator.not_: 'not',
    operator.truth: 'is_true',
}

HAS_MATMUL_OPERATOR = sys.version_info >= (3, 5)

if not IS_PY3:
    BINOPS_TO_OPERATORS['/?'] = operator.div
    INPLACE_BINOPS_TO_OPERATORS['/?='] = operator.idiv
    OPERATORS_TO_BUILTINS[operator.div] = '/?'
    OPERATORS_TO_BUILTINS[operator.idiv] = '/?'
if HAS_MATMUL_OPERATOR:
    BINOPS_TO_OPERATORS['@'] = operator.matmul
    INPLACE_BINOPS_TO_OPERATORS['@='] = operator.imatmul

_shutting_down = False

def _at_shutdown():
    global _shutting_down
    _shutting_down = True

def shutting_down(globals=globals):
    """
    Whether the interpreter is currently shutting down.
    For use in finalizers, __del__ methods, and similar; it is advised
    to early bind this function rather than look it up when calling it,
    since at shutdown module globals may be cleared.
    """
    # At shutdown, the attribute may have been cleared or set to None.
    v = globals().get('_shutting_down')
    return v is True or v is None

# weakref.finalize registers an exit function that runs all finalizers for
# which atexit is True. Some of these finalizers may call shutting_down() to
# check whether the interpreter is shutting down. For this to behave correctly,
# we need to make sure that _at_shutdown is called before the finalizer exit
# function. Since atexit operates as a LIFO stack, we first contruct a dummy
# finalizer then register atexit to ensure this ordering.
weakref.finalize(lambda: None, lambda: None)
atexit.register(_at_shutdown)

class ConfigOptions(object):
    OPTIONS = {}

    def __init__(self):
        self._values = self.OPTIONS.copy()

    def set(self, name, value=True):
        if name not in self.OPTIONS:
            raise NameError("Invalid flag: %s" % name)
        self._values[name] = value

    def unset(self, name):
        self.set(name, False)

    def _check_attr(self, name):
        if name not in self.OPTIONS:
            raise AttributeError("Invalid flag: %s" % name)

    def __getattr__(self, name):
        self._check_attr(name)
        return self._values[name]

    def __setattr__(self, name, value):
        if name.startswith('_'):
            super(ConfigOptions, self).__setattr__(name, value)
        else:
            self._check_attr(name)
            self._values[name] = value

    def __repr__(self):
        return "Flags(%s)" % ', '.join('%s=%s' % (k, v)
                                       for k, v in self._values.items()
                                       if v is not False)

    def copy(self):
        copy = type(self)()
        copy._values = self._values.copy()
        return copy

    def __eq__(self, other):
        return isinstance(other, ConfigOptions) and other._values == self._values

    def __ne__(self, other):
        return not self == other

    def __hash__(self):
        return hash(tuple(sorted(self._values.items())))

class SortedMap(Mapping):
    """Immutable
    """

    def __init__(self, seq):
        self._values = []
        self._index = {}
        for i, (k, v) in enumerate(sorted(seq)):
            self._index[k] = i
            self._values.append((k, v))

    def __getitem__(self, k):
        i = self._index[k]
        return self._values[i][1]

    def __len__(self):
        return len(self._values)

    def __iter__(self):
        return iter(k for k, v in self._values)

class UniqueDict(dict):
    def __setitem__(self, key, value):
        if key in self:
            raise AssertionError("key already in dictionary: %r" % (key,))
        super(UniqueDict, self).__setitem__(key, value)

# Django's cached_property
# see https://docs.djangoproject.com/en/dev/ref/utils/#django.utils.functional.cached_property

class cached_property(object):
    """
    Decorator that converts a method with a single self argument into a
    property cached on the instance.

    Optional ``name`` argument allows you to make cached properties of other
    methods. (e.g.  url = cached_property(get_absolute_url, name='url') )
    """
    def __init__(self, func, name=None):
        self.func = func
        self.name = name or func.__name__

    def __get__(self, instance, type=None):
        if instance is None:
            return self
        res = instance.__dict__[self.name] = self.func(instance)
        return res

def runonce(fn):
    @functools.wraps(fn)
    def inner():
        if not inner._ran:
            res = fn()
            inner._result = res
            inner._ran = True
        return inner._result

    inner._ran = False
    return inner

def bit_length(intval):
    """
    Return the number of bits necessary to represent integer `intval`.
    """
    assert isinstance(intval, INT_TYPES)
    if intval >= 0:
        return len(bin(intval)) - 2
    else:
        return len(bin(-intval - 1)) - 2

def stream_list(lst):
    """
    Given a list, return an infinite iterator of iterators.
    Each iterator iterates over the list from the last seen point up to
    the current end-of-list.

    In effect, each iterator will give the newly appended elements from the
    previous iterator instantiation time.
    """
    def sublist_iterator(start, stop):
        return iter(lst[start:stop])

    start = 0
    while True:
        stop = len(lst)
        yield sublist_iterator(start, stop)
        start = stop

class BenchmarkResult(object):
    def __init__(self, func, records, loop):
        self.func = func
        self.loop = loop
        self.records = np.array(records) / loop
        self.best = np.min(self.records)

    def __repr__(self):
        name = getattr(self.func, "__name__", self.func)
        args = (name, self.loop, self.records.size, format_time(self.best))
        return "%20s: %10d loops, best of %d: %s per loop" % args

def format_time(tm):
    units = "s ms us ns ps".split()
    base = 1
    for unit in units[:-1]:
        if tm >= base:
            break
        base /= 1000
    else:
        unit = units[-1]
    return "%.1f%s" % (tm / base, unit)

def benchmark(func, maxsec=1):
    timer = timeit.Timer(func)
    number = 1
    result = timer.repeat(1, number)
    # Too fast to be measured
    while min(result) / number == 0:
        number *= 10
        result = timer.repeat(3, number)
    best = min(result) / number
    if best >= maxsec:
        return BenchmarkResult(func, result, number)
        # Scale it up to make it close the maximum time
    max_per_run_time = maxsec / 3 / number
    number = max(max_per_run_time / best / 3, 1)
    # Round to the next power of 10
    number = int(10 ** math.ceil(math.log10(number)))
    records = timer.repeat(3, number)
    return BenchmarkResult(func, records, number)

RANGE_ITER_OBJECTS = (builtins.range,)
if PYVERSION < (3, 0):
    RANGE_ITER_OBJECTS += (builtins.xrange,)
    try:
        from future.types.newrange import newrange
        RANGE_ITER_OBJECTS += (newrange,)
    except ImportError:
        pass

def logger_hasHandlers(logger):
    # Backport from python3.5 logging implementation of `.hasHandlers()`
    c = logger
    rv = False
    while c:
        if c.handlers:
            rv = True
            break
        if not c.propagate:
            break
        else:
            c = c.parent
    return rv

# A dummy module for dynamically-generated functions
_dynamic_modname = '<dynamic>'
_dynamic_module = ModuleType(_dynamic_modname)
_dynamic_module.__builtins__ = moves.builtins

def chain_exception(new_exc, old_exc):
    """Set the __cause__ attribute on *new_exc* for explicit exception
    chaining.  Returns the inplace modified *new_exc*.
    """
    if DEVELOPER_MODE:
        new_exc.__cause__ = old_exc
    return new_exc


from speech_analysis_raspi.support import saving_loading as sl
from speech_analysis_raspi.support import configurations_variables as confv
from speech_analysis_raspi.support import configurations_methods as confm
from speech_analysis_raspi.support import directory_file_checking as dfc
import sys
import tflite_runtime.interpreter as tflite
from speech_analysis_raspi.support import calculations as calc
from speech_analysis_raspi.support import configuration_classes as confc
from speech_analysis_raspi.support import  recording_configurations as rconf
from python_speech_features import mfcc
import numpy as np

def load_all_helpers_dict(database_dict, gender):
    # print(database_dict)

    dataset_helpers_dict = {}
    for db, mode in database_dict.items():
        # print(db, mode)
        dataset_helpers_dict[db] = get_helper(database=db, gender=gender, mode=mode)

    return dataset_helpers_dict

def get_helper(database, gender, mode):
    mconf = sl.load_model_config(database=database, gender=gender, mode=mode)
    # print(mconf.database)
    # print(mconf.gender)
    # print(mconf.mode)
    # print(mconf.nfilt)
    # print(mconf.nfeat)
    # print(mconf.nfft)
    # print(mconf.step)
    # print(mconf.classes)
    # print(mconf.features_save_name)
    # print(mconf.model_config_save_name)
    # print(mconf.training_log_name)
    # print(mconf.model_save_name)
    # print(mconf.model_h5_save_name)
    # print(mconf.model_tflite_save_name)
    # print(mconf.feature_path)
    # print(mconf.model_config_path)
    # print(mconf.training_log_path)
    # print(mconf.model_path)
    # print(mconf.model_h5_path)
    # print(mconf.model_tflite_path)

    model_tflite_path = confm.get_universal_path(mconf.model_tflite_path)
    if not dfc.check_file_existence(model_tflite_path):
        print("Saved model does not exist for {database} - {gender}".format(database=database, gender=gender))
        sys.exit()

    # TFLite - Load the TFLite model and allocate tensors.
    interpreter = tflite.Interpreter(model_path=model_tflite_path)

    # Get validation accuracy
    val_acc = calc.get_validation_accuracy(confm.get_universal_path(mconf.training_log_path))

    # Creating dataset_helper to hold respective ModelConfig and interpreter objects with validation accuracies
    dataset_helper = confc.DatasetHelper(modelconfig=mconf, interpreter=interpreter, val_acc=val_acc)

    return dataset_helper

def predict_real_time(dataset_helpers_dict, signal):
    dict_act_map_list = []
    for db, helper in dataset_helpers_dict.items():
        # print(db, helper, helper.modelconfig, helper.interpreter, helper.val_acc)
        dict_act_map = predict(helper=helper, signal=signal)
        # print('dict_act_map', dict_act_map)
        # print('dict_act_map.dict', dict_act_map.dict, 'dict_act_map.val_acc', dict_act_map.val_acc)
        dict_act_map_list.append(dict_act_map)
        # print(dict_act_map_list)

    final_result = calc.get_biased_probs(dict_act_map_list)
    # print('final_result', final_result)

    return final_result

def predict(helper: confc.DatasetHelper, signal):
    # TFLite - Load the TFLite model and allocate tensors.
    interpreter = helper.interpreter
    interpreter.allocate_tensors()

    # TFLite - Get input and output tensors.
    input_details = interpreter.get_input_details()
    output_details = interpreter.get_output_details()

    mconf = helper.modelconfig
    y_pred = []
    y_prob = []
    count = 0
    for i in range(0, signal.shape[0] - mconf.step, mconf.step):
        count = count + 1   # Which 10th of a second number
        sample = signal[i:i + mconf.step]
        x = mfcc(sample, rconf.RATE_realtime, numcep=mconf.nfeat, nfilt=mconf.nfilt, nfft=mconf.nfft)
        x = (x - mconf.min) / (mconf.max - mconf.min)

        x = x.reshape(1, x.shape[0], x.shape[1], 1)

        # TFLite - convert array to relevant datatype noted in input_details[0]['dtype']
        x = np.array(x, dtype=np.float32)
        # TFLite - set the value of the input tensor
        interpreter.set_tensor(input_details[0]['index'], x)
        # TFLite - invoke the interpreter
        interpreter.invoke()
        # TFLite - get the value of the input tensor
        y_hat = interpreter.get_tensor(output_details[0]['index'])

        y_prob.append(y_hat)

    fn_prob = np.mean(y_prob, axis=0).flatten()
    # print("The mean of all probabilities of", count, "1/10 second chunks from the single audio file for the classes ", mconf.classes, " are ", fn_prob)

    results_dict = calc.get_dict(mconf.classes, fn_prob)
    # print('results_dict', results_dict)

    dict_act_map = confc.DictToValAccMap(dict=results_dict, val_acc=helper.val_acc)
    # print('dict_act_map', dict_act_map)
    # print('dict_act_map.dict', dict_act_map.dict, 'dict_act_map.val_acc', dict_act_map.val_acc)

    return dict_act_map


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Script to plot CPU usage from profiling data.

This script requires the matplotlib and numpy Python modules.
"""

import argparse
import glob
import os
import sys

import numpy

from matplotlib import pyplot

def Main():
  """The main program function.

  Returns:
    bool: True if successful or False if not.
  """
  argument_parser = argparse.ArgumentParser(description=(
      'Plots CPU usage from profiling data.'))

  argument_parser.add_argument(
      '--output', dest='output_file', type=str, help=(
          'path of the output file to write the graph to instead of using '
          'interactive mode. The output format deduced from the extension '
          'of the filename.'))

  argument_parser.add_argument(
      '--process', dest='process', type=str, default='', help=(
          'comma separated list of names of processes to graph.'))

  argument_parser.add_argument(
      '--profiler', dest='profiler', type=str, default='processing', help=(
          'name of the CPU profiler to graph.'))

  argument_parser.add_argument(
      'profile_path', type=str, help=(
          'path to the directory containing the profiling data.'))

  options = argument_parser.parse_args()

  if not os.path.isdir(options.profile_path):
    print('No such directory: {0:s}'.format(options.profile_path))
    return False

  if options.profiler not in ('analyzers', 'parsers', 'processing'):
    print('Unsupported profiler: {0:s}'.format(options.profiler))
    return False

  processes = []
  if options.process:
    processes = options.process.split(',')

  if options.profiler == 'analyzers':
    name_prefix = 'analyzers'
    name_suffix = 'analyzers'

  elif options.profiler == 'parsers':
    name_prefix = 'cputime'
    name_suffix = 'parsers'

  elif options.profiler == 'processing':
    name_prefix = 'processing'
    name_suffix = 'processing'

  names = ['time', 'name', 'cpu']

  glob_pattern = '{0:s}-*-{1:s}.csv.gz'.format(name_prefix, name_suffix)
  glob_expression = os.path.join(options.profile_path, glob_pattern)
  for csv_file_name in glob.glob(glob_expression):
    process_name = os.path.basename(csv_file_name)
    process_name_prefix = '{0:s}-'.format(name_prefix)
    process_name_suffix = '-{0:s}.csv.gz'.format(name_suffix)
    process_name = process_name.replace(process_name_prefix, '').replace(
        process_name_suffix, '')
    if processes and process_name not in processes:
      continue

    data = numpy.genfromtxt(
        csv_file_name, delimiter='\t', dtype=None, encoding='utf-8',
        names=names, skip_header=1)

    if data.size > 0:
      for name in numpy.unique(data['name']):
        # Ignore process_sources since it is a single sample that contains
        # the cumulative CPU time.
        if options.profiler == 'processing' and name == 'process_sources':
          continue

        data_by_name = numpy.extract(data['name'] == name, data)
        label = '-'.join([name, process_name])
        pyplot.plot(data_by_name['time'], data_by_name['cpu'], label=label)

  pyplot.title('CPU usage over time')

  pyplot.xlabel('Time')
  pyplot.xscale('linear')

  pyplot.ylabel('Used CPU')
  pyplot.yscale('linear')

  pyplot.legend()

  if options.output_file:
    pyplot.savefig(options.output_file)
  else:
    pyplot.show()

  return True

if __name__ == '__main__':
  if not Main():
    sys.exit(1)
  else:
    sys.exit(0)


#!/usr/bin/env python
from __future__ import print_function, division

import argparse, os, sys
import re
import numpy as np
import cv2 as cv
import pytesseract
import platform
import itertools
from tqdm import tqdm

from config import config
from datetime import datetime

import engine

def say(message):
    sys.stdout.write(message + "\n")

def ask(message):
    try:
        return input(message)
    except KeyboardInterrupt:
        say("\nBye.")
        sys.exit(0)

def print_board(board, active, clear_screen=True, spymaster=True):

    if clear_screen:
        if platform.system() == "Windows":
            os.system("cls")
        else:
            sys.stdout.write(chr(27) + "[2J")

    for row in range(5):
        for col in range(5):
            index = row * 5 + col
            word = board[index]
            tag = " "
            if not active[word]:
                word = word + "-" * (11 - len(word))
                tag = "-"
            sys.stdout.write("{0}{1:11s} ".format(tag, word))
        sys.stdout.write("\n")

# TODO: this needs to be piped into engine, not copied and pasted

def play_computer_spymaster(
    engine,
    player_words,
    opponent_words,
    neutral_words,
    assassin_word,
    given_clues,
    gamma=1.0,
    verbose=True,
):

    say("Thinking...")
    sys.stdout.flush()

    # # Loop over all permutations of words.
    # num_words = len(player_words)
    # best_score, saved_clues = [], []
    # counts = range(num_words, 0, -1)
    # groups_and_count = []
    # for count in counts:
    #     for group in itertools.combinations(range(num_words), count):
    #         groups_and_count.append((group, count,))
    # for group, count in tqdm(groups_and_count):
    #     # Multiply similarity scores by this factor for any clue
    #     # corresponding to this many words.
    #     bonus_factor = count ** gamma
    #     # print(type(player_words), player_words)
    #     words = player_words[list(group)]
    #     clue, score = engine.model.get_clue(
    #         clue_words=words,
    #         pos_words=player_words,
    #         neg_words=np.concatenate((opponent_words, neutral_words)),
    #         veto_words=assassin_word,
    #     )
    #     if clue:
    #         best_score.append(score * bonus_factor)
    #         saved_clues.append((clue, words))
    # num_clues = len(saved_clues)
    # order = sorted(range(num_clues), key=lambda k: best_score[k], reverse=True)

    max_group = 2
    does_stretch = [2]
    num_words = len(player_words)
    best_score, saved_clues = [], []
    counts = range(min(num_words, max_group), 0, -1)
    groups_and_count = []
    for count in counts:
        for group in itertools.combinations(range(num_words), count):
            groups_and_count.append((group, count,))
    groups_and_count = tqdm(groups_and_count)
    for group, count in groups_and_count:
        # Multiply similarity scores by this factor for any clue
        # corresponding to this many words.
        bonus_factor = count ** gamma
        words = player_words[list(group)]
        clue, score, stretch = engine.model.get_clue_stretch(
            clue_words=words,
            pos_words=player_words,
            neg_words=opponent_words,
            neut_words=neutral_words,
            veto_words=assassin_word,
            given_clues=given_clues,
            give_stretch=(count in does_stretch),
        )
        if clue:
            best_score.append(score * bonus_factor)
            clue_words = words
            if stretch:
                clue_words = np.concatenate((words, np.asarray(stretch)))
            saved_clues.append((clue, clue_words))

    num_clues = len(saved_clues)
    order = sorted(range(num_clues), key=lambda k: best_score[k], reverse=True)

    if not os.path.exists(config.logs_folder):
        os.makedirs(config.logs_folder)
    with open(engine.logs_filename, 'a+') as f:
        for i in order[:10]:
            clue, words = saved_clues[i]
            f.write(
                u"{0:.3f} {2} {3} = {1}\n".format(
                    best_score[i],
                    " + ".join([str(w).upper()[2:-1] for w in words]),
                    str(clue)[2:-1],
                    len(words),
                )
            )
        f.write("\n")

    # print_board(board, active, spymaster=True)
    values = []
    for i in order[:10]:
        values.append(
            [saved_clues[i][0], saved_clues[i][1], best_score[i],]
        )
    return values

colors = [
    (92.5, 66.7, 66.7),
    (65.9, 84.7, 92.2),
    (97.3, 99.2, 100.0),
    (40.8, 42.4, 42.7),
]
colors = list([np.uint8(c * 2.55) for c in color[::-1]] for color in colors)
color_names = ["red", "blue", "white", "black"]

def main():
    parser = argparse.ArgumentParser(
        description="Play the CodeNames game.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument(
        "--nohide",
        type=bool,
        default=False,
        help="Shows only 1 clue, hides what the words for each clue",
    )
    args = parser.parse_args()

    e = engine.GameEngine()

    # manual init
    pytesseract.pytesseract.tesseract_cmd = (
        r"C:\\Program Files\\Tesseract-OCR\\tesseract.exe"
    )
    config = "-l eng --oem 1 --psm 3"
    e.size = 5
    pad = 20
    img = cv.imread("board.png")
    colors_loc = [[], [], [], []]
    word_color_map = {}
    board = []
    cx, cy = img.shape[1] // 5, img.shape[0] // 5
    for y in range(5):
        for x in range(5):
            img_sec = img[
                y * cy + pad : (y + 1) * cy - pad, x * cx + pad : (x + 1) * cx - pad
            ]
            img_gray = cv.cvtColor(img_sec, cv.COLOR_BGR2GRAY)
            avg = np.average(img_gray)
            mask1 = img_gray[:, :] < avg
            mask2 = img_gray[:, :] > avg
            img_gray[mask2] = 0
            img_gray[mask1] = 255
            kernel = np.ones((3, 3), np.uint8)
            erosion = cv.erode(img_gray, kernel, iterations=1)
            text = pytesseract.image_to_string(img_sec, config=config).replace(" ", "_")
            board.append(text)
            dev = [0, 0, 0, 0]
            for c in range(3):
                avg = np.median(img_sec[:, :, c])
                for i, color in enumerate(colors):
                    # print(avg, color[c])
                    dev[i] += (avg - color[c]) ** 2
            index, smol = 0, dev[0]
            if not args.nohide:
                print("({},{}) -> {}".format(x + 1, y + 1, text))
            else:
                print("({},{}) -> {} = {}".format(x + 1, y + 1, text, dev))
            for i in range(1, 4):
                if dev[i] < smol:
                    smol = dev[i]
                    index = i
            colors_loc[index].append((x, y,))
            word_color_map[text] = color_names[index]
            # cv.imshow('img', erosion)
            # cv.imshow('img_pre', img_sec)
            # cv.waitKey()

    # play the game
    active = {}
    for word in board:
        active[word] = True
    input("Loaded successfully! Hit enter to continue: ")
    print_board(board, active)
    values = None
    index = 0
    side = ""
    given_words = []
    while True:
        print_board(board, active)
        if values:
            if not args.nohide:
                clue, words, best_score = values[index]
                say(u"{3}{0:.3f} {1} {2}".format(best_score, str(clue)[2:-1], len(words), side))
            else:
                for clue, words, best_score in values:
                    say(
                        u"{0:.3f} {2} {3} = {1}".format(
                            best_score,
                            " + ".join([str(w).upper()[2:-1] for w in words]),
                            str(clue)[2:-1],
                            len(words),
                        )
                    )
        text = input().upper()
        if text in ("EXIT", "QUIT"):
            return
        if text in board:
            active[text] = not active[text]
        if text in ("USE", "USED"):
            given_words.append(clue)
        elif text == "RED" or text == "BLUE":
            red_words = np.asarray(
                [
                    word.lower().encode("utf8")
                    for word in board
                    if active[word] and word_color_map[word] == "red"
                ]
            )
            blue_words = np.asarray(
                [
                    word.lower().encode("utf8")
                    for word in board
                    if active[word] and word_color_map[word] == "blue"
                ]
            )
            neutral_words = np.asarray(
                [
                    word.lower().encode("utf8")
                    for word in board
                    if active[word] and word_color_map[word] == "white"
                ]
            )
            assassin_word = [
                word.lower().encode("utf8")
                for word in board
                if active[word] and word_color_map[word] == "black"
            ]
            if text == "RED":
                values = play_computer_spymaster(
                    e, red_words, blue_words, neutral_words, assassin_word, given_words
                )
                side = "red: "
            else:
                values = play_computer_spymaster(
                    e, blue_words, red_words, neutral_words, assassin_word, given_words
                )
                side = "blue: "
            index = 0
        elif text == "NEXT":
            index += 1
            if index >= len(values):
                index -= len(values)
        elif text == "PREV":
            index -= 1
            if index < 0:
                index += len(values)

if __name__ == "__main__":
    main()


import torch
import torch.nn as nn
import torch.nn.functional as F
import math
import numpy as np
import cv2
import time
from ..loss import build_loss, ohem_batch, iou
from ..post_processing import pa

from tqdm import tqdm

from multiprocessing.pool import ThreadPool
from multiprocessing import Process, Queue

from scipy.sparse import csr_matrix

def compute_M(data):
    cols = np.arange(data.size)
    return csr_matrix((cols, (data.ravel(), cols)),
                      shape=(data.max()+1, data.size))

def get_indices_sparse(data):
    M = compute_M(data)
    return [np.unravel_index(row.data, data.shape) for row in M[1:]]

class PA_Head(nn.Module):
    def __init__(self,
                 in_channels,
                 hidden_dim,
                 num_classes,
                 loss_text,
                 loss_kernel,
                 loss_emb):
        super(PA_Head, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, hidden_dim, kernel_size=3, stride=1, padding=1)
        self.bn1 = nn.BatchNorm2d(hidden_dim)
        self.relu1 = nn.ReLU(inplace=True)

        self.conv2 = nn.Conv2d(hidden_dim, num_classes, kernel_size=1, stride=1, padding=0)

        self.text_loss = build_loss(loss_text)
        self.kernel_loss = build_loss(loss_kernel)
        self.emb_loss = build_loss(loss_emb)

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()

    def forward(self, f):
        out = self.conv1(f)
        out = self.relu1(self.bn1(out))
        out = self.conv2(out)

        return out

    def get_results(self, out, img_meta, cfg):
        outputs = dict()

        if not self.training and cfg.report_speed:
            torch.cuda.synchronize()
            start = time.time()

        score = torch.sigmoid(out[:, 0, :, :])
        kernels = out[:, :2, :, :] > 0
        text_mask = kernels[:, :1, :, :]
        kernels[:, 1:, :, :] = kernels[:, 1:, :, :] * text_mask
        emb = out[:, 2:, :, :]
        emb = emb * text_mask.float()

        score = score.data.cpu().numpy()[0].astype(np.float32)
        kernels = kernels.data.cpu().numpy()[0].astype(np.uint8)
        emb = emb.cpu().numpy()[0].astype(np.float32)

        # pa
        label = pa(kernels, emb)

        # image size
        org_img_size = img_meta['org_img_size'][0]
        img_size = img_meta['img_size'][0]

        label_num = np.max(label) + 1
        label = cv2.resize(label, (int(img_size[1]), int(img_size[0])), interpolation=cv2.INTER_NEAREST)
        score = cv2.resize(score, (int(img_size[1]), int(img_size[0])), interpolation=cv2.INTER_NEAREST)

        if not self.training and cfg.report_speed:
            torch.cuda.synchronize()
            outputs.update(dict(
                det_pa_time=time.time() - start
            ))

        scale = (float(org_img_size[1]) / float(img_size[1]),
                 float(org_img_size[0]) / float(img_size[0]))

        with_rec = hasattr(cfg.model, 'recognition_head')

        if with_rec:
            bboxes_h = np.zeros((1, label_num, 4), dtype=np.int32)
            instances = [[]]

        bboxes = []
        scores = []
        for i in range(1, label_num):
            ind = label == i
            points = np.array(np.where(ind)).transpose((1, 0))

            if points.shape[0] < cfg.test_cfg.min_area:
                label[ind] = 0
                continue

            score_i = np.mean(score[ind])
            if score_i < cfg.test_cfg.min_score:
                label[ind] = 0
                continue

            if with_rec:
                tl = np.min(points, axis=0)
                br = np.max(points, axis=0) + 1
                bboxes_h[0, i] = (tl[0], tl[1], br[0], br[1])
                instances[0].append(i)

            if cfg.test_cfg.bbox_type == 'rect':
                rect = cv2.minAreaRect(points[:, ::-1])
                bbox = cv2.boxPoints(rect) * scale
            elif cfg.test_cfg.bbox_type == 'poly':
                binary = np.zeros(label.shape, dtype='uint8')
                binary[ind] = 1
                contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
                bbox = contours[0] * scale

            bbox = bbox.astype('int32')
            bboxes.append(bbox.reshape(-1))
            scores.append(score_i)

        outputs.update(dict(
            bboxes=bboxes,
            scores=scores
        ))
        if with_rec:
            outputs.update(dict(
                label=label,
                bboxes_h=bboxes_h,
                instances=instances
            ))

        return outputs

    def loss(self, out, gt_texts, gt_kernels, training_masks, gt_instances, gt_bboxes):
        # output
        texts = out[:, 0, :, :]
        kernels = out[:, 1:2, :, :]
        embs = out[:, 2:, :, :]

        # text loss
        selected_masks = ohem_batch(texts, gt_texts, training_masks)
        loss_text = self.text_loss(texts, gt_texts, selected_masks, reduce=False)
        iou_text = iou((texts > 0).long(), gt_texts, training_masks, reduce=False)
        losses = dict(
            loss_text=loss_text,
            iou_text=iou_text
        )

        # kernel loss
        loss_kernels = []
        selected_masks = gt_texts * training_masks
        for i in range(kernels.size(1)):
            kernel_i = kernels[:, i, :, :]
            gt_kernel_i = gt_kernels[:, i, :, :]
            loss_kernel_i = self.kernel_loss(kernel_i, gt_kernel_i, selected_masks, reduce=False)
            loss_kernels.append(loss_kernel_i)
        loss_kernels = torch.mean(torch.stack(loss_kernels, dim=1), dim=1)
        iou_kernel = iou(
            (kernels[:, -1, :, :] > 0).long(), gt_kernels[:, -1, :, :], training_masks * gt_texts, reduce=False)
        losses.update(dict(
            loss_kernels=loss_kernels,
            iou_kernel=iou_kernel
        ))

        # embedding loss
        loss_emb = self.emb_loss(embs, gt_instances, gt_kernels[:, -1, :, :], training_masks, gt_bboxes, reduce=False)
        losses.update(dict(
            loss_emb=loss_emb
        ))

        return losses

    def sigmoid(self, x):
        return 1 / (1 + np.exp(-x))

    def get_single_result(self, out, img_meta, cfg):
        time0 = time.time()

        outputs = dict()

        if not self.training and cfg.report_speed:
            torch.cuda.synchronize()
            start = time.time()

        temp = out.cpu().numpy()
        score = self.sigmoid(temp[:,0,:,:])
        kernels = temp[:, :2, :, :] > 0
        text_mask = kernels[:, :1, :, :]
        kernels[:, 1:, :, :] = kernels[:, 1:, :, :] * text_mask
        emb = temp[:, 2:, :, :]
        emb = emb * text_mask.astype(float)

        score = score[0].astype(np.float32)
        kernels = kernels[0].astype(np.uint8)
        emb = emb[0].astype(np.float32)

        cv2.imwrite('kernel.jpg', kernels[1,:,:]*128)

        # score = torch.sigmoid(out[:, 0, :, :])
        # kernels = out[:, :2, :, :] > 0
        # text_mask = kernels[:, :1, :, :]
        # kernels[:, 1:, :, :] = kernels[:, 1:, :, :] * text_mask
        # emb = out[:, 2:, :, :]
        # emb = emb * text_mask.float()

        # score = score.data.cpu().numpy()[0].astype(np.float32)
        # kernels = kernels.data.cpu().numpy()[0].astype(np.uint8)
        # emb = emb.cpu().numpy()[0].astype(np.float32)

        # print(np.max(score - score1))
        # print(np.max(kernels - kernels1))
        # print(np.max(emb - emb1))
        time1 = time.time()
        # print(time1-time0)
        # pa
        label = pa(kernels, emb)
        time2 = time.time()
        # print(time2-time1)
        # image size
        org_img_size = img_meta['org_img_size']
        img_size = img_meta['img_size']

        print('org_size', org_img_size)
        print('img_size', img_size)

        label_num = np.max(label) + 1
        label = cv2.resize(label, (int(img_size[1]), int(img_size[0])), interpolation=cv2.INTER_NEAREST)
        score = cv2.resize(score, (int(img_size[1]), int(img_size[0])), interpolation=cv2.INTER_NEAREST)

        label1 = label.copy()

        if not self.training and cfg.report_speed:
            torch.cuda.synchronize()
            outputs.update(dict(
                det_pa_time=time.time() - start
            ))

        scale = (float(org_img_size[1]) / float(img_size[1]),
                 float(org_img_size[0]) / float(img_size[0]))
        print('scale', scale)

        with_rec = hasattr(cfg.model, 'recognition_head')

        if with_rec:
            bboxes_h = np.zeros((1, label_num, 4), dtype=np.int32)
            instances = [[]]

        # bboxes = []
        # scores = []

        # for i in range(1, label_num):
        #     ind = label == i
        #     points = np.array(np.where(ind)).transpose((1, 0))

        #     if points.shape[0] < cfg.test_cfg.min_area:
        #         label[ind] = 0
        #         continue

        #     score_i = np.mean(score[ind])
        #     if score_i < cfg.test_cfg.min_score:
        #         label[ind] = 0
        #         continue

        #     if with_rec:
        #         tl = np.min(points, axis=0)
        #         br = np.max(points, axis=0) + 1
        #         bboxes_h[0, i] = (tl[0], tl[1], br[0], br[1])
        #         instances[0].append(i)

        #     if cfg.test_cfg.bbox_type == 'rect':
        #         rect = cv2.minAreaRect(points[:, ::-1])
        #         bbox = cv2.boxPoints(rect) * scale
        #     elif cfg.test_cfg.bbox_type == 'poly':
        #         binary = np.zeros(label.shape, dtype='uint8')
        #         binary[ind] = 1
        #         contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        #         bbox = contours[0] * scale

        #     bbox = bbox.astype('int32')
        #     bboxes.append(bbox.reshape(-1))
        #     scores.append(score_i)

        # bboxes0 = bboxes.copy()
        # scores0 = scores.copy()
        time3 = time.time()
        print('for loop time', time3-time2)

        # label = label1.copy()
        # label[label > 0] = 1

        # cv2.imwrite('label.jpg', label*128)
        # cv2.imwrite('label1.jpg', label1)
        # label_num = np.max(label) + 1
        # print(label_num)

        ###################### multi thread ########################
        # bboxes = []
        # scores = []

        # # num_thread = label_num -1
        # pool = ThreadPool(processes=label_num-1)

        # result = []
        # for i in range(1, label_num):
        #     async_result = pool.apply_async(thread_worker, (label, score, scale, cfg.test_cfg.min_area, cfg.test_cfg.min_score, i)) # tuple of args for foo
        #     result.append(async_result)

        # for res in result:
        #     box = res.get()[0]
        #     if np.any(box != None):
        #         bboxes.append(box)
        #         scores.append(res.get()[1])

        # print('num word', len(bboxes))

        # time4 = time.time()
        # print('for loop time',time4-time3)

        ######################### sparse matrix #############################
        # bboxes = []
        # scores = []

        # list_points = get_indices_sparse(label)

        # for i in range(1, label_num):
        #     # ind = label == i
        #     t51 = time.time()
        #     # points = np.array(np.where(ind)).transpose((1, 0))
        #     points = np.array(list_points[i-1]).transpose((1, 0))
        #     ind = list_points[i-1]
        #     t6 = time.time()
        #     if points.shape[0] < cfg.test_cfg.min_area:
        #         label[ind] = 0
        #         continue
        #     t7 = time.time()
        #     score_i = np.mean(score[ind])
        #     t8 = time.time()
        #     if score_i < cfg.test_cfg.min_score:
        #         label[ind] = 0
        #         continue
        #     t9 = time.time()
        #     if cfg.test_cfg.bbox_type == 'rect':
        #         rect = cv2.minAreaRect(points[:, ::-1])
        #         bbox = cv2.boxPoints(rect) * scale
        #     elif cfg.test_cfg.bbox_type == 'poly':
        #         binary = np.zeros(label.shape, dtype='uint8')
        #         binary[ind] = 1
        #         contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        #         bbox = contours[0] * scale

        #     bbox = bbox.astype('int32')
        #     bboxes.append(bbox.reshape(-1))
        #     scores.append(score_i)

        # time4 = time.time()
        # print('for loop time',time4-time3)

        # print('num word', len(bboxes))

        ####################### merge code ########################
        bboxes = []
        scores = []

        list_points = get_indices_sparse(label)

        def thread_worker_inner(i):
            points = np.array(list_points[i-1]).transpose((1, 0))
            ind = list_points[i-1]

            if points.shape[0] < cfg.test_cfg.min_area:
                label[ind] = 0
                return (None, None)

            score_i = np.mean(score[ind])
            if score_i < cfg.test_cfg.min_score:
                label[ind] = 0
                return (None, None)

            rect = cv2.minAreaRect(points[:, ::-1])
            bbox = cv2.boxPoints(rect) * scale

            bbox = bbox.astype('int32')
            return (bbox.reshape(-1), score_i)

        # num_thread = label_num -1
        pool = ThreadPool(processes=label_num-1)

        result = []
        for i in range(1, label_num):
            # async_result = pool.apply_async(thread_worker0, (list_points, score, scale, cfg.test_cfg.min_area, cfg.test_cfg.min_score, i))
            async_result = pool.apply_async(thread_worker_inner, (i,))
            result.append(async_result)

        for res in result:
            box = res.get()[0]
            if np.any(box != None):
                bboxes.append(box)
                scores.append(res.get()[1])

        print('num word', len(bboxes))

        time4 = time.time()
        print('for loop time',time4-time3)

        ################# multi process #######################
        # bboxes = []
        # scores = []

        # list_points = get_indices_sparse(label)

        # q = Queue()

        # result = []
        # for i in range(1, label_num):
        #     p = Process(target=process_worker, args=(list_points, score, scale, cfg.test_cfg.min_area, cfg.test_cfg.min_score, i, q, )) # tuple of args for foo
        #     p.start()

        # for i in result:
        #     box, score = q.get()
        #     if np.any(box != None):
        #         bboxes.append(box)
        #         scores.append(score)

        # print('num word', len(bboxes))

        # time4 = time.time()
        # print('for loop time',time4-time3)

        outputs.update(dict(
            bboxes=bboxes,
            scores=scores
        ))
        if with_rec:
            outputs.update(dict(
                label=label,
                bboxes_h=bboxes_h,
                instances=instances
            ))

        return outputs

def thread_worker(label, score, scale, min_area, min_score, i):
    ind = label == i
    points = np.array(np.where(ind)).transpose((1, 0))

    if points.shape[0] < min_area:
        label[ind] = 0
        return (None, None)

    score_i = np.mean(score[ind])
    if score_i < min_score:
        label[ind] = 0
        return (None, None)

    rect = cv2.minAreaRect(points[:, ::-1])
    bbox = cv2.boxPoints(rect) * scale

    bbox = bbox.astype('int32')
    return (bbox.reshape(-1), score_i)

def thread_worker0(list_points, score, scale, min_area, min_score, i):
    points = np.array(list_points[i-1]).transpose((1, 0))
    ind = list_points[i-1]

    if points.shape[0] < min_area:
        label[ind] = 0
        return (None, None)

    score_i = np.mean(score[ind])
    if score_i < min_score:
        label[ind] = 0
        return (None, None)

    rect = cv2.minAreaRect(points[:, ::-1])
    bbox = cv2.boxPoints(rect) * scale

    bbox = bbox.astype('int32')
    return (bbox.reshape(-1), score_i)

def process_worker(list_points, score, scale, min_area, min_score, i, q):
    points = np.array(list_points[i-1]).transpose((1, 0))
    ind = list_points[i-1]

    if points.shape[0] < min_area:
        label[ind] = 0
        return (None, None)

    score_i = np.mean(score[ind])
    if score_i < min_score:
        label[ind] = 0
        return (None, None)

    rect = cv2.minAreaRect(points[:, ::-1])
    bbox = cv2.boxPoints(rect) * scale

    bbox = bbox.astype('int32')
    # return (bbox.reshape(-1), score_i)
    q.put((bbox.reshape(-1), score_i))


import numpy as np
import librosa
import math
import sys

print("Loading file")
audio, sample_rate = librosa.load(sys.argv[1], duration=60, offset=0, sr=15360)

print("Getting spectrum")
spectrum = librosa.stft(audio)
S = np.abs(spectrum)

fout = open("spectrum.h", "w")

print("Writing file")
fn = 36
fs = int(len(S) / fn)
fout.write("const uint16_t spectrum[][4] = {\n")
for t in range(0,len(S[0]-1)):
  fout.write("{ ")
  f_prev = 0
  for f in [8, 45, 300, 600]:
    v = 0
    for i in range(f_prev, f): v += S[i][t]
    if v != 0: v = int(v/30)
    if v < 0:  v = 0
    f_prev = f
    fout.write(str(int(v)) + ", ")
  fout.write("},\n")
fout.write("};\n")
fout.close()

print("Finished")


from __future__ import print_function
import abc
import types
import os
import random
from functools import wraps, partial
from itertools import chain
from collections import Container
from six import add_metaclass
import warnings
import math

from gbdxtools.ipe.io import to_geotiff
from gbdxtools.ipe.util import RatPolyTransform, AffineTransform, pad_safe_positive, pad_safe_negative, IPE_TO_DTYPE, preview

from shapely import ops, wkt
from shapely.geometry import box, shape, mapping
from shapely.geometry.base import BaseGeometry
try:
    from rio_hist.match import histogram_match
    has_rio = True
except ImportError:
    has_rio = False
import mercantile

import skimage.transform as tf

import pyproj
import dask
from dask import sharedict, optimize
from dask.delayed import delayed
import dask.array as da
from dask.base import is_dask_collection
from dask.rewrite import RewriteRule, RuleSet
import numpy as np

import operator
from toolz import valmap

from affine import Affine

try:
    from matplotlib import pyplot as plt
    has_pyplot = True
except:
    has_pyplot = False

try:
    xrange
except NameError:
    xrange = range

threads = int(os.environ.get('GBDX_THREADS', 8))
threaded_get = partial(dask.threaded.get, num_workers=threads)

@add_metaclass(abc.ABCMeta)
class DaskMeta(object):
    """
    A DaskMeta is an interface for the required attributes for initializing a dask Array
    """
    @abc.abstractproperty
    def dask(self):
        pass

    @abc.abstractproperty
    def name(self):
        pass

    @abc.abstractproperty
    def chunks(self):
        pass

    @abc.abstractproperty
    def dtype(self):
        pass

    @abc.abstractproperty
    def shape(self):
        pass

    def infect(self, target):
        assert isinstance(target, da.Array), "DaskMeta can only be attached to Dask Arrays"
        assert len(target.shape) in [2, 3], "target must be a dask array with 2 or 3 dimensions"
        target.__dict__["__daskmeta__"] = property(lambda s: self, DaskImage.__set_daskmeta__)
        return target

@add_metaclass(abc.ABCMeta)
class DaskImage(da.Array):
    """
    A DaskImage is a 2 or 3 dimension dask array that contains implements the `__daskmeta__` interface.
    """
    @property
    def __daskmeta__(self):
        """ Should return a DaskMeta """
        pass

    @__daskmeta__.setter
    def __set_daskmeta__(self, obj):
        self.__dict__["__daskmeta__"] = property(lambda s: obj, self.__set_daskmeta__)

    @classmethod
    def __subclasshook__(cls, C):
        if cls is DaskImage:
            try:
                if(is_dask_collection(C) and any("__daskmeta__" in B.__dict__ for B in C.__mro__)):
                    return True
            except AttributeError:
                pass
        return NotImplemented

    @staticmethod
    def _cull(dsk, keys=None):
        if not keys:
            return dsk
        elif "token" not in dsk:
            return dsk
        if "token" not in keys:
            keys.append("token")
        dsk1, _ = optimize.cull(dsk, keys)
        return dsk1

    def __getattribute__(self, name):
        fn = object.__getattribute__(self, name)
        if(isinstance(fn, types.MethodType) and
           any(name in C.__dict__ for C in self.__class__.__mro__)):
            @wraps(fn)
            def wrapped(*args, **kwargs):
                result = fn(*args, **kwargs)
                if isinstance(result, da.Array) and len(result.shape) in [2,3]:
                    dsk = self._cull(result.dask, result.__dask_keys__())
                    copy = super(DaskImage, self.__class__).__new__(self.__class__,
                                                                    dsk, result.name, result.chunks,
                                                                    result.dtype, result.shape)
                    copy.__dict__.update(self.__dict__)
                    try:
                        copy.__dict__.update(result.__dict__)
                    except AttributeError:
                        # this means result was an object with __slots__
                        pass
                    return copy
                return result
            return wrapped
        else:
            return fn

    @classmethod
    def create(cls, dm, **kwargs):
        """
        Given a dask meta object, construct a dask array, attach dask meta object.
        """
        assert isinstance(dm, DaskMeta), "argument must be an instance of a DaskMeta subclass"
        with dask.set_options(array_plugins=[dm.infect]):

            return da.Array.__new__(cls, dm.dask, dm.name, dm.chunks, dm.dtype, dm.shape)

    def read(self, bands=None, **kwargs):
        """
        Reads data from a dask array and returns the computed ndarray matching the given bands

        kwargs:
            bands (list): band indices to read from the image. Returns bands in the order specified in the list of bands.

        Returns:
            array (ndarray): a numpy array of image data
        """
        arr = self
        if bands is not None:
            arr = self[bands, ...]
        return arr.compute(get=threaded_get)

    def randwindow(self, window_shape):
        """
        Get a random window of a given shape from withing an image

        kwargs:
            window_shape (tuple): The desired shape of the returned image as (height, width) in pixels.

        Returns:
            image (dask): a new image object of the specified shape
        """
        row = random.randrange(window_shape[0], self.shape[1])
        col = random.randrange(window_shape[1], self.shape[2])
        return self[:, row-window_shape[0]:row, col-window_shape[1]:col]

    def iterwindows(self, count=64, window_shape=(256, 256)):
        """
        Iterate over random windows of an image

        kwargs:
            count (int): the number of the windows to generate. Defaults to 64, if `None` with continue to iterate over random windows until stopped.
            window_shape (tuple): The desired shape of each image as (height, width) in pixels.

        Returns:
            windows (generator): a generator of windows of the given shape
        """
        if count is None:
            while True:
                yield self.randwindow(window_shape)
        else:
            for i in xrange(count):
                yield self.randwindow(window_shape)

@add_metaclass(abc.ABCMeta)
class GeoImage(Container):
    _default_proj = "EPSG:4326"

    # def __geo_interface__(self):
    #     pass

    # def __geo_transform__(self):
    #     pass

    @classmethod
    def __subclasshook__(cls, C):
        # Must be a numpy-like, with __geo_transform__, and __geo_interface__
        if(issubclass(C, DaskImage) or issubclass(C, np.ndarray) and
           any("__geo_transform__" in B.__dict__ for B in C.__mro__) and
           any("__geo_interface__" in B.__dict__ for B in C.__mro__)):
            return True
        return NotImplemented

    @property
    def affine(self):
        """ The geo transform of the image

        Returns:
            affine (dict): The image's affine transform
        """
        # TODO add check for Ratpoly or whatevs
        return self.__geo_transform__._affine

    @property
    def bounds(self):
        """ Access the spatial bounding box of the image

        Returns:
            bounds (list): list of bounds in image projected coordinates (minx, miny, maxx, maxy)
        """
        return shape(self).bounds

    @property
    def proj(self):
        """ The projection of the image """
        return self.__geo_transform__.proj

    def aoi(self, **kwargs):
        """ Subsets the Image by the given bounds

        kwargs:
            bbox: optional. A bounding box array [minx, miny, maxx, maxy]
            wkt: optional. A WKT geometry string
            geojson: optional. A GeoJSON geometry dictionary

        Returns:
            image (ndarray): an image instance
        """
        g = self._parse_geoms(**kwargs)
        if g is None:
            return self
        else:
            return self[g]

    def geotiff(self, **kwargs):
        """ Creates a geotiff on the filesystem

        kwargs:
            path (str): optional. The path to save the geotiff to.
            bands (list): optional. A list of band indices to save to the output geotiff ([4,2,1])
            dtype (str): optional. The data type to assign the geotiff to ("float32", "uint16", etc)
            proj (str): optional. An EPSG proj string to project the image data into ("EPSG:32612")

        Returns:
            path (str): the path to created geotiff
        """
        if 'proj' not in kwargs:
            kwargs['proj'] = self.proj
        return to_geotiff(self, **kwargs)

    def preview(self, **kwargs):
        preview(self, **kwargs)

    def warp(self, dem=None, proj="EPSG:4326", **kwargs):
        """
        Delayed warp across an entire AOI or Image
        creates a new dask image by deferring calls to the warp_geometry on chunks

        kwargs:
            dem (ndarray): optional. A DEM for warping to specific elevation planes
            proj (str): optional. An EPSG proj string to project the image data into ("EPSG:32612")

        Returns:
            image (dask): a warped image as deferred image array (a dask)
        """
        try:
            img_md = self.ipe.metadata["image"]
            x_size = img_md["tileXSize"]
            y_size = img_md["tileYSize"]
        except (AttributeError, KeyError):
            x_size = kwargs.get("chunk_size", 256)
            y_size = kwargs.get("chunk_size", 256)

        # Create an affine transform to convert between real-world and pixels
        if self.proj is None:
            from_proj = "EPSG:4326"
        else:
            from_proj = self.proj

        try:
            # NOTE: this only works on images that have IPE rpcs metadata
            center = wkt.loads(self.ipe.metadata["image"]["imageBoundsWGS84"]).centroid
            g = box(*(center.buffer(self.ipe.metadata["rpcs"]["gsd"] / 2).bounds))
            # print "Input GSD (deg):", self.ipe.metadata["rpcs"]["gsd"]
            tfm = partial(pyproj.transform, pyproj.Proj(init="EPSG:4326"), pyproj.Proj(init=proj))
            gsd = kwargs.get("gsd", ops.transform(tfm, g).area ** 0.5)
            current_bounds = wkt.loads(self.ipe.metadata["image"]["imageBoundsWGS84"]).bounds
        except (AttributeError, KeyError, TypeError):
            tfm = partial(pyproj.transform, pyproj.Proj(init=self.proj), pyproj.Proj(init=proj))
            gsd = kwargs.get("gsd", (ops.transform(tfm, shape(self)).area / (self.shape[1] * self.shape[2])) ** 0.5 )
            current_bounds = self.bounds

        tfm = partial(pyproj.transform, pyproj.Proj(init=from_proj), pyproj.Proj(init=proj))
        itfm = partial(pyproj.transform, pyproj.Proj(init=proj), pyproj.Proj(init=from_proj))
        output_bounds = ops.transform(tfm, box(*current_bounds)).bounds
        gtf = Affine.from_gdal(output_bounds[0], gsd, 0.0, output_bounds[3], 0.0, -1 * gsd)

        ll = ~gtf * (output_bounds[:2])
        ur = ~gtf * (output_bounds[2:])
        x_chunks = int((ur[0] - ll[0]) / x_size) + 1
        y_chunks = int((ll[1] - ur[1]) / y_size) + 1

        num_bands = self.shape[0]

        try:
            dtype = IPE_TO_DTYPE[img_md["dataType"]]
        except:
            dtype = 'uint8'

        daskmeta = {
            "dask": {},
            "chunks": (num_bands, y_size, x_size),
            "dtype": dtype,
            "name": "warp-{}".format(self.name),
            "shape": (num_bands, y_chunks * y_size, x_chunks * x_size)
        }

        def px_to_geom(xmin, ymin):
            xmax = int(xmin + x_size)
            ymax = int(ymin + y_size)
            bounds = list((gtf * (xmin, ymax)) + (gtf * (xmax, ymin)))
            return box(*bounds)

        full_bounds = box(*output_bounds)

        dasks = []
        if isinstance(dem, GeoImage):
            if dem.proj != proj:
                dem = dem.warp(proj=proj, dem=dem)
            dasks.append(dem.dask)

        for y in xrange(y_chunks):
            for x in xrange(x_chunks):
                xmin = x * x_size
                ymin = y * y_size
                geometry = px_to_geom(xmin, ymin)
                daskmeta["dask"][(daskmeta["name"], 0, y, x)] = (self._warp, geometry, gsd, dem, proj, dtype, 5)
        daskmeta["dask"], _ = optimize.cull(sharedict.merge(daskmeta["dask"], *dasks), list(daskmeta["dask"].keys()))

        result = GeoDaskWrapper(daskmeta, self)
        result.__geo_interface__ = mapping(full_bounds)
        result.__geo_transform__ = AffineTransform(gtf, proj)
        return GeoImage.__getitem__(result, box(*output_bounds))

    def _warp(self, geometry, gsd, dem, proj, dtype, buf=0):
        transpix = self._transpix(geometry, gsd, dem, proj)
        xmin, xmax, ymin, ymax = (int(max(transpix[0,:,:].min() - buf, 0)),
                                  int(min(transpix[0,:,:].max() + buf, self.shape[1])),
                                  int(max(transpix[1,:,:].min() - buf, 0)),
                                  int(min(transpix[1,:,:].max() + buf, self.shape[2])))
        transpix[0,:,:] = transpix[0,:,:] - xmin
        transpix[1,:,:] = transpix[1,:,:] - ymin
        data = self[:,xmin:xmax, ymin:ymax].compute(get=dask.get) # read(quiet=True)

        if data.shape[1]*data.shape[2] > 0:
            return np.rollaxis(np.dstack([tf.warp(data[b,:,:], transpix, preserve_range=True, order=3, mode="edge") for b in xrange(data.shape[0])]).astype(dtype), 2, 0)
        else:
            return np.zeros((data.shape[0], transpix.shape[1], transpix.shape[2]))

    def _transpix(self, geometry, gsd, dem, proj):
        xmin, ymin, xmax, ymax = geometry.bounds
        x = np.linspace(xmin, xmax, num=int((xmax-xmin)/gsd))
        y = np.linspace(ymax, ymin, num=int((ymax-ymin)/gsd))
        xv, yv = np.meshgrid(x, y, indexing='xy')

        if self.proj is None:
            from_proj = "EPSG:4326"
        else:
            from_proj = self.proj

        itfm = partial(pyproj.transform, pyproj.Proj(init=proj), pyproj.Proj(init=from_proj))

        xv, yv = itfm(xv, yv) # if that works

        if isinstance(dem, GeoImage):
            g = box(xv.min(), yv.min(), xv.max(), yv.max())
            try:
                dem = dem[g].compute(get=dask.get) # read(quiet=True)
            except AssertionError:
                dem = 0 # guessing this is indexing by a 0 width geometry.

        if isinstance(dem, np.ndarray):
            dem = tf.resize(np.squeeze(dem), xv.shape, preserve_range=True, order=1, mode="edge")

        return self.__geo_transform__.rev(xv, yv, z=dem, _type=np.float32)[::-1]

    def _parse_geoms(self, **kwargs):
        """ Finds supported geometry types, parses them and returns the bbox """
        bbox = kwargs.get('bbox', None)
        wkt_geom = kwargs.get('wkt', None)
        geojson = kwargs.get('geojson', None)
        if bbox is not None:
            g = box(*bbox)
        elif wkt_geom is not None:
            g = wkt.loads(wkt_geom)
        elif geojson is not None:
            g = shape(geojson)
        else:
            return None
        if self.proj is None:
            return g
        else:
            return self._reproject(g, from_proj=kwargs.get('from_proj', 'EPSG:4326'))

    def _reproject(self, geometry, from_proj=None, to_proj=None):
        if from_proj is None:
            from_proj = self._default_proj
        if to_proj is None:
            to_proj = self.proj if self.proj is not None else "EPSG:4326"
        tfm = partial(pyproj.transform, pyproj.Proj(init=from_proj), pyproj.Proj(init=to_proj))
        return ops.transform(tfm, geometry)

    def _slice_padded(self, _bounds):
        pads = (max(-_bounds[0], 0), max(-_bounds[1], 0),
                max(_bounds[2]-self.shape[2], 0), max(_bounds[3]-self.shape[1], 0))
        bounds = (max(_bounds[0], 0),
                  max(_bounds[1], 0),
                  max(min(_bounds[2], self.shape[2]), 0),
                  max(min(_bounds[3], self.shape[1]), 0))

        # NOTE: image is a dask array that implements daskmeta interface (via op)
        result = self[:, bounds[1]:bounds[3], bounds[0]:bounds[2]]
        if pads[0] > 0:
            dims = (result.shape[0], result.shape[1], pads[0])
            result = da.concatenate([da.zeros(dims, chunks=dims, dtype=result.dtype),
                                     result], axis=2)
        if pads[2] > 0:
            dims = (result.shape[0], result.shape[1], pads[2])
            result = da.concatenate([result,
                                     da.zeros(dims, chunks=dims, dtype=result.dtype)], axis=2)
        if pads[1] > 0:
            dims = (result.shape[0], pads[1], result.shape[2])
            result = da.concatenate([da.zeros(dims, chunks=dims, dtype=result.dtype),
                                     result], axis=1)
        if pads[3] > 0:
            dims = (result.shape[0], pads[3], result.shape[2])
            result = da.concatenate([result,
                                     da.zeros(dims, chunks=dims, dtype=result.dtype)], axis=1)

        image = super(DaskImage, self.__class__).__new__(self.__class__,
                                                         result.dask, result.name, result.chunks,
                                                         result.dtype, result.shape)

        image.__geo_transform__ = self.__geo_transform__ + (_bounds[0], _bounds[1])
        return image

    def __getitem__(self, geometry):
        g = shape(geometry)

        bounds = ops.transform(self.__geo_transform__.rev, g).bounds
        try:
            assert g in self, "Image does not contain specified geometry {} not in {}".format(g.bounds, self.bounds)
        except AssertionError as ae:
            warnings.warn(ae.args)

        image = self._slice_padded(bounds)
        image.__geo_interface__ = mapping(g)
        return image

    def __contains__(self, g):
        geometry = ops.transform(self.__geo_transform__.rev, g)
        img_bounds = box(0, 0, *self.shape[2:0:-1])
        return img_bounds.contains(geometry)

class DaskMetaWrapper(DaskMeta):
    def __init__(self, dask):
        self.da = dask

    @property
    def dask(self):
        return self.da["dask"]

    @property
    def name(self):
        return self.da["name"]

    @property
    def chunks(self):
        return self.da["chunks"]

    @property
    def dtype(self):
        return self.da["dtype"]

    @property
    def shape(self):
        return self.da["shape"]

# Mixin class that defines plotting methods and rgb/ndvi methods
# used as a mixin to provide access to the plot method on
# GeoDaskWrapper images and ipe images
class PlotMixin(object):
    @property
    def _rgb_bands(self):
        return [4, 2, 1]

    @property
    def _ndvi_bands(self):
        return [6, 4]

    def base_layer_match(self, blm=False, **kwargs):
        rgb = self.rgb(**kwargs)
        if not blm:
            return rgb
        assert has_rio, "Have you installed the rio_hist library? If not, try pip install rio_hist and reload gbdxtools."
        from gbdxtools.images.tms_image import TmsImage
        bounds = self._reproject(box(*self.bounds), from_proj=self.proj, to_proj="EPSG:4326").bounds
        tms = TmsImage(zoom=self._calc_tms_zoom(self.affine[0]), bbox=bounds, **kwargs)
        ref = np.rollaxis(tms.read(), 0, 3)
        out = np.dstack([histogram_match(rgb[:,:,idx], ref[:,:,idx].astype(np.double)/255.0)
                        for idx in xrange(rgb.shape[-1])])
        return out

    def rgb(self, **kwargs):
        data = self._read(self[kwargs.get("bands", self._rgb_bands),...], **kwargs)
        data = np.rollaxis(data.astype(np.float32), 0, 3)
        lims = np.percentile(data, kwargs.get("stretch", [2, 98]), axis=(0, 1))
        for x in xrange(len(data[0,0,:])):
            top = lims[:,x][1]
            bottom = lims[:,x][0]
            data[:,:,x] = (data[:,:,x] - bottom) / float(top - bottom)
        return np.clip(data, 0, 1)

    def ndvi(self, **kwargs):
        data = self._read(self[self._ndvi_bands,...]).astype(np.float32)
        return (data[0,:,:] - data[1,:,:]) / (data[0,:,:] + data[1,:,:])

    def plot(self, spec="rgb", **kwargs):
        if self.shape[0] == 1 or ("bands" in kwargs and len(kwargs["bands"]) == 1):
            if "cmap" in kwargs:
                cmap = kwargs["cmap"]
                del kwargs["cmap"]
            else:
                cmap = "Greys_r"
            self._plot(tfm=self._single_band, cmap=cmap, **kwargs)
        else:
            if spec == "rgb" and self._has_token(**kwargs):
                self._plot(tfm=self.base_layer_match, **kwargs)
            else:
                self._plot(tfm=getattr(self, spec), **kwargs)

    def _has_token(self, **kwargs):
        if "access_token" in kwargs or "MAPBOX_API_KEY" in os.environ:
            return True
        else:
            return False

    def _plot(self, tfm=lambda x: x, **kwargs):
        assert has_pyplot, "To plot images please install matplotlib"
        assert self.shape[1] and self.shape[-1], "No data to plot, dimensions are invalid {}".format(str(self.shape))

        f, ax1 = plt.subplots(1, figsize=(kwargs.get("w", 10), kwargs.get("h", 10)))
        ax1.axis('off')
        plt.imshow(tfm(**kwargs), interpolation='nearest', cmap=kwargs.get("cmap", None))
        plt.show(block=False)

    def _read(self, data, **kwargs):
        if hasattr(data, 'read'):
            return data.read(**kwargs)
        else:
            return data.compute(get=threaded_get)

    def _single_band(self, **kwargs):
        return self._read(self[0,:,:], **kwargs)

    def _calc_tms_zoom(self, scale):
        for z in range(15,20):
            b = mercantile.bounds(0,0,z)
            if scale > math.sqrt((b.north - b.south)*(b.east - b.west) / (256*256)):
                return z

class GeoDaskWrapper(DaskImage, GeoImage, PlotMixin):
    def __new__(cls, daskmeta, img):
        dm = DaskMetaWrapper(daskmeta)
        self = super(GeoDaskWrapper, cls).create(dm)
        self._dm = dm
        self.__geo_interface__ = img.__geo_interface__
        self.__geo_transform__ = img.__geo_transform__
        return self

    @property
    def __daskmeta__(self):
        return self._dm

    def __getitem__(self, geometry):
        if isinstance(geometry, BaseGeometry) or getattr(geometry, "__geo_interface__", None) is not None:
            image = GeoImage.__getitem__(self, geometry)
            return image
        else:
            result = DaskImage.__getitem__(self, geometry)
            image = super(GeoDaskWrapper, self.__class__).__new__(self.__class__,
                                                            result.dask, result.name, result.chunks,
                                                            result.dtype, result.shape)

            if all([isinstance(e, slice) for e in geometry]) and len(geometry) == len(self.shape):
                xmin, ymin, xmax, ymax = geometry[2].start, geometry[1].start, geometry[2].stop, geometry[1].stop
                xmin = 0 if xmin is None else xmin
                ymin = 0 if ymin is None else ymin
                xmax = self.shape[2] if xmax is None else xmax
                ymax = self.shape[1] if ymax is None else ymax

                g = ops.transform(self.__geo_transform__.fwd, box(xmin, ymin, xmax, ymax))
                image.__geo_interface__ = mapping(g)
                image.__geo_transform__ = self.__geo_transform__ + (xmin, ymin)
            else:
                image.__geo_interface__ = self.__geo_interface__
                image.__geo_transform__ = self.__geo_transform__
            return image


# import packages
from dogclassifierapp.prediction_scripts.detector_functions import face_detector, dog_detector
from dogclassifierapp.prediction_scripts.data_functions import path_to_tensor, extract_Resnet50, Resnet50_predict_breed
from keras.models import load_model
from keras.applications.resnet50 import ResNet50, preprocess_input
import pickle
import numpy as np
from os import listdir
from os.path import isfile, join
from keras import backend as K

def predict_breed(img_path, checkpoint='static/resnet50_model.h5', names_pkl='static/dog_names.pkl'):
	"""
	This function loads a trained network from a checkpoint file and uses the model to
	predict the dog breed for an input image. When the input image is of a dog, it
	returns the predicted breed in a friendly message. When the input image is of a
	human, it returns the most resembling breed in a corresponding message.

	Inputs:    			- img_path: type str, filepath to single image
						- checkpoint: type str, checkpoint filepath for trained network
						- names_pkl: type str, path to pkl file of list of dog names

	Output:				- prediction message, path to breed image (if human) or empty string (if dog)

	"""

	# clear Keras session to reset before loading model
	K.clear_session()

	# load model
	model = load_model(checkpoint)

	# load list of dog names from pickle file
	with open(names_pkl, 'rb') as f:
		dog_names = pickle.load(f)

	# return dog breed that is predicted by the model
	prediction = Resnet50_predict_breed(img_path, model, dog_names)

	# get correct article to use in front of breed name
	vowels = ['a','e','i','o','u']
	if prediction.lower()[0] in vowels:
		article = 'an'
	else:
		article = 'a'

	# return prediction with friendly message
	if face_detector(img_path) == True and dog_detector(img_path) == False:

		# get list of filenames for breed images in breeds folder
		breed_files = [f for f in listdir('static/breeds') if isfile(join('static/breeds', f))]

		# get filename that matches breed prediction
		for filename in breed_files:
			if prediction.replace(' ', '_') in filename:
				breed_img = 'static/breeds/' + filename
		# return prediction message and breed image
		return "Hi there, human! Your furry doppelganger is {} {}. I'd take that as a compliment! :)".format(article, prediction), breed_img

	elif dog_detector(img_path) == True:
		return "This doggo appears to be {} {}. What a good boye!".format(article, prediction), 'static/img/dog.jpg'

	else:
		return "Oops, the algorithm won't work with this image! Please make sure the image is either of a dog or a front-facing human face.", 'static/img/oops.png'