# coding: utf-8
# pylint: disable= too-many-lines, redefined-builtin, protected-access
"""NDArray API of mxnet."""
from __future__ import absolute_import
from __future__ import division

import ctypes
import warnings
import sys
import functools
import operator
import numpy as np
from .base import _LIB, string_types, numeric_types
from .base import c_array, mx_float, py_str, c_str, mx_real_t
from .base import mx_uint, NDArrayHandle, FunctionHandle
from .base import ctypes2buffer
from .base import check_call, ctypes2docstring
from .context import Context
from . import _ndarray_internal as _internal

# pylint: disable= no-member
_DTYPE_NP_TO_MX = {
    np.float32 : 0,
    np.float64 : 1,
    np.float16 : 2,
    np.uint8   : 3,
    np.int32   : 4
}

_DTYPE_MX_TO_NP = {
    0 : np.float32,
    1 : np.float64,
    2 : np.float16,
    3 : np.uint8,
    4 : np.int32
}
# pylint: enable= no-member

def _new_empty_handle():
    """Return a new empty handle.

    Empty handle can be used to hold result

    Returns
    -------
    a new empty ndarray handle
    """
    hdl = NDArrayHandle()
    check_call(_LIB.MXNDArrayCreateNone(ctypes.byref(hdl)))
    return hdl

def _new_alloc_handle(shape, ctx, delay_alloc, dtype=mx_real_t):
    """Return a new handle with specified shape and context.

    Empty handle is only used to hold results

    Returns
    -------
    a new empty ndarray handle
    """
    hdl = NDArrayHandle()
    check_call(_LIB.MXNDArrayCreateEx(
        c_array(mx_uint, shape),
        mx_uint(len(shape)),
        ctypes.c_int(ctx.device_typeid),
        ctypes.c_int(ctx.device_id),
        ctypes.c_int(int(delay_alloc)),
        ctypes.c_int(int(_DTYPE_NP_TO_MX[np.dtype(dtype).type])),
        ctypes.byref(hdl)))
    return hdl

def waitall():
    """Wait all async operation to finish in MXNet

    This function is used for benchmark only
    """
    check_call(_LIB.MXNDArrayWaitAll())

class NDArray(object):
    """NDArray object in mxnet.

    NDArray is basic ndarray/Tensor like data structure in mxnet.
    """
    # pylint: disable= no-member
    def __init__(self, handle, writable=True):
        """initialize a new NDArray

        Parameters
        ----------
        handle : NDArrayHandle
            NDArray handle of C API
        """
        assert isinstance(handle, NDArrayHandle)
        self.handle = handle
        self.writable = writable

    def __repr__(self):
        shape_info = 'x'.join(['%d' % x for x in self.shape])
        return '<%s %s @%s>' % (self.__class__.__name__,
                                shape_info, self.context)

    def __del__(self):
        check_call(_LIB.MXNDArrayFree(self.handle))

    def __add__(self, other):
        return add(self, other)

    def __iadd__(self, other):
        if not self.writable:
            raise ValueError('trying to add to a readonly NDArray')
        if isinstance(other, NDArray):
            return _internal._plus(self, other, out=self)
        elif isinstance(other, numeric_types):
            return _internal._plus_scalar(self, float(other), out=self)
        else:
            raise TypeError('type %s not supported' % str(type(other)))

    def __radd__(self, other):
        return self.__add__(other)

    def __sub__(self, other):
        return subtract(self, other)

    def __isub__(self, other):
        if not self.writable:
            raise ValueError('trying to subtract from a readonly NDArray')
        if isinstance(other, NDArray):
            return _internal._minus(self, other, out=self)
        elif isinstance(other, numeric_types):
            return _internal._minus_scalar(self, float(other), out=self)
        else:
            raise TypeError('type %s not supported' % str(type(other)))

    def __rsub__(self, other):
        return subtract(other, self)

    def __mul__(self, other):
        return multiply(self, other)

    def __neg__(self):
        return _internal._mul_scalar(self, -1.0)

    def __imul__(self, other):
        if not self.writable:
            raise ValueError('trying to multiply to a readonly NDArray')
        if isinstance(other, NDArray):
            return _internal._mul(self, other, out=self)
        elif isinstance(other, numeric_types):
            return _internal._mul_scalar(self, float(other), out=self)
        else:
            raise TypeError('type %s not supported' % str(type(other)))

    def __rmul__(self, other):
        return self.__mul__(other)

    def __div__(self, other):
        return divide(self, other)

    def __rdiv__(self, other):
        return divide(other, self)

    def __idiv__(self, other):
        if not self.writable:
            raise ValueError('trying to divide from a readonly NDArray')
        if isinstance(other, NDArray):
            return _internal._div(self, other, out=self)
        elif isinstance(other, numeric_types):
            return _internal._div_scalar(self, float(other), out=self)
        else:
            raise TypeError('type %s not supported' % str(type(other)))

    def __truediv__(self, other):
        return divide(self, other)

    def __rtruediv__(self, other):
        return divide(other, self)

    def __itruediv__(self, other):
        return self.__idiv__(other)

    def __pow__(self, other):
        return power(self, other)

    def __rpow__(self, other):
        return power(other, self)

    def __getstate__(self):
        this = self.__dict__.copy()
        handle = this['handle']
        if handle is not None:
            length = ctypes.c_size_t()
            cptr = ctypes.POINTER(ctypes.c_char)()
            check_call(_LIB.MXNDArraySaveRawBytes(self.handle,
                                                  ctypes.byref(length),
                                                  ctypes.byref(cptr)))
            this['handle'] = ctypes2buffer(cptr, length.value)
        return this

    def __setstate__(self, state):
        handle = state['handle']
        if handle is not None:
            buf = handle
            handle = NDArrayHandle()
            ptr = (ctypes.c_char * len(buf)).from_buffer(buf)
            length = ctypes.c_size_t(len(buf))
            check_call(_LIB.MXNDArrayLoadFromRawBytes(ptr, length, ctypes.byref(handle)))
            state['handle'] = handle
        self.__dict__.update(state)

    def __setitem__(self, in_slice, value):
        """Set ndarray value.

        `value` can be a scalar, an `NDArray` or numpy array of compatible shape.
        The following modes are supported:

        - `array[:] = value`: set all the contents
        - `array[i] = value`: set the i-th slice. If the array is of dimension
          `(d1, d2, d3)`, it sets value of a slice of shape `(1, d2, d3)`.
        - `array[i:j] = value`: similarly, if the array is of dimension
          `(d1, d2, d3)`, it sets value of a slice of shape `(j-i, d2, d3)`.

        Fully-dimensional indexing is also supported. For example, if array is
        of shape `(d1, d2, d3)`, one can do

        - `array[:, :, :] = value`: achieving the same effect of `array[:] = value`
        - `array[:, i, j:k] = value`: each index could be a python slice or an int.
        """
        # pylint: disable=too-many-branches
        if not self.writable:
            raise ValueError('trying to assign to a readonly NDArray')
        if isinstance(in_slice, int):
            sliced_arr = self._at(in_slice)
            sliced_arr[:] = value
            return
        if isinstance(in_slice, slice):
            if in_slice.step is not None:
                raise ValueError('NDArray only support continuous slicing on axis 0')
            if in_slice.start is not None or in_slice.stop is not None:
                sliced_arr = self._slice(in_slice.start, in_slice.stop)
                sliced_arr[:] = value
                return
            if isinstance(value, NDArray):
                if value.handle is not self.handle:
                    value.copyto(self)
            elif isinstance(value, numeric_types):
                _internal._set_value(float(value), out=self)
            elif isinstance(value, (np.ndarray, np.generic)):
                self._sync_copyfrom(value)
            else:
                raise TypeError('type %s not supported' % str(type(value)))
        if isinstance(in_slice, tuple):
            # multi-dimension indexing
            my_shape = self.shape
            assert len(in_slice) == len(my_shape)
            for slice_i in in_slice:
                assert isinstance(slice_i, (slice, int))
            begin = [0 for _ in my_shape]
            end = [x for x in my_shape]
            for i, slice_i in enumerate(in_slice):
                if isinstance(slice_i, int):
                    assert slice_i < my_shape[i]
                    begin[i] = slice_i
                    end[i] = slice_i + 1
                if isinstance(slice_i, slice):
                    # only support continuous slicing
                    assert slice_i.step is None
                    begin[i] = slice_i.start or 0
                    end[i] = slice_i.stop or my_shape[i]
                    assert begin[i] < end[i]
                    assert end[i] <= my_shape[i]
            begin = tuple(begin)
            end = tuple(end)
            if isinstance(value, NDArray):
                value = value.as_in_context(self.context)
                _internal._crop_assign(self, value, out=self,
                                       begin=begin, end=end)
            elif isinstance(value, numeric_types):
                _internal._crop_assign_scalar(self, out=self,
                                              begin=begin, end=end,
                                              scalar=value)
            elif isinstance(value, (np.ndarray, np.generic)):
                value = array(value, ctx=self.context)
                _internal._crop_assign(self, value, out=self,
                                       begin=begin, end=end)
            else:
                raise TypeError('type %s not supported' % str(type(value)))
        # pylint: enable=too-many-branches

    def __getitem__(self, in_slice):
        """Get ndarray"""
        if isinstance(in_slice, int):
            return self._at(in_slice)
        if not isinstance(in_slice, slice) or in_slice.step is not None:
            raise ValueError('NDArray only support continuous slicing on axis 0')
        if in_slice.start is not None or in_slice.stop is not None:
            return self._slice(in_slice.start, in_slice.stop)
        else:
            return self

    def _sync_copyfrom(self, source_array):
        """Peform an synchronize copy from the array.

        Parameters
        ----------
        source_array : array_like
            The data source we should like to copy from.
        """
        if not isinstance(source_array, np.ndarray):
            try:
                source_array = np.array(source_array, dtype=self.dtype)
            except:
                raise TypeError('array must be an array_like data,' +
                                'type %s is not supported' % str(type(array)))
        source_array = np.ascontiguousarray(source_array, dtype=self.dtype)
        if source_array.shape != self.shape:
            raise ValueError('Shape inconsistant: expected %s vs got %s'%(
                str(self.shape), str(source_array.shape)))
        check_call(_LIB.MXNDArraySyncCopyFromCPU(
            self.handle,
            source_array.ctypes.data_as(ctypes.c_void_p),
            ctypes.c_size_t(source_array.size)))

    def _slice(self, start, stop):
        """Return a sliced NDArray that shares memory with current one.

        Parameters
        ----------
        start : int
            Starting index of slice.
        stop : int
            Finishing index of slice.
        """
        handle = NDArrayHandle()
        start = mx_uint(start) if start else mx_uint(0)
        stop = mx_uint(stop) if stop else mx_uint(self.shape[0])
        check_call(_LIB.MXNDArraySlice(
            self.handle, start, stop, ctypes.byref(handle)))
        return NDArray(handle=handle, writable=self.writable)

    def _at(self, idx):
        """Return a sub NDArray that shares memory with current one.

        Parameters
        ----------
        idx : int
            index of sub array.
        """
        handle = NDArrayHandle()
        idx = mx_uint(idx)
        check_call(_LIB.MXNDArrayAt(
            self.handle, idx, ctypes.byref(handle)))
        return NDArray(handle=handle, writable=self.writable)

    def reshape(self, new_shape):
        """Return a reshaped NDArray that shares memory with current one.

        Parameters
        ----------
        new_shape : iterable of int
            new shape of NDArray
        """
        handle = NDArrayHandle()
        check_call(_LIB.MXNDArrayReshape(self.handle,
                                         len(new_shape),
                                         c_array(ctypes.c_int, new_shape),
                                         ctypes.byref(handle)))
        return NDArray(handle=handle, writable=self.writable)

    # pylint: disable= undefined-variable
    def broadcast_to(self, shape):
        """ Broadcasting the current NDArray into the given shape. The semantics is
        the same with `numpy`'s broadcasting

        Parameters
        ---------
        shape : the shape to broadcast
            the broadcast shape
        """
        cur_shape = self.shape
        err_str = 'operands could not be broadcast together with remapped shapes' \
                  '[original->remapped]: {} and requested shape {}'.format(cur_shape, shape)
        if len(shape) < len(cur_shape):
            raise ValueError(err_str)
        cur_shape = (1,) * (len(shape) - len(cur_shape)) + cur_shape
        cur_shape_arr = np.array(cur_shape)
        broadcasting_axes = np.nonzero(cur_shape_arr != np.array(shape))
        if (cur_shape_arr[broadcasting_axes] != 1).any():
            raise ValueError(err_str)
        if cur_shape != self.shape:
            return broadcast_to(self.reshape(cur_shape), shape=shape)
        else:
            return broadcast_to(self, shape=tuple(shape))
    # pylint: enable= undefined-variable

    def wait_to_read(self):
        """Block until all pending writes operations on current NDArray are finished.

        This function will return when all the pending writes to the current
        NDArray finishes. There can still be pending read going on when the
        function returns.
        """
        check_call(_LIB.MXNDArrayWaitToRead(self.handle))

    @property
    def shape(self):
        """Get shape of current NDArray.

        Returns
        -------
        a tuple representing shape of current ndarray
        """
        ndim = mx_uint()
        pdata = ctypes.POINTER(mx_uint)()
        check_call(_LIB.MXNDArrayGetShape(
            self.handle, ctypes.byref(ndim), ctypes.byref(pdata)))
        return tuple(pdata[:ndim.value])

    @property
    def size(self):
        """Get size of current NDArray.

        Returns
        -------
        an int representing size of current ndarray
        """
        return np.prod(self.shape)

    @property
    def context(self):
        """Get context of current NDArray.

        Returns
        -------
        context : mxnet.Context
            The context of current NDArray.
        """
        dev_typeid = ctypes.c_int()
        dev_id = ctypes.c_int()
        check_call(_LIB.MXNDArrayGetContext(
            self.handle, ctypes.byref(dev_typeid), ctypes.byref(dev_id)))
        return Context(Context.devtype2str[dev_typeid.value], dev_id.value)

    @property
    def dtype(self):
        """Get data type of current NDArray.

        Returns
        -------
        an numpy.dtype object representing type of current ndarray
        """
        mx_dtype = ctypes.c_int()
        check_call(_LIB.MXNDArrayGetDType(
            self.handle, ctypes.byref(mx_dtype)))
        return _DTYPE_MX_TO_NP[mx_dtype.value]

    @property
    # pylint: disable= invalid-name, undefined-variable
    def T(self):
        """Get transpose of current NDArray"""
        if len(self.shape) != 2:
            raise ValueError('Only 2D matrix is allowed to be transposed')
        return transpose(self)
    # pylint: enable= invalid-name, undefined-variable

    def asnumpy(self):
        """Return a copied numpy array of current array.

        Returns
        -------
        array : numpy.ndarray
            A copy of array content.
        """
        data = np.empty(self.shape, dtype=self.dtype)
        check_call(_LIB.MXNDArraySyncCopyToCPU(
            self.handle,
            data.ctypes.data_as(ctypes.c_void_p),
            ctypes.c_size_t(data.size)))
        return data

    def asscalar(self):
        """Return a CPU scalar(float) of current ndarray.

        This ndarray must have shape (1,)

        Returns
        -------
        scalar : np.float
            The scalar representation of the ndarray.
        """
        if self.shape != (1,):
            raise ValueError("The current array is not a scalar")
        return self.asnumpy()[0]

    def astype(self, dtype):
        """Return a copied numpy array of current array with specified type.

        Parameters
        ----------
        dtype : numpy.dtype or string
            Desired type of result array.

        Returns
        -------
        array : numpy.ndarray
            A copy of array content.
        """
        res = empty(self.shape, ctx=self.context, dtype=dtype)
        self.copyto(res)
        return res

    def copyto(self, other):
        """Copy the content of current array to other.

        When other is NDArray, the content is copied over.
        When other is a Context, a new NDArray in the context
        will be created as target

        Parameters
        ----------
        other : NDArray or Context
            Target NDArray or context we want to copy data to.

        Returns
        -------
        dst : NDArray
            The copy target NDArray
        """
        if isinstance(other, NDArray):
            if other.handle is self.handle:
                warnings.warn('copy an array to itself, is it intended?',
                              RuntimeWarning)
                return
            return _internal._copyto(self, out=other)
        elif isinstance(other, Context):
            hret = NDArray(_new_alloc_handle(self.shape, other, True, self.dtype))
            return _internal._copyto(self, out=hret)
        else:
            raise TypeError('copyto do not support type ' + str(type(other)))

    def copy(self):
        """Make a copy of the current ndarray on the same context

        Return
        ------
        cpy : NDArray
            The copy
        """
        return self.copyto(self.context)

    # pylint: enable= no-member

    def as_in_context(self, context):
        """Return an `NDArray` that lives in the target context. If the array
        is already in that context, `self` is returned. Otherwise, a copy is
        made.

        Parameters
        ----------
        context : Context
            The target context we want the return value to live in.

        Returns
        -------
        A copy or `self` as an `NDArray` that lives in the target context.
        """
        if self.context == context:
            return self
        return self.copyto(context)

def onehot_encode(indices, out):
    """One hot encoding indices into matrix out.

    Parameters
    ----------
    indices: NDArray
        An NDArray containing indices of the categorical features.

    out: NDArray
        The result holder of the encoding.

    Returns
    -------
    out: Array
        Same as out.
    """
    # pylint: disable= no-member, protected-access
    return _internal._onehot_encode(indices, out, out=out)
    # pylint: enable= no-member, protected-access

def empty(shape, ctx=None, dtype=mx_real_t):
    """Create an empty uninitialized new NDArray, with specified shape.

    Parameters
    ----------
    shape : tuple
        shape of the NDArray.

    ctx : Context, optional
        The context of the NDArray, default to current default context.

    Returns
    -------
    out: Array
        The created NDArray.
    """
    if isinstance(shape, int):
        shape = (shape, )
    if ctx is None:
        ctx = Context.default_ctx
    return NDArray(handle=_new_alloc_handle(shape, ctx, False, dtype))

#pylint: disable= too-many-arguments, no-member, protected-access
def _ufunc_helper(lhs, rhs, fn_array, fn_scalar, lfn_scalar, rfn_scalar=None):
    """ Helper function for element-wise operation
    The function will perform numpy-like broadcasting if needed and call different functions

    Parameters
    ----------
    lhs : NDArray or numeric value
        left hande side operand

    rhs : NDArray or numeric value
        right hand side operand

    fn_array : function
        function to be called if both lhs and rhs are of NDArray type

    fn_scalar : function
        function to be called if both lhs and rhs are numeric values

    lfn_scalar : function
        function to be called if lhs is NDArray while rhs is numeric value

    rfn_scalar : function
        function to be called if lhs is numeric value while rhs is NDArray;
        if none is provided, then the function is commutative, so rfn_scalar is equal to lfn_scalar

    Returns
    -------
    out: NDArray
        result array
    """
    if isinstance(lhs, numeric_types):
        if isinstance(rhs, numeric_types):
            return fn_scalar(lhs, rhs)
        else:
            if rfn_scalar is None:
                # commutative function
                return lfn_scalar(rhs, float(lhs))
            else:
                return rfn_scalar(rhs, float(lhs))
    elif isinstance(rhs, numeric_types):
        return lfn_scalar(lhs, float(rhs))
    elif isinstance(rhs, NDArray):
        # check whether broadcasting is needed
        lsize = functools.reduce(operator.mul, lhs.shape)
        rsize = functools.reduce(operator.mul, rhs.shape)
        if lsize < rsize:
            lhs = lhs.broadcast_to(rhs.shape)
        elif lsize > rsize:
            rhs = rhs.broadcast_to(lhs.shape)
        return fn_array(lhs, rhs)
    else:
        raise TypeError('type %s not supported' % str(type(rhs)))
#pylint: enable= too-many-arguments, no-member, protected-access

def add(lhs, rhs):
    """ Perform element-wise addition

    Parameters
    ----------
    lhs : Array or float value
        left hand side operand

    rhs : Array of float value
        right hand side operand

    Returns
    -------
    out: Array
        result array
    """
    # pylint: disable= no-member, protected-access
    return _ufunc_helper(
        lhs,
        rhs,
        _internal._plus,
        operator.add,
        _internal._plus_scalar,
        None)
    # pylint: enable= no-member, protected-access

def subtract(lhs, rhs):
    """ Perform element-wise subtract

    Parameters
    ----------
    lhs : Array or float value
        left hand side operand

    rhs : Array of float value
        right hand side operand

    Returns
    -------
    out: Array
        result array
    """
    # pylint: disable= no-member, protected-access
    return _ufunc_helper(
        lhs,
        rhs,
        _internal._minus,
        operator.sub,
        _internal._minus_scalar,
        _internal._rminus_scalar)
    # pylint: enable= no-member, protected-access

def multiply(lhs, rhs):
    """ Perform element-wise multiplication

    Parameters
    ----------
    lhs : Array or float value
        left hand side operand

    rhs : Array of float value
        right hand side operand

    Returns
    -------
    out: Array
        result array
    """
    # pylint: disable= no-member, protected-access
    return _ufunc_helper(
        lhs,
        rhs,
        _internal._mul,
        operator.mul,
        _internal._mul_scalar,
        None)
    # pylint: enable= no-member, protected-access

def divide(lhs, rhs):
    """ Perform element-wise divide

    Parameters
    ----------
    lhs : Array or float value
        left hand side operand

    rhs : Array of float value
        right hand side operand

    Returns
    -------
    out: Array
        result array
    """
    # pylint: disable= no-member, protected-access
    return _ufunc_helper(
        lhs,
        rhs,
        _internal._div,
        operator.truediv,
        _internal._div_scalar,
        _internal._rdiv_scalar)
    # pylint: enable= no-member, protected-access

def power(lhs, rhs):
    """ Perform power operator

    Parameters
    ----------
    lhs : Array or float value
        left hand side operand

    rhs : Array of float value
        right hand side operand

    Returns
    -------
    out: Array
        result array
    """
    # pylint: disable= no-member, protected-access
    return _ufunc_helper(
        lhs,
        rhs,
        _internal._power,
        operator.pow,
        _internal._power_scalar,
        _internal._rpower_scalar)
    # pylint: enable= no-member, protected-access

def maximum(lhs, rhs):
    """ Perform maximum operator

    Parameters
    ----------
    lhs : Array or float value
        left hand side operand

    rhs : Array of float value
        right hand side operand

    Returns
    -------
    out: Array
        result array
    """
    # pylint: disable= no-member, protected-access
    return _ufunc_helper(
        lhs,
        rhs,
        _internal._maximum,
        lambda x, y: x if x > y else y,
        _internal._maximum_scalar,
        None)
    # pylint: enable= no-member, protected-access

def minimum(lhs, rhs):
    """ Perform minimum operator

    Parameters
    ----------
    lhs : Array or float value
        left hand side operand

    rhs : Array of float value
        right hand side operand

    Returns
    -------
    out: Array
        result array
    """
    # pylint: disable= no-member, protected-access
    return _ufunc_helper(
        lhs,
        rhs,
        _internal._minimum,
        lambda x, y: x if x < y else y,
        _internal._minimum_scalar,
        None)
    # pylint: enable= no-member, protected-access

def true_divide(lhs, rhs):
    """ Same as numpy's true_divide. It adjusts the output type to present the best answer,
    regardless of input types.
    """
    return divide(lhs, rhs)

def negative(arr):
    """ Return the negation of array values """
    return multiply(arr, -1.0)

def zeros(shape, ctx=None, dtype=mx_real_t):
    """Create a new NDArray filled with 0, with specified shape.

    Parameters
    ----------
    shape : tuple
        shape of the NDArray.
    ctx : Context, optional.
        The context of the NDArray, default to current default context.

    Returns
    -------
    out: Array
        The created NDArray.
    """
    arr = empty(shape, ctx, dtype)
    arr[:] = 0.0
    return arr

def ones(shape, ctx=None, dtype=mx_real_t):
    """Create a new NDArray filled with 1, with specified shape.

    Parameters
    ----------
    shape : tuple
        shape of the NDArray.
    ctx : Context, optional
        The context of the NDArray, default to current default context.

    Returns
    -------
    out: Array
        The created NDArray.
    """
    arr = empty(shape, ctx, dtype)
    arr[:] = 1.0
    return arr

def full(shape, val, ctx=None):
    """Create a new NDArray filled with given value, with specified shape.

    Parameters
    ----------
    shape : tuple
        shape of the NDArray.
    val : float
        value to be filled with.
    ctx : Context, optional
        The context of the NDArray, default to current default context.

    Returns
    -------
    out: Array
        The created NDArray.
    """
    arr = empty(shape, ctx)
    arr[:] = val
    return arr

def array(source_array, ctx=None, dtype=mx_real_t):
    """Create a new NDArray that copies content from source_array.

    Parameters
    ----------
    source_array : array_like
        Source data to create NDArray from.

    ctx : Context, optional
        The context of the NDArray, default to current default context.

    Returns
    -------
    out: Array
        The created NDArray.
    """

    if not isinstance(source_array, np.ndarray):
        try:
            source_array = np.array(source_array, dtype=dtype)
        except:
            raise TypeError('source_array must be array like object')
    arr = empty(source_array.shape, ctx, dtype)
    arr[:] = source_array
    return arr

def concatenate(arrays, axis=0, always_copy=True):
    """Concatenate a list of NDArrays along the first dimension.

    Parameters
    ----------
    arrays : list of NDArray
        Arrays to be concatenate. They must have identical shape except
        the first dimension. They also must have the same data type.
    axis : int
        The axis along which to concatenate.
    always_copy : bool
        Default `True`. When not `True`, if the arrays only contain one
        `NDArray`, that element will be returned directly, avoid copying.

    Returns
    -------
    An `NDArray` that lives on the same context as `arrays[0].context`.
    """
    assert isinstance(arrays, list)
    assert len(arrays) > 0
    assert isinstance(arrays[0], NDArray)

    if not always_copy and len(arrays) == 1:
        return arrays[0]

    shape_axis = arrays[0].shape[axis]
    shape_rest1 = arrays[0].shape[0:axis]
    shape_rest2 = arrays[0].shape[axis+1:]
    dtype = arrays[0].dtype
    for arr in arrays[1:]:
        shape_axis += arr.shape[axis]
        assert shape_rest1 == arr.shape[0:axis]
        assert shape_rest2 == arr.shape[axis+1:]
        assert dtype == arr.dtype
    ret_shape = shape_rest1 + (shape_axis,) + shape_rest2
    ret = empty(ret_shape, ctx=arrays[0].context, dtype=dtype)

    idx = 0
    begin = [0 for _ in ret_shape]
    end = list(ret_shape)
    for arr in arrays:
        if axis == 0:
            ret[idx:idx+arr.shape[0]] = arr
        else:
            begin[axis] = idx
            end[axis] = idx+arr.shape[axis]
            # pylint: disable=no-member,protected-access
            _internal._crop_assign(ret, arr, out=ret,
                                   begin=tuple(begin),
                                   end=tuple(end))
            # pylint: enable=no-member,protected-access
        idx += arr.shape[axis]

    return ret

def load(fname):
    """Load ndarray from binary file.

    You can also use pickle to do the job if you only work on python.
    The advantage of load/save is the file is language agnostic.
    This means the file saved using save can be loaded by other language binding of mxnet.
    You also get the benefit being able to directly load/save from cloud storage(S3, HDFS)

    Parameters
    ----------
    fname : str
        The name of the file.Can be S3 or HDFS address (remember built with S3 support).
        Example of fname:

        - `s3://my-bucket/path/my-s3-ndarray`
        - `hdfs://my-bucket/path/my-hdfs-ndarray`
        - `/path-to/my-local-ndarray`

    Returns
    -------
    out : list of NDArray or dict of str to NDArray
        List of NDArray or dict of str->NDArray, depending on what was saved.
    """
    if not isinstance(fname, string_types):
        raise TypeError('fname need to be string')
    out_size = mx_uint()
    out_name_size = mx_uint()
    handles = ctypes.POINTER(NDArrayHandle)()
    names = ctypes.POINTER(ctypes.c_char_p)()
    check_call(_LIB.MXNDArrayLoad(c_str(fname),
                                  ctypes.byref(out_size),
                                  ctypes.byref(handles),
                                  ctypes.byref(out_name_size),
                                  ctypes.byref(names)))
    if out_name_size.value == 0:
        return [NDArray(NDArrayHandle(handles[i])) for i in range(out_size.value)]
    else:
        assert out_name_size.value == out_size.value
        return dict(
            (py_str(names[i]), NDArray(NDArrayHandle(handles[i]))) for i in range(out_size.value))

def save(fname, data):
    """Save list of NDArray or dict of str->NDArray to binary file.

    You can also use pickle to do the job if you only work on python.
    The advantage of load/save is the file is language agnostic.
    This means the file saved using save can be loaded by other language binding of mxnet.
    You also get the benefit being able to directly load/save from cloud storage(S3, HDFS)

    Parameters
    ----------
    fname : str
        The name of the file.Can be S3 or HDFS address (remember built with S3 support).
        Example of fname:

        - `s3://my-bucket/path/my-s3-ndarray`
        - `hdfs://my-bucket/path/my-hdfs-ndarray`
        - `/path-to/my-local-ndarray`

    data : list of NDArray or dict of str to NDArray
        The data to be saved.
    """
    handles = []
    if isinstance(data, dict):
        keys = []
        for key, val in data.items():
            if not isinstance(key, string_types):
                raise TypeError('save only accept dict str->NDArray or list of NDArray')
            if not isinstance(val, NDArray):
                raise TypeError('save only accept dict str->NDArray or list of NDArray')
            keys.append(c_str(key))
            handles.append(val.handle)
        keys = c_array(ctypes.c_char_p, keys)
    else:
        for val in data:
            if not isinstance(val, NDArray):
                raise TypeError('save only accept dict str->NDArray or list of NDArray')
            handles.append(val.handle)
        keys = None
    check_call(_LIB.MXNDArraySave(c_str(fname),
                                  mx_uint(len(handles)),
                                  c_array(NDArrayHandle, handles),
                                  keys))

def imdecode(str_img, clip_rect=(0, 0, 0, 0), out=None, index=0, channels=3, mean=None):
    """Decode an image from string. Requires OpenCV to work.

    Parameters
    ----------
    str_img : str
        binary image data
    clip_rect : iterable of 4 int
        clip decoded image to rectangle (x0, y0, x1, y1)
    out : NDArray
        output buffer. can be 3 dimensional (c, h, w) or 4 dimensional (n, c, h, w)
    index : int
        output decoded image to i-th slice of 4 dimensional buffer
    channels : int
        number of channels to output. Decode to grey scale when channels = 1.
    mean : NDArray
        subtract mean from decode image before outputing.
    """
    # pylint: disable= no-member, protected-access, too-many-arguments
    if mean is None:
        mean = NDArray(_new_empty_handle())
    if out is None:
        return _internal._imdecode(mean, index,
                                   clip_rect[0],
                                   clip_rect[1],
                                   clip_rect[2],
                                   clip_rect[3],
                                   channels,
                                   len(str_img),
                                   str_img=str_img)
    else:
        return _internal._imdecode(mean, index,
                                   clip_rect[0],
                                   clip_rect[1],
                                   clip_rect[2],
                                   clip_rect[3],
                                   channels,
                                   len(str_img),
                                   str_img=str_img,
                                   out=out)

# pylint: disable=too-many-locals, invalid-name
def _make_ndarray_function(handle):
    """Create a NDArray function from the FunctionHandle."""
    NDARRAY_ARG_BEFORE_SCALAR = 1
    ACCEPT_EMPTY_MUTATE_TARGET = 1 << 2
    # Get the property of NDArray
    n_used_vars = mx_uint()
    n_scalars = mx_uint()
    n_mutate_vars = mx_uint()
    type_mask = ctypes.c_int()
    check_call(_LIB.MXFuncDescribe(
        handle,
        ctypes.byref(n_used_vars),
        ctypes.byref(n_scalars),
        ctypes.byref(n_mutate_vars),
        ctypes.byref(type_mask)))
    n_mutate_vars = n_mutate_vars.value
    n_used_vars = n_used_vars.value
    n_scalars = n_scalars.value
    type_mask = type_mask.value
    accept_empty_mutate = (type_mask & ACCEPT_EMPTY_MUTATE_TARGET) != 0
    # infer type of the function
    if (type_mask & NDARRAY_ARG_BEFORE_SCALAR) != 0:
        use_vars_range = range(0, n_used_vars)
        scalar_range = range(n_used_vars, n_used_vars + n_scalars)
    else:
        scalar_range = range(0, n_scalars)
        use_vars_range = range(n_scalars, n_used_vars + n_scalars)

    # Get the information from the function
    name = ctypes.c_char_p()
    desc = ctypes.c_char_p()
    num_args = mx_uint()
    arg_names = ctypes.POINTER(ctypes.c_char_p)()
    arg_types = ctypes.POINTER(ctypes.c_char_p)()
    arg_descs = ctypes.POINTER(ctypes.c_char_p)()
    ret_type = ctypes.c_char_p()

    check_call(_LIB.MXFuncGetInfo(
        handle, ctypes.byref(name), ctypes.byref(desc),
        ctypes.byref(num_args),
        ctypes.byref(arg_names),
        ctypes.byref(arg_types),
        ctypes.byref(arg_descs),
        ctypes.byref(ret_type)))
    func_name = py_str(name.value)
    param_str = ctypes2docstring(num_args, arg_names, arg_types, arg_descs)
    doc_str = ('%s\n\n' +
               '%s\n' +
               'out : NDArray, optional\n' +
               '    The output NDArray to hold the result.\n\n'+
               'Returns\n' +
               '-------\n' +
               'out : NDArray\n'+
               '    The output of binary function.')
    doc_str = doc_str % (py_str(desc.value), param_str)

    # Definition of internal functions.
    def binary_ndarray_function(lhs, rhs, out=None, **kwargs):
        """Internal binary function
        """
        if out:
            if not isinstance(out, NDArray):
                raise TypeError('out must be NDArray')
            if not out.writable:
                raise TypeError('out must be writable')
        else:
            if not accept_empty_mutate:
                raise TypeError('argument out is required to call %s' % func_name)
            out = NDArray(_new_empty_handle())
        check_call(_LIB.MXFuncInvokeEx( \
                handle, \
                c_array(NDArrayHandle, (lhs.handle, rhs.handle)), \
                c_array(mx_float, ()), \
                c_array(NDArrayHandle, (out.handle,)), \
                ctypes.c_int(len(kwargs)), \
                c_array(ctypes.c_char_p, [c_str(key) for key in kwargs.keys()]), \
                c_array(ctypes.c_char_p, [c_str(str(i)) for i in kwargs.values()])))
        return out

    def unary_ndarray_function(src, out=None, *args, **kwargs):
        """internal NDArray function"""
        if out:
            if not isinstance(out, NDArray):
                raise TypeError('out must be NDArray')
            if not out.writable:
                raise TypeError('out must be writable')
        else:
            if not accept_empty_mutate:
                raise TypeError('argument out is required to call %s' % func_name)
            out = NDArray(_new_empty_handle())
        check_call(_LIB.MXFuncInvokeEx( \
                handle, \
                c_array(NDArrayHandle, (src.handle,)), \
                c_array(mx_float, [args[i] for i in scalar_range]), \
                c_array(NDArrayHandle, (out.handle,)), \
                ctypes.c_int(len(kwargs)), \
                c_array(ctypes.c_char_p, [c_str(key) for key in kwargs.keys()]), \
                c_array(ctypes.c_char_p, [c_str(str(i)) for i in kwargs.values()])))
        return out

    def generic_ndarray_function(*args, **kwargs):
        """Invoke this function by passing in parameters

        Parameters
        ----------
        *args
            Positional arguments of input scalars and NDArray
        out : NDArray or tuple of NDArray, optional
            Output NDArray, used to hold the output result.

        Returns
        -------
        out : NDArray
            The result NDArray(tuple) of result of computation.
        """
        if 'out' in kwargs:
            mutate_vars = kwargs['out']
            if isinstance(mutate_vars, NDArray):
                mutate_vars = (mutate_vars,)
            if len(mutate_vars) != n_mutate_vars:
                raise TypeError('expect %d out in %s', n_mutate_vars, func_name)
            del kwargs['out']
        else:
            if accept_empty_mutate:
                mutate_vars = tuple(
                    NDArray(_new_empty_handle()) for i in range(n_mutate_vars))
            else:
                raise TypeError('argument out is required to call %s' % func_name)
        check_call(_LIB.MXFuncInvokeEx( \
            handle, \
            c_array(NDArrayHandle, [args[i].handle for i in use_vars_range]), \
            c_array(mx_float, [args[i] for i in scalar_range]), \
            c_array(NDArrayHandle, [v.handle for v in mutate_vars]), \
            ctypes.c_int(len(kwargs)), \
            c_array(ctypes.c_char_p, [c_str(key) for key in kwargs.keys()]), \
            c_array(ctypes.c_char_p, [c_str(str(i)) for i in kwargs.values()])))
        if n_mutate_vars == 1:
            return mutate_vars[0]
        else:
            return mutate_vars
    # End of function declaration
    if n_mutate_vars == 1 and n_used_vars == 2 and n_scalars == 0:
        ret_function = binary_ndarray_function
    elif n_mutate_vars == 1 and n_used_vars == 1 and n_scalars == 0:
        ret_function = unary_ndarray_function
    else:
        ret_function = generic_ndarray_function
    ret_function.__name__ = func_name
    ret_function.__doc__ = doc_str
    return ret_function

# pylint: enable=too-many-locals, invalid-name

def _init_ndarray_module():
    """List and add all the ndarray functions to current module."""
    plist = ctypes.POINTER(FunctionHandle)()
    size = ctypes.c_uint()
    check_call(_LIB.MXListFunctions(ctypes.byref(size),
                                    ctypes.byref(plist)))

    module_obj = sys.modules[__name__]
    module_internal = sys.modules["mxnet._ndarray_internal"]
    for i in range(size.value):
        hdl = FunctionHandle(plist[i])
        function = _make_ndarray_function(hdl)
        # if function name starts with underscore, register as internal namespace
        if function.__name__.startswith('_'):
            setattr(module_internal, function.__name__, function)
        else:
            fname = function.__name__
            fn_obj = getattr(module_obj, fname, None)
            if fn_obj is None:
                setattr(module_obj, fname, function)
            else:
                setattr(module_obj, fname + '_internal', function)

# Initialize the NDArray module
_init_ndarray_module()


import pickle as pkl
from mmdet.core import BboxOverlaps2D, bbox_overlaps
import torch as t

import itertools
from collections import OrderedDict
from functools import partial

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.ticker as mticker
from cycler import cycler

def get_iou(pr,gt):
    if pr.shape[1] == 5:
        pr = t.FloatTensor(pr)[:,1:]
    else:
        pr = t.FloatTensor(pr)

    gt = t.FloatTensor(gt)[:,1:]
    ious = bbox_overlaps(pr,gt)
    v, idx = t.max(ious, dim=1)

    return v.view(-1)

def y_fmt(tick_val, pos):
    if tick_val > 1000000:
        val = int(tick_val)/1000000
        return '{:d} M'.format(int(val))
    elif tick_val > 1000:
        val = int(tick_val) / 1000
        return '{:d} k'.format(int(val))
    else:
        return tick_val

def filled_hist(ax, edges, values, bottoms=None, orientation='v',
                **kwargs):
    """
    Draw a histogram as a stepped patch.

    Extra kwargs are passed through to `fill_between`

    Parameters
    ----------
    ax : Axes
        The axes to plot to

    edges : array
        A length n+1 array giving the left edges of each bin and the
        right edge of the last bin.

    values : array
        A length n array of bin counts or values

    bottoms : scalar or array, optional
        A length n array of the bottom of the bars.  If None, zero is used.

    orientation : {'v', 'h'}
       Orientation of the histogram.  'v' (default) has
       the bars increasing in the positive y-direction.

    Returns
    -------
    ret : PolyCollection
        Artist added to the Axes
    """
    print(orientation)
    if orientation not in 'hv':
        raise ValueError("orientation must be in {{'h', 'v'}} "
                         "not {o}".format(o=orientation))

    kwargs.setdefault('step', 'post')
    edges = np.asarray(edges)
    values = np.asarray(values)
    if len(edges) - 1 != len(values):
        raise ValueError('Must provide one more bin edge than value not: '
                         'len(edges): {lb} len(values): {lv}'.format(
                             lb=len(edges), lv=len(values)))

    if bottoms is None:
        bottoms = 0
    bottoms = np.broadcast_to(bottoms, values.shape)

    values = np.r_[values, values[-1]]
    bottoms = np.r_[bottoms, bottoms[-1]]
    if orientation == 'h':
        return ax.fill_betweenx(edges, values, bottoms,
                                **kwargs)
    elif orientation == 'v':
        return ax.fill_between(edges, values, bottoms,
                               **kwargs)
    else:
        raise AssertionError("you should never be here")

def stack_hist(ax, stacked_data, sty_cycle, bottoms=None,
               hist_func=None, labels=None,
               plot_func=None, plot_kwargs=None):
    """
    Parameters
    ----------
    ax : axes.Axes
        The axes to add artists too

    stacked_data : array or Mapping
        A (N, M) shaped array.  The first dimension will be iterated over to
        compute histograms row-wise

    sty_cycle : Cycler or operable of dict
        Style to apply to each set

    bottoms : array, optional
        The initial positions of the bottoms, defaults to 0

    hist_func : callable, optional
        Must have signature `bin_vals, bin_edges = f(data)`.
        `bin_edges` expected to be one longer than `bin_vals`

    labels : list of str, optional
        The label for each set.

        If not given and stacked data is an array defaults to 'default set {n}'

        If stacked_data is a mapping, and labels is None, default to the keys
        (which may come out in a random order).

        If stacked_data is a mapping and labels is given then only
        the columns listed by be plotted.

    plot_func : callable, optional
        Function to call to draw the histogram must have signature:

          ret = plot_func(ax, edges, top, bottoms=bottoms,
                          label=label, **kwargs)

    plot_kwargs : dict, optional
        Any extra kwargs to pass through to the plotting function.  This
        will be the same for all calls to the plotting function and will
        over-ride the values in cycle.

    Returns
    -------
    arts : dict
        Dictionary of artists keyed on their labels
    """
    # deal with default binning function
    if hist_func is None:
        hist_func = np.histogram

    # deal with default plotting function
    if plot_func is None:
        plot_func = filled_hist

    # deal with default
    if plot_kwargs is None:
        plot_kwargs = {}
    print(plot_kwargs)
    try:
        l_keys = stacked_data.keys()
        label_data = True
        if labels is None:
            labels = l_keys

    except AttributeError:
        label_data = False
        if labels is None:
            labels = itertools.repeat(None)

    if label_data:
        loop_iter = enumerate((stacked_data[lab], lab, s)
                              for lab, s in zip(labels, sty_cycle))
    else:
        loop_iter = enumerate(zip(stacked_data, labels, sty_cycle))

    arts = {}
    for j, (data, label, sty) in loop_iter:
        if label is None:
            label = 'dflt set {n}'.format(n=j)
        label = sty.pop('label', label)
        vals, edges = hist_func(data)

        bottoms = np.zeros_like(vals)
        top = bottoms + vals
        print(sty)
        sty.update(plot_kwargs)
        print(sty)
        ret = plot_func(ax, edges, top, bottoms=bottoms,
                        label=label, **sty)
        bottoms = top
        arts[label] = ret
    ax.legend(fontsize=10)
    return arts

def draw(tag, stack_data):
    # set up histogram function to fixed bins
    edges = np.linspace(0.1, 1, 20, endpoint=True)
    hist_func = partial(np.histogram, bins=edges)

    # set up style cycles
    color_cycle = cycler(facecolor=['#f0fc03','#0356fc'])
    label_cycle = cycler(label=tag)
    alpha_cycle = cycler(alpha=[0.7, 0.3])
    # Fixing random state for reproducibility

    dict_data = OrderedDict(zip((c['label'] for c in label_cycle), stack_data))

    fig = plt.figure(figsize=(5,5), tight_layout=True)
    ax = plt.axes()
    arts = stack_hist(ax, stack_data, color_cycle + label_cycle + alpha_cycle,
                    hist_func=hist_func)
    ax.set_ylabel('counts')
    ax.set_xlabel('IoU')
    ax.yaxis.set_major_formatter(mticker.FuncFormatter(y_fmt))
    plt.xlim([0.1, 1])
    plt.ylim(bottom=0)
    plt.tight_layout()

    plt.savefig('xmTool/%s_%s' %(tag[0], tag[1]) , dpi=100)
    plt.close()

if __name__ == '__main__':

    boxfile = 'xmTool/cascade_prp.pkl'

    with open(boxfile, 'rb') as f:
        boxdict =pkl.load(f)

    gtfile = 'xmTool/gt_test.pkl'

    with open(gtfile, 'rb') as f1:
        gtdict =pkl.load(f1)

    tag1 = ['rpn', 'stage3']

    pr = boxdict[tag1[0]]
    pr2 = boxdict[tag1[1]]

    gt = gtdict['test']

    count = t.tensor([])
    count2 = t.tensor([])

    for i in range(len(pr)):

        ious = get_iou(pr[i], gt[i])
        count = t.cat((count, ious))

        ious2 = get_iou(pr2[i], gt[i])
        count2 = t.cat((count2, ious2))

    count = count.numpy()
    count2 = count2.numpy()

    stack_data = np.stack((count,count2))

    draw(tag1, stack_data)


"""
Deep Q Network agent.
"""

import random
from collections import namedtuple, deque
import numpy as np
import torch
import torch.nn.functional as F
import torch.optim as optim
#from visualize import show_frames_3d
import dill
from libs.agent_util import ReplayBuffer, PrioritizedReplayBuffer

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

class Agent():
    """Interacts with and learns from the environment."""

    def __init__(self, model, action_size, seed=0, load_file=None,
                 buffer_size=int(1e5),
                 batch_size=64,
                 gamma=0.99,
                 tau=1e-3,
                 lr=5e-4,
                 update_every=4,
                 use_double_dqn=True,
                 use_prioritized_experience_replay=False,
                 alpha_start=0.5,
                 alpha_decay=0.9992,
                 action_map=None):
        """
        Params
        ======
            model: model object
            action_size (int): dimension of each action
            seed (int): Random seed
            load_file (str): path of checkpoint file to load
            buffer_size (int): replay buffer size
            batch_size (int): minibatch size
            gamma (float): discount factor
            tau (float): for soft update of target parameters
            lr (float): learning rate
            update_every (int): how often to update the network
            use_double_dqn (bool): wheter to use double DQN algorithm
            use_prioritized_experience_replay (bool): wheter to use PER algorithm
            alpha_start (float): initial value for alpha, used in PER
            alpha_decay (float): decay rate for alpha, used in PER
            action_map (dict): how to map action indexes from model output to gym environment
        """
        random.seed(seed)

        self.action_size = action_size
        self.buffer_size = buffer_size
        self.batch_size = batch_size
        self.gamma = gamma
        self.tau = tau
        self.lr = lr
        self.update_every = update_every
        self.use_double_dqn = use_double_dqn
        self.use_prioritized_experience_replay = use_prioritized_experience_replay

        self.loss_list = []       # track loss across steps
        self.entropy_list = []    # track entropy across steps

        # Q-Network
        self.qnetwork_local = model.local
        self.qnetwork_target = model.target

        # DEBUG weight initialization
        #print(self.qnetwork_local.fc_s.weight.data[0])
        #print(self.qnetwork_target.fc_s.weight.data[0])
        #self.qnetwork_local.fc_s.weight.data[0] = torch.tensor([0.0, 0.0, 0.0, 0.0])
        #print(self.qnetwork_local.fc_s.weight.data[0])
        #print(self.qnetwork_target.fc_s.weight.data[0])
        #input('->')

        self.optimizer = optim.Adam(self.qnetwork_local.parameters(), lr=lr)
        #self.optimizer = optim.RMSprop(self.qnetwork_local.parameters(), lr=.00025, momentum=0.95)

        # Replay memory
        if use_prioritized_experience_replay:
            self.memory = PrioritizedReplayBuffer(action_size, self.buffer_size, self.batch_size, seed)
        else:
            self.memory = ReplayBuffer(action_size, self.buffer_size, self.batch_size, seed)
        # Initialize time step (for updating every update_every steps)
        self.t_step = 0
        # initalize alpha (used in prioritized experience sampling probability)
        self.alpha_start = alpha_start
        self.alpha_decay = alpha_decay
        self.alpha = self.alpha_start

        if load_file:
            self.qnetwork_local.load_state_dict(torch.load(load_file + '.pth'))
            self.qnetwork_target.load_state_dict(torch.load(load_file + '.pth'))
            #self.memory = dill.load(open(load_file + '.buffer.pck','rb'))
            print('Loaded: {}'.format(load_file))

        self.action_map = action_map

    def step(self, state, action, reward, next_state, done):
        # Save experience in replay memory
        if self.use_prioritized_experience_replay:
            priority = 100.0   # set initial priority to max value
            self.memory.add(state, action, reward, next_state, done, priority)
        else:
            self.memory.add(state, action, reward, next_state, done)

        # Learn every update_every time steps.
        self.t_step = (self.t_step + 1) % self.update_every
        if self.t_step == 0:
            # If enough samples are available in memory, get random subset and learn
            if len(self.memory) > self.batch_size:
                # if prioritized experience replay is enabled
                if self.use_prioritized_experience_replay:
                    self.memory.sort()
                    indexes, experiences = self.memory.sample(self.alpha)
                    self.learn(experiences, self.gamma, indexes)
                    self.alpha = self.alpha_decay*self.alpha
                else:
                    experiences = self.memory.sample()
                    self.learn(experiences, self.gamma)

    def act(self, state, eps=0.):
        """Returns actions for given state as per current policy.

        Params
        ======
            state (array_like): current state
            eps (float): epsilon, for epsilon-greedy action selection
        """
        if len(state.shape) == 1:   # reshape 1-D states into 2-D (as expected by the model)
            state = np.expand_dims(state, axis=0)
        state = torch.from_numpy(state).float().to(device)
        # calculate action values
        self.qnetwork_local.eval()
        with torch.no_grad():
            action_values = self.qnetwork_local(state)
        self.qnetwork_local.train()

        # Epsilon-greedy action selection
        if random.random() > eps:
            return np.argmax(action_values.cpu().data.numpy())
        else:
            return random.choice(np.arange(self.action_size))

    def learn(self, experiences, gamma, indexes=None):
        """Update value parameters using given batch of experience tuples.

        Params
        ======
            experiences (Tuple[torch.Tensor]): tuple of (s, a, r, s', done) tuples
            gamma (float): discount factor
        """
        if self.use_prioritized_experience_replay:
            states, actions, rewards, next_states, dones, priorities = experiences
        else:
            states, actions, rewards, next_states, dones = experiences

        # DEBUG replay memory
        #print('learning:')
        #show_frames(states)
        #show_frames(next_states)

        # Select double DQN or regular DQN
        if self.use_double_dqn:
            # get greedy actions (for next states) from local model
            q_local_argmax = self.qnetwork_local(next_states).detach().argmax(dim=1).unsqueeze(1)
            # get predicted q values (for next states) from target model indexed by q_local_argmax
            q_targets_next = self.qnetwork_target(next_states).gather(1, q_local_argmax).detach()
        else:
            # get max predicted q values (for next states) from target model
            q_targets_next = self.qnetwork_target(next_states).detach().max(1)[0].unsqueeze(1)

        # get q values from local model
        q_local = self.qnetwork_local(states)
        # get q values for chosen action
        predictions = q_local.gather(1, actions)
        # calculate td targets
        targets = rewards + (gamma * q_targets_next * (1 - dones))

        # calculate new priorities
        if self.use_prioritized_experience_replay:
            with torch.no_grad():
                new_priorities = torch.abs(targets - predictions).to(device)
                self.memory.batch_update(indexes, (states, actions, rewards, next_states, dones, new_priorities))

        # calculate loss using mean squared error: (targets - predictions).pow(2).mean()
        loss = F.mse_loss(predictions, targets)
        # minimize loss
        self.optimizer.zero_grad()
        loss.backward()
        # clip gradients
        #for param in self.qnetwork_local.parameters():
        #    param.grad.data.clamp_(-10, 10)
        self.optimizer.step()

        # update stats
        with torch.no_grad():
            self.loss_list.append(loss.item())
            # calculate sparse softmax cross entropy
            self.entropy_list.append(F.cross_entropy(q_local, actions.squeeze(1)))

        # update target network
        self.soft_update(self.qnetwork_local, self.qnetwork_target, self.tau)

    def soft_update(self, local_model, target_model, tau):
        """Soft update model parameters.
        θ_target = τ*θ_local + (1 - τ)*θ_target

        Params
        ======
            local_model (PyTorch model): weights will be copied from
            target_model (PyTorch model): weights will be copied to
            tau (float): interpolation parameter
        """
        for target_param, local_param in zip(target_model.parameters(), local_model.parameters()):
            target_param.data.copy_(tau*local_param.data + (1.0-tau)*target_param.data)


import torch.nn as nn
import torchvision
from haven import haven_img
import torch
import cv2
import numba
# from . import utils as ut
from haven import haven_img as hi
from skimage import morphology as morph
import numpy as nopython
import torch.utils.model_zoo as model_zoo
# from . import _losses as _losses
from torch import optim
import torch.nn.functional as F
from torchvision.transforms import functional as FT
# from . import lcs_utils
# from models.scoring_net import lcs_utils
from torchvision.models.detection.rpn import concat_box_prediction_layers
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor

import torch
import numpy as np
import torch.nn.functional as F
import copy
from pycocotools import mask as maskUtils
# from datasets import _transforms
from torchvision.transforms import functional as FT
import numba
import torch.nn as nn
import torchvision
from haven import haven_utils as hu
import torch
# from haven._toolbox import misc as ms
from skimage import morphology as morph
import numpy as np
# from haven.managers import exp_manager as em
import torch.utils.model_zoo as model_zoo
from torchvision.models import detection
# from . import _losses as _losses
from torch import optim
import torch.nn.functional as F
# from maskrcnn_benchmark.structures.image_list import to_image_list
from src import models
from collections import OrderedDict

# import utils as ut

from scipy.spatial.distance import cdist
# import cvxpy as cp
# from cvxpylayers.torch import CvxpyLayer
import os

class AttModel(nn.Module):
    def __init__(self, exp_dict):
        super().__init__()
        self.att_dict = exp_dict['attention']

        self.rcnn = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
        self.rcnn.rpn.nms_thresh = 0.7

        self.box_score_method = self.att_dict['box_score_method']
        self.agg_score_method = self.att_dict['agg_score_method']
        self.threshold = self.att_dict['agg_score_method']

    @torch.no_grad()
    def get_attention_dict(self, images_original, counts, probs, return_roi=True):

        self.eval()

        assert(counts > 0)

        # 1. get bbox and their scores
        _, H_img, W_img = images_original[0].shape
        images, targets = self.rcnn.transform([images_original.cuda()[0]], None)
        bbox, obj_scores, fpn_features = rpn_forward(self.rcnn, images)

        # 2. filter based on max area
        # ---------------------------
        bbox, obj_scores = filter_by_area(bbox, obj_scores, images, counts=counts, T=0.01)

        # box_score_method
        if self.box_score_method == 'mean':
            box_scores = score_bbox_heatmap(probs, bbox, images,
                                    reduction='mean')

        elif self.box_score_method == 'center':
            box_scores = score_bbox_heatmap(probs, bbox, images,
                                    reduction='center')

        elif self.box_score_method == 'objectness':
            box_scores = obj_scores
        else:
            stop

        # agg_score_method
        if self.agg_score_method == 'max':
            scores = torch.max(obj_scores, box_scores)
        elif self.agg_score_method == 'mean':
            scores = (obj_scores + box_scores) / 2.
        elif self.agg_score_method == 'replace_0.8':
            scores[box_scores>0.8] = box_scores[box_scores>0.8]
        elif self.agg_score_method == 'replace_0.7':
            scores[box_scores>0.7] = box_scores[box_scores>0.7]
        elif self.agg_score_method == 'replace':
            scores = box_scores
        else:
            stop

        print()

        # 3. rescore bbox and sort
        # if np.random.rand() < 0.3:
        ind = scores.argsort(descending=True)
        bbox, scores = bbox[ind], scores[ind]

        scores_top = copy.deepcopy(scores)

        # 5. normalize bbox
        bbox_norm = normalize_bbox(bbox.clone(), images)

        # 4. get top bbox
        ind_fg = get_top_bbox(counts, bbox, scores_top, T=0.01, ratio_curr=1.0)
        if np.random.rand() <= self.att_dict['select_prob']:
        # if np.random.rand() <= 1.:
            end = max(int((scores[ind_fg] > 0.7).float().sum()), 1)
            ind_fg = ind_fg[:end]

            # get roi_mask
            mask_bg = 1 - bbox_to_mask(bbox_norm, H_img, W_img)
            mask_fg = bbox_to_mask(bbox_norm[ind_fg], H_img, W_img)
            roi_mask = mask_fg + mask_bg

        else:
            roi_mask = None

        points = bbox2points(bbox_norm[ind_fg], H_img, W_img)

        return {
                # 'bbox_norm':bbox_norm,
                # 'scores': scores,
                # 'ind_fg':ind_fg,
                'points':torch.LongTensor(points[None]),
                'roi_mask':roi_mask
                }

def score_bbox_points(probs, points, bbox, batch, images):

    if points.sum() == 0:
        return torch.zeros(bbox.shape[0]).cuda()

    points_probs = points * probs.cpu().numpy()

    Hp, Wp = points.squeeze().shape
    yp, xp = np.where(points.squeeze())
    points_probs_list = points_probs.squeeze()[yp, xp]
    bbox_norm = normalize_bbox(bbox.clone(), images)
    xmin, ymin, xmax, ymax = torch.chunk(bbox_norm, 4, dim=1)
    yp = torch.from_numpy(yp).cuda().float()[None] / Hp
    xp = torch.from_numpy(xp).cuda().float()[None] / Wp

    D = (xp - xmin) >= 0
    D &= (yp - ymin) >= 0
    D &= (xmax - xp) >= 0
    D &= (ymax - yp) >= 0

    scores = (D.float() * torch.FloatTensor(points_probs_list)[None].cuda())
    scores = scores.max(dim=1)[0]

    return scores

def get_roi_mask(bbox, ind_fg, mask_bg, images):
    H, W = mask_bg.shape
    bbox_norm = normalize_bbox(bbox.clone(), images)
    mask_fg = bbox_to_mask(bbox_norm[ind_fg], H, W)

    roi_mask = ((mask_bg) + mask_fg)

    return roi_mask

def bbox_on_image(batch, bbox_norm, scores, ind_fg):
    img_org = hu.denormalize(batch['images'], mode='rgb')
    img_org = img_org / img_org.max()
    img_box = bbox_on_image(bbox_norm[ind_fg], img_org.squeeze(), text_list=scores)

    return np.hstack([img_box])

@torch.no_grad()
def score_bbox_heatmap(probs, bbox, images, reduction):
    # import ipdb; ipdb.set_trace()
    # points = lcfcn.predict_on_batch(batch, method='points')
    bbox_norm = normalize_bbox(bbox.clone(), images)
    scores = torch.zeros(bbox.shape[0]).cuda()

    H, W = probs.shape
    for i, bb in enumerate(bbox_norm):
        xmin, ymin, xmax, ymax = bb

        prob_bb = probs[int(ymin*H):int(ymax*H), int(xmin*W):int(xmax*W)]
        if reduction == 'mean':
            scores[i] = torch.mean(prob_bb)
        elif reduction == 'center':
            hi, wi = prob_bb.shape
            scores[i] = prob_bb[hi//2, wi//2]
        elif reduction == 'max':
            scores[i] = torch.max(prob_bb)
        else:
            raise ValueError('heatmap_method does not exist')

    return scores

    if points.sum() == 0:
        return torch.zeros(bbox.shape[0]).cuda()

    points_probs = points * probs.cpu().numpy()

    Hp, Wp = points.squeeze().shape
    yp, xp = np.where(points.squeeze())
    points_probs_list = points_probs.squeeze()[yp, xp]
    bbox_norm = normalize_bbox(bbox.clone(), images)
    xmin, ymin, xmax, ymax = torch.chunk(bbox_norm, 4, dim=1)
    yp = torch.from_numpy(yp).cuda().float()[None] / Hp
    xp = torch.from_numpy(xp).cuda().float()[None] / Wp

    D = (xp - xmin) >= 0
    D &= (yp - ymin) >= 0
    D &= (xmax - xp) >= 0
    D &= (ymax - yp) >= 0

    scores = (D.float() * torch.FloatTensor(points_probs_list)[None].cuda())
    scores = scores.max(dim=1)[0]

    return scores

def filter_by_area(bbox, scores, images, counts, T, top_void=None):
    ind_fg = get_top_bbox(counts, bbox, scores, T=T, top_void=top_void)
    bbox_norm = normalize_bbox(bbox.clone(), images)
    areas = get_bbox_areas(bbox_norm)
    if len(ind_fg):

        max_area = areas[ind_fg].max()
    else:
        max_area = 0
    ind = areas <= max_area
    bbox, scores = bbox[ind], scores[ind]

    return bbox, scores

def normalize_bbox(bbox, images):
    H, W = images.image_sizes[0]

    bbox[:, [0,2]] = bbox[:, [0,2]] / W
    bbox[:, [1,3]] = bbox[:, [1,3]] / H

    return bbox

def bbox2points(bbox, H, W):
    points = np.zeros((H, W), dtype=int)
    for bb in bbox:
        x1, y1, x2, y2 = bb
        x1, y1 = (int(x1*W), int(y1*H),)
        x2, y2 = (int(x2*W), int(y2*H),)

        xc, yc = (x2 + x1)//2, (y2 + y1)//2
        points[int(yc), int(xc)] = 1

    return points

def rpn_forward(rcnn, images):
    features = rcnn.backbone(images.tensors)

    if isinstance(features, torch.Tensor):
        features = OrderedDict([('0', features)])

    features_list = list(features.values())

    objectness, pred_bbox_deltas = rcnn.rpn.head(features_list)
    anchors = rcnn.rpn.anchor_generator(images, features_list)

    num_images = len(anchors)
    num_anchors_per_level = [o[0].numel() for o in objectness]
    objectness, pred_bbox_deltas = \
        concat_box_prediction_layers(objectness, pred_bbox_deltas)

    proposals = rcnn.rpn.box_coder.decode(pred_bbox_deltas.detach(), anchors)
    proposals = proposals.view(num_images, -1, 4)
    boxes, scores = rcnn.rpn.filter_proposals(proposals, objectness, images.image_sizes,
                            num_anchors_per_level)

    # normalize boxes
    bbox = boxes[0]

    # normalize scores
    scores = scores[0]
    scores += abs(scores.min())
    scores /= scores.max()

    return bbox, scores, features

def get_top_bbox(counts, bbox, scores,  T, top_void=None, ratio_curr=None):
    if counts == 0:
        return []

    if ratio_curr:
        counts_top = max(1, int(counts * ratio_curr))
    elif top_void:
        counts_top = min(int(counts), (scores >= top_void).sum().item())
    else:
        counts_top = counts

    ind_fg = []
    while len(ind_fg) == 0 and T <= 1.0 and counts_top != 0:
        ind_fg = filter_c_wsl(bbox, scores, counts=counts_top, T=T)
        T = T *1.1

    # ind_bg = np.setdiff1d(np.arange(bbox.shape[0]), ind_fg)

    return ind_fg

def filter_c_wsl(bbox, scores, counts, T=0.5):
    # assert ordered
    # assert(1 == torch.mean((torch.arange(bbox.shape[0]) == scores.argsort(descending=True).cpu()).float()).item())
    bbox_overlaps = torchvision.ops.box_iou(bbox, bbox)

    ind = c_wsl_fast(bbox_overlaps.cpu().numpy(), scores.cpu().numpy(),
                        int(counts), T=T)

    return ind

@numba.jit(nopython=True)
def c_wsl_fast(bbox_overlaps, scores, counts, T=0.5):
    N = scores.shape[0]

    G_score_best = 0

    G_matrix = np.ones((N, N), dtype=np.int32) * -1
    G_size_vector = np.zeros(N, dtype=np.int32)

    # Find best proposals
    for i in range(N):
        box_i = i
        G = G_matrix[i]
        G_size = G_size_vector[i]

        G[0] = box_i
        G_size += 1
        G_score = scores[box_i]

        # Find the best compatible proposals to add
        for j in range(i + 1, N):
            # Check if we selected enough proposals
            if G_size == counts:  # j == N-1?
                # Check if the total score is hier than the previous one
                if G_score >= G_score_best:
                    G_score_best = G_score
                    G_best = G
                    G_size_best = G_size
                break

            box_j = j
            # Check if this proposal is compatible with the selected ones in G
            compatible = True
            for k in range(G_size):
                box_k = G[k]

                overlapping = bbox_overlaps[box_k, box_j]

                if overlapping > T:
                    compatible = False
                    break

            if not compatible:
                continue

            # Add this proposal to the list of selected ones
            G[G_size] = box_j
            G_size += 1
            G_score += scores[j]

    best_ind = G_best[:G_size_best]
    return best_ind

import torch
import tqdm
import argparse
import pandas as pd
import pickle, os
import numpy as np
from torch.utils.data import sampler
import json
import cv2
from haven import haven_utils as hu

def get_bbox_areas(bbox_norm):
    area_list = np.ones(bbox_norm.shape[0])*-1

    for i, bb in enumerate(bbox_norm):
        x1, y1, x2, y2 = bb

        area = (y2-y1) * (x2-x1)
        area_list[i] = area

    return area_list

def bbox_to_mask(bbox_norm, H, W):
    mask = np.zeros((H, W))
    for i, bb in enumerate(bbox_norm):
        x1, y1, x2, y2 = bb
        x1, y1 = int(x1*W), int(y1*H)
        x2, y2 = int(x2*W), int(y2*H)

        mask[y1:y2, x1:x2] = 1
    return mask

def bbox_on_image(bbox_xyxy, image, text_list=None):
    image = hu.f2l(image.squeeze())
    image_uint8 = (image*254).astype("uint8").copy()
    H, W, _ = image.shape
    for i, bb in enumerate(bbox_xyxy):
        x1, y1, x2, y2 = bb
        x1, y1 = int(x1*W), int(y1*H)
        x2, y2 = int(x2*W), int(y2*H)

        # Blue color in BGR
        color = (255, 0, 0)

        # Line thickness of 2 px
        thickness = 2

        # Using cv2.rectangle() method
        # Draw a rectangle with blue line borders of thickness of 2 px
        image_uint8 = cv2.rectangle(image_uint8, (x1, y1), (x2, y2), color, thickness)
        if text_list is not None:
            image_uint8 = cv2.putText(image_uint8, str(int(text_list[i]*100)),
                        (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX,  0.6, (0, 0, 255) , 1, cv2.LINE_AA)
    # hu.save_image("/mnt/datasets/public/issam/prototypes/wscl/tmp.jpg", image_uint8)
    return image_uint8 / 255.

import numpy,from scipy,import numba,import cvxpy,from cvxpy


#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Module to drive robot with a given acceleration profile. Should be
called directly from the main.py wrapper. Arguments are defined in
the AccControl class docstring below.

Questions? sherwinl@bu.edu, BU CODES Lab
"""
from __future__ import division
import rospy
import numpy as np
import std_msgs.msg
import geometry_msgs.msg
import tf
import config as cfg
import os
import sys
sys.path.append('../')

from lateral_control import lateral_control
from trials import eco_and, human_trial
from helper.trajectory import liveplot
from helper.logprint import logprint
from geometry_msgs.msg import Twist, PoseStamped
from tf.transformations import quaternion_from_euler, euler_from_quaternion

hostname = os.uname()[1]
MIN_V = 0
MAX_V = 0.4
glob_cmd_vel = 0

def yaw_to_vec(angle):
    """
    Converts yaw component of euler angle to
    vector in Cartesian space. Assumes a 0
    angle heading is pointing toward the
    positive x-axis in the coordinate system.

    Args:
      angle (float): Angle of yaw orientation

    Returns:
      (x,y) (tuple): Cartesian vector
    """

    if abs(angle) > np.pi/2:
        x = -1
        y = x*np.tan(angle)
    elif abs(angle) < np.pi/2:
        x = 1
        y = x*np.tan(angle)
    elif angle == np.pi/2:
        x = 0
        y = 1
    else:
        x = 0
        y = -1

    return (x,y)

class AccControl():
    """
    Acceleration control class where each object instance is a discrete
    acceleration profile for the robot.

    Args:
      robot (str): Name of robot
      accel (float): Input acceleration in m/s^2
      vel (float): Input velocity in m/s (for constant speed)
      zero_hd (float): Heading (in deg) when robot is pointed at south wall
      delay (int): Time delay for robot motion begins (in sec)
      trial (str): Control scenario being trialed ('ECO-AND' or 'HUMAN')
    """
    def __init__(self, robot, accel, vel=0, zero_hd=0,
                 delay=0, trial=None):
        rospy.init_node('acc_controller')
        self.hz = 10
        self.rate = rospy.Rate(self.hz)

        # Control trial
        self.trial = trial

        # Publishers
        self.vel_pub = rospy.Publisher('/cmd_vel',Twist,queue_size=10)
        self.mocap = rospy.Subscriber('/vrpn_client_node/' + robot + '/pose',
                                      PoseStamped,
                                      self.mocap_cb)

        # Message datatypes
        self.cmd_vel = Twist()
        self.mocap = PoseStamped()

        # Input arguments
        self.init_accel = accel
        self.accel = self.init_accel
        if vel != 0:
            self.vel = vel
        else:
            self.vel = glob_cmd_vel
        self.delay = delay * self.hz
        self.zero_hd = np.radians(zero_hd)

        # Calibration adjustment
        self.calibrate = 0

        # Initialize mocap data
        self.initialize()

    def initialize(self):
        """
        Initial conditions.
        """
        self.rate.sleep()

        # Initial pose and orientation
        self.init_pos_x = self.mocap.pose.position.x
        self.init_pos_y = self.mocap.pose.position.y
        init_or_z = self.mocap.pose.orientation.z
        init_or_w = self.mocap.pose.orientation.w
        init_yaw = euler_from_quaternion((0,0,init_or_z,init_or_w))[2]
        init_yaw += np.radians(self.calibrate) # CALIBRATION EPSILON

        # Coordinate transform if 0 heading is not south-facing
        if self.zero_hd == 0:
            self.init_heading = init_yaw
        else:
            if (init_yaw - self.zero_hd) < -np.pi:
                self.init_heading = (init_yaw - self.zero_hd) % np.pi
            else:
                self.init_heading = init_yaw - self.zero_hd

    def mocap_cb(self, data):
        """
        Callback function for mocap data.
        """
        self.mocap = data

    def update_acc(self, accel):
        """
        Update acceleration state.
        """
        self.accel = accel

    def control(self):
        """
        Primary controller method.
        """
        global glob_cmd_vel
        self.timer = 0

        # Initialize arrays to store robot pose
        self.x = []
        self.y = []

        # Trajectory state
        prev_state = 0

        while not rospy.is_shutdown():
            self.timer += 1
            self.rate.sleep()

            if self.timer >= self.delay:

                # Current pose and orientation
                pose_x = self.mocap.pose.position.x
                pose_y = self.mocap.pose.position.y
                orient_z = self.mocap.pose.orientation.z
                orient_w = self.mocap.pose.orientation.w
                yaw = euler_from_quaternion((0,0,orient_z,orient_w))[2]
                yaw += np.radians(self.calibrate) # CALIBRATION EPSILON

                # Stop if robot exceeds floor boundaries
                if ((pose_x < cfg.X_LABLIM[0]) or (pose_x > cfg.X_LABLIM[1]) or
                        (pose_y < cfg.Y_LABLIM[0]) or (pose_y > cfg.Y_LABLIM[1])):
                    break

                # Save pose data for trajectory plot
                self.x.append(pose_x)
                self.y.append(pose_y)

                # Coordinate transform if 0 heading is not south-facing
                if self.zero_hd != 0:
                    if (yaw - self.zero_hd) < -np.pi:
                        yaw = (yaw - self.zero_hd) % np.pi
                    else:
                        yaw = yaw - self.zero_hd

                # Update acceleration based on trial
                if self.trial == 'ECO-AND':
                    eco_and(pose_y, self.timer, self.hz)
                    self.accel = cfg.accel
                elif self.trial == 'HUMAN':
                    human_trial(pose_y, self.timer, self.hz, self.vel)
                    self.accel = cfg.accel

                # Current orientation vector (info only)
                orient_vec = yaw_to_vec(yaw)
                ov_unit = orient_vec / np.linalg.norm(orient_vec)

                # Compute velocity
                if self.accel >= 0:
                    if MIN_V <= round(self.vel,2) < MAX_V:
                        self.vel += (self.accel / self.hz)
                else:
                    if MIN_V < round(self.vel,2) <= MAX_V:
                        self.vel += (self.accel / self.hz)
                    else:
                        self.vel = 0

                # Zero out velocity if calculation is negative
                if self.vel < 0:
                    self.vel = 0

                # Adjust path trajectory with lateral control
                lc = lateral_control(prev_state, [self.init_pos_x,
                                     self.init_pos_y], [pose_x,pose_y],
                                     self.init_heading, yaw, self.vel)

                prev_state = lc[1]
                self.cmd_vel.angular.z = lc[0]
                self.cmd_vel.linear.x = self.vel
                glob_cmd_vel = self.cmd_vel.linear.x
                self.vel_pub.publish(self.cmd_vel)

                # Print data to log file
                logprint(self.timer,[self.init_pos_x,self.init_pos_y],
                         [pose_x,pose_y],[orient_z,orient_w],
                         np.degrees(yaw),ov_unit,
                         [self.cmd_vel.linear.x,
                         self.cmd_vel.angular.z])


import streamlit as st
import librosa
import librosa.display
import os
import mido
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib.colors import colorConverter
import plotly.graph_objects as go
import plotly.express as px

from turtle import bgcolor
import re
from pathlib import Path

class MidiFile(mido.MidiFile):
    def __init__(self, midifile, verbose=False):
        self.sr = 10   # down sampling rate from MIDI to time axis
        self.meta = {}
        self.max_nch = 16

        mido.MidiFile.__init__(self, midifile)
        self.fpath = Path(midifile)

        self.events, self.nch = self.get_events(verbose)
        self.roll, self.note_range, self.intensity_range = self.get_roll(self.events)

        self.length_ticks = self.get_total_ticks()
        self.length_seconds = mido.tick2second(
            self.length_ticks, self.ticks_per_beat, self.get_tempo())
        self.ticks_per_sec = self.length_ticks/self.length_seconds # miditicks/sec
        self.xticks_per_sec = self.ticks_per_sec / self.sr

        st.sidebar.write('## midi file')
        st.sidebar.write("Num. of tracks: ", len(self.tracks))
        st.sidebar.write("Num. of active channels: ", self.nch)
        st.sidebar.write("Intensity range [0, 100]: [{}, {}]".format(
            self.intensity_range[0], self.intensity_range[1]))
        st.sidebar.write("Note range [0, 127]: [{}, {}]".format(
            self.note_range[0], self.note_range[1]))
        st.sidebar.write("ticks/beat: ", self.ticks_per_beat)
        st.sidebar.write("ticks/second: ", self.ticks_per_sec)
        st.sidebar.write("Tick length: [ticks]", self.length_ticks)
        st.sidebar.write("Time length [s]: ", self.length_seconds)

    @st.cache
    def get_tempo(self):
        try:
            return self.meta["set_tempo"]["tempo"]
        except:
            return 500000

    @st.cache
    def get_total_ticks(self):
        max_ticks = 0
        for channel in range(self.nch):
            ticks = sum(msg.time for msg in self.events[channel])
            if ticks > max_ticks:
                max_ticks = ticks
        return max_ticks

    @st.cache
    def get_events(self, verbose=False):
        """
        Extract self.max_nch (default: 16) channel data from MIDI and return a list.
        Lyrics and meta data used in extra channels are not include in the list.

        Returns:
            list : [[ch1],[ch2]....[ch16]] # Note that empty channel is removed!
        """
        if verbose:
            print(self)

        mid = self
        events =  [[] for i in range(self.max_nch)]

        for track in mid.tracks:
            for msg in track:
                try:
                    channel = msg.channel
                    events[channel].append(msg)
                except AttributeError:
                    try:
                        if type(msg) != type(mido.UnknownMetaMessage):
                            self.meta[msg.type] = msg.dict()
                        else:
                            pass
                    except:
                        print("error", type(msg))
        events = list(filter(None, events)) # remove emtpy channel

        return events, len(events)

    @st.cache
    def get_roll(self, events, verbose=False):
        """
        Convert event (channel) data to piano roll data
        """
        intensity_range = [100,0] # [min, max] adjusted by get_roll()
        note_range = [127,0] # [min, max] adjusted by get_roll()

        length_ticks = self.get_total_ticks()  # get total length in tick unit

        roll = np.zeros(
            (self.nch, 128, length_ticks // self.sr), dtype="int8")
        register_note = [int(-1)]*128        # register the state (on/off) of each key
        register_timbre = np.ones(self.nch)  # register the state (program_change) of each channel

        for idx, channel in enumerate(events):
            time_counter = 0
            volume = 100

            if verbose:
                print("channel", idx, "start")

            for msg in channel:
                if msg.type == "control_change":
                    if msg.is_cc(7):  # if msg.control == 7: Main Volume
                        volume = 100*msg.value //127  # [0, 100]

                    if msg.is_cc(11):  # if msg.control == 11: Expression Controller
                        # volume[0,100] x expression[0,127]/127
                        volume *= msg.value // 127

                if msg.type == "program_change":
                    register_timbre[idx] = msg.program
                    if verbose:
                        print("channel", idx, "pc", msg.program, "time",
                              time_counter, "duration", msg.time)

                if msg.type == "note_on":
                    if verbose:
                        print("on ", msg.note, "time", time_counter,
                              "duration", msg.time, "velocity", msg.velocity)

                    # note_on_start_time = time_counter // self.sr
                    note_on_end_time = (time_counter + msg.time) // self.sr
                    intensity = volume * msg.velocity // 127

                    if intensity_range[0] > intensity: # update minimum intensity
                        intensity_range[0] = intensity
                    if intensity_range[1] < intensity: # update maximum intensity
                        intensity_range[1] = intensity

                    if register_note[msg.note] != -1:  # not after note_off
                        last_end_time, last_intensity = register_note[msg.note]
                        roll[idx, msg.note,
                             last_end_time:note_on_end_time] = last_intensity

                    register_note[msg.note] = (note_on_end_time, intensity)

                    if note_range[0] > msg.note: # update minimum note
                        note_range[0] = msg.note
                    if  note_range[1] < msg.note: # update maximum note
                        note_range[1] = msg.note

                if msg.type == "note_off":
                    if verbose:
                        print("off", msg.note, "time", time_counter,
                              "duration", msg.time, "velocity", msg.velocity)

                    # note_off_start_time = time_counter // self.sr
                    note_off_end_time = (time_counter + msg.time) // self.sr
                    last_end_time, last_intensity = register_note[msg.note]
                    roll[idx, msg.note,
                         last_end_time:note_off_end_time] = last_intensity

                    register_note[msg.note] = -1  # reinitialize register

                time_counter += msg.time

            # if there is a note not closed at the end of a channel, close it
            for key, data in enumerate(register_note):
                if data != -1:
                    note_on_end_time = data[0]
                    intensity = data[1]
                    # note_off_start_time = time_counter // self.sr
                    roll[idx, key, note_on_end_time:] = intensity
                register_note[idx] = -1

        return roll, note_range, intensity_range

    def _grp_init(self, figsize=(15, 9), xlim=None, ylim=None, bgcolor='white'):
        """
        Display basic information and initialize graphics.
        Called by draw_roll()
        """
        dsp_len_seconds = xlim[1]-xlim[0]

        # x ticks
        xticks_interval_sec = dsp_len_seconds // 10 if dsp_len_seconds > 10 else dsp_len_seconds/10
        xticks_interval =  xticks_interval_sec * self.xticks_per_sec
        #xticks_interval = mido.second2tick(
        #    xticks_interval_sec, self.ticks_per_beat, self.get_tempo()) / self.sr  # [ticks/interval]
        #print("xticks_interval_sec: ", xticks_interval_sec)
        #print("xticks_interval: {} [ticks/label]".format(xticks_interval))

        # Initialize graphics
        plt.rcParams["font.size"] = 20
        fig = plt.figure(figsize=figsize)
        ax = fig.add_subplot(111)
        ax.axis("equal")
        ax.set_facecolor(bgcolor)

        nxticks = int(self.length_ticks//xticks_interval)
        plt.xticks(
            [int(x * xticks_interval) for x in range(nxticks)],
            [round(x * xticks_interval_sec, 2) for x in range(nxticks)]
        )
        plt.yticks([y*12 for y in range(9)], [y*12 for y in range(9)])

        ax.set_xlabel("time [s]")
        ax.set_ylabel("note")
        xlim_ticks=[0,self.length_ticks-1]
        if xlim != None:
            #xticks_per_sec = xticks_interval/xticks_interval_sec
            xlim_ticks=np.array(xlim)*self.xticks_per_sec
            ax.set_xlim(xlim_ticks)

        if ylim == None:
            ylim=[0, 127]
        elif ylim == "Auto" or ylim == "auto":
            ylim=[self.note_range[0]-1, self.note_range[1]+1]
        ax.set_ylim(ylim)

        ax.set_xlabel("time [s]")
        ax.set_ylabel("note")

        return fig, ax, xlim_ticks

    def get_colormap_selector(self, cmap_name=None, bgcolor='white'):
        """ Define color map for each channel """
        cmap_list=(None,'Greys', 'Purples', 'Blues', 'Greens', 'Oranges', 'Reds',
            'YlOrBr', 'YlOrRd', 'OrRd', 'PuRd', 'RdPu', 'BuPu',
            'GnBu', 'PuBu', 'YlGnBu', 'PuBuGn', 'BuGn', 'YlGn')

        try:
            default_idx=cmap_list.index(cmap_name)
        except ValueError:
            default_idx=None

        cmap_name=st.sidebar.selectbox('colormap', cmap_list, index=default_idx)

        if cmap_name==None:
            transparent = colorConverter.to_rgba(bgcolor)
            colors = [
                mpl.colors.to_rgba(mpl.colors.hsv_to_rgb(
                    (i / self.nch, 1, 1)), alpha=1)
                for i in range(self.nch)
            ]
            cmaps = [
                mpl.colors.LinearSegmentedColormap.from_list(
                    'my_cmap', [transparent, colors[i]], 128)
                for i in range(self.nch)
            ]
        else:
            cmap = plt.cm.get_cmap(cmap_name)
            cmaps = [ cmap for i in range(self.nch) ]

        """
        make look up table (LUT) data, e.g., (K=3)
            array([[0. , 0. , 0. , 0. ],
                [0.5, 0. , 0. , 0.2],
                [1. , 0. , 0. , 0.4],
                [0. , 0. , 0. , 0.6],
                [1. , 0. , 0. , 0.8],
                [0. , 0. , 0. , 1. ]])
            The first 3 rows are colormap, and the last 3 rows are the colours
            for data low and high out-of-range values and for masked values.
            https://stackoverflow.com/questions/18035411/meaning-of-the-colormap-lut-list-in-matplotlib-color
        """
        for i in range(self.nch):
            cmaps[i]._init()
            alphas = np.linspace(0, 1, cmaps[i].N + 3) # about 3 extra rows, see the example above
            cmaps[i]._lut[:, -1] = alphas

        return cmaps

    def get_bgcolor_slider(self, bgcolor='white'):
        bgcolors=('white','black')
        default_idx=bgcolors.index(bgcolor)
        bgcolor=st.sidebar.selectbox('background color', bgcolors, index=default_idx)
        return bgcolor

    def get_xlim_slider(self,xlim):
        if xlim == None:
            xlim = [0, int(self.length_seconds)]

        xlim = st.sidebar.slider('Time range [s]: ', min_value=0, max_value=int(
            self.length_seconds), value=(xlim[0], xlim[1]))

        return xlim

    def draw_roll(self, figsize=(15, 9), xlim=None, ylim=None, cmaps=None, bgcolor='white', vlines=None, hlines=False, colorbar=False):
        """Create piano roll image.

        Args:
            figsize (tuple or list): figure size

            xlim: Time range to be displayed [s]
                None (not specified) : Full range
                tuple or list : (xmin, xmax) [s]

            ylim: Range of notes to be displayed in vertical axis
                None (not specified) : Full range of notes
                "Auto" or "auto" : automatic range adjustment
                tuple or list : range of notes to be displayed [s]

            bgcolor (string): name of background color

            colorbar (boolean): enable colorbar of intensity
        """

        if xlim == None:
            xlim = [0,int(self.length_seconds)]

        fig, ax1, xlim_ticks = self._grp_init(
            figsize=figsize, xlim=xlim, ylim=ylim, bgcolor=bgcolor)

        if cmaps == None:
            self.get_colormap_selector('Purple')

        for i in range(self.nch):
            try:
                target_roll = self.roll[i, :, :int(xlim_ticks[1])]
                #target_roll = self.roll[i, :, :]

                max_intensity = np.max(np.array(target_roll))
                #print("max_intensity:", max_intensity)
                im = ax1.imshow(self.roll[i], origin="lower",
                                interpolation='nearest', cmap=cmaps[i], aspect='auto', clim=[0, max_intensity])
                if hlines != False:
                    ax1.hlines([12*(i+1) for i in range(9)], xlim_ticks[0],
                           xlim_ticks[1], colors='b', linewidth=1, linestyles='dotted')
                if vlines != None:
                    ax1.vlines(np.array(vlines)*self.xticks_per_sec,
                               ylim[0], ylim[1], color='r', linewidth=1, colors='b')
                if colorbar:
                    fig.colorbar(im)
            except IndexError:
                pass

        # draw color bar for channel color
        # if colorbar:
        #     cmap = mpl.colors.LinearSegmentedColormap.from_list(
        #         'my_cmap', colors, self.nch)
        #     ax2 = fig.add_axes([0.1, 0.9, 0.8, 0.05])
        #     cbar = mpl.colorbar.ColorbarBase(ax2, cmap=cmap,
        #                                      orientation='horizontal',
        #                                      ticks=list(range(self.nch)))

        #ax1.set_title(self.fpath.name)
        #plt.draw()
        #plt.ion()
        with st.container():
            st.pyplot(fig)
        plt.savefig("outputs/"+self.fpath.name+".png", bbox_inches="tight")
        #plt.show(block=True)

def get_dirs(folder_path):
    dirs = [ f for f in os.listdir(folder_path) if os.path.isdir(folder_path+"/"+f) ]
    return (sorted(dirs))

# def file_selector(folder_path):
#     filenames = os.listdir(folder_path)
#     selected_filename = st.selectbox('Select a file', filenames)
#     return os.path.join(folder_path, selected_filename)

def show_wav(file):
    wav,sr = librosa.load(file)
    wav_seconds=int(len(wav)/sr)
    st.sidebar.write('## audio file')
    st.sidebar.write('sampling rate [Hz]: ', sr)
    st.audio(file)

def main():
    dir = "data/pedb2_v0.0.1.b/"
    #target="bac-inv001-o-p2"
    target = "bac-wtc101-p-a-p1"

    st.set_page_config(layout='wide')

    dirs = get_dirs(dir)
    target = st.sidebar.selectbox('Select file to visualize', dirs)

    st.write(target)

    path_wav = "{0}/{1}/{1}.wav".format(dir, target)
    show_wav(path_wav)

    path_mid = "{0}/{1}/{1}.mid".format(dir, target)
    mid = MidiFile(path_mid)

    path_pdf = "{0}/{1}/{1}.pdf".format(dir, target)
    st.sidebar.write('[PDF]({})'.format(path_pdf))

    # events = mid.get_events()
    # roll = mid.get_roll(verbose=False)

    st.sidebar.write("## parameters")
    # cmap_list: colormap name
    # https://matplotlib.org/stable/tutorials/colors/colormaps.html
    bgcolor = mid.get_bgcolor_slider(bgcolor='white')
    cmaps = mid.get_colormap_selector(cmap_name='Purples',bgcolor=bgcolor)
    hlines=st.sidebar.checkbox('Draw lines on C')

    xlim=[0,4]
    xlim=mid.get_xlim_slider(xlim)

    params={
        'figsize': (20, 4),
        'ylim': [30,92],
        'cmaps': cmaps,
        'bgcolor': bgcolor,
        'vlines': xlim,
        'hlines': hlines,
        'colorbar': None
    }
    mid.draw_roll(xlim=None, **params)
    mid.draw_roll(xlim=xlim, **params)

    st.sidebar.write("## MIDI database")
    st.sidebar.write("[PEDB: Music Performance Expression with Phrase Structure](https://crestmuse.jp/pedb_edition2/)")
if __name__ == "__main__":
    main()