"""Base class for sparse matrix formats using compressed storage."""
__all__ = []

from warnings import warn
import operator

import numpy as np
from scipy._lib._util import _prune_array

from .base import spmatrix, isspmatrix, SparseEfficiencyWarning
from .data import _data_matrix, _minmax_mixin
from .dia import dia_matrix
from . import _sparsetools
from ._sparsetools import (get_csr_submatrix, csr_sample_offsets, csr_todense,
                           csr_sample_values, csr_row_index, csr_row_slice,
                           csr_column_index1, csr_column_index2)
from ._index import IndexMixin
from .sputils import (upcast, upcast_char, to_native, isdense, isshape,
                      getdtype, isscalarlike, isintlike, get_index_dtype,
                      downcast_intp_index, get_sum_dtype, check_shape,
                      matrix, asmatrix, is_pydata_spmatrix)

class _cs_matrix(_data_matrix, _minmax_mixin, IndexMixin):
    """base matrix class for compressed row- and column-oriented matrices"""

    def __init__(self, arg1, shape=None, dtype=None, copy=False):
        _data_matrix.__init__(self)

        if isspmatrix(arg1):
            if arg1.format == self.format and copy:
                arg1 = arg1.copy()
            else:
                arg1 = arg1.asformat(self.format)
            self._set_self(arg1)

        elif isinstance(arg1, tuple):
            if isshape(arg1):
                # It's a tuple of matrix dimensions (M, N)
                # create empty matrix
                self._shape = check_shape(arg1)
                M, N = self.shape
                # Select index dtype large enough to pass array and
                # scalar parameters to sparsetools
                idx_dtype = get_index_dtype(maxval=max(M, N))
                self.data = np.zeros(0, getdtype(dtype, default=float))
                self.indices = np.zeros(0, idx_dtype)
                self.indptr = np.zeros(self._swap((M, N))[0] + 1,
                                       dtype=idx_dtype)
            else:
                if len(arg1) == 2:
                    # (data, ij) format
                    from .coo import coo_matrix
                    other = self.__class__(coo_matrix(arg1, shape=shape))
                    self._set_self(other)
                elif len(arg1) == 3:
                    # (data, indices, indptr) format
                    (data, indices, indptr) = arg1

                    # Select index dtype large enough to pass array and
                    # scalar parameters to sparsetools
                    maxval = None
                    if shape is not None:
                        maxval = max(shape)
                    idx_dtype = get_index_dtype((indices, indptr),
                                                maxval=maxval,
                                                check_contents=True)

                    self.indices = np.array(indices, copy=copy,
                                            dtype=idx_dtype)
                    self.indptr = np.array(indptr, copy=copy, dtype=idx_dtype)
                    self.data = np.array(data, copy=copy, dtype=dtype)
                else:
                    raise ValueError("unrecognized {}_matrix "
                                     "constructor usage".format(self.format))

        else:
            # must be dense
            try:
                arg1 = np.asarray(arg1)
            except Exception as e:
                raise ValueError("unrecognized {}_matrix constructor usage"
                                 "".format(self.format)) from e
            from .coo import coo_matrix
            self._set_self(self.__class__(coo_matrix(arg1, dtype=dtype)))

        # Read matrix dimensions given, if any
        if shape is not None:
            self._shape = check_shape(shape)
        else:
            if self.shape is None:
                # shape not already set, try to infer dimensions
                try:
                    major_dim = len(self.indptr) - 1
                    minor_dim = self.indices.max() + 1
                except Exception as e:
                    raise ValueError('unable to infer matrix dimensions') from e
                else:
                    self._shape = check_shape(self._swap((major_dim,
                                                          minor_dim)))

        if dtype is not None:
            self.data = self.data.astype(dtype, copy=False)

        self.check_format(full_check=False)

    def getnnz(self, axis=None):
        if axis is None:
            return int(self.indptr[-1])
        else:
            if axis < 0:
                axis += 2
            axis, _ = self._swap((axis, 1 - axis))
            _, N = self._swap(self.shape)
            if axis == 0:
                return np.bincount(downcast_intp_index(self.indices),
                                   minlength=N)
            elif axis == 1:
                return np.diff(self.indptr)
            raise ValueError('axis out of bounds')

    getnnz.__doc__ = spmatrix.getnnz.__doc__

    def _set_self(self, other, copy=False):
        """take the member variables of other and assign them to self"""

        if copy:
            other = other.copy()

        self.data = other.data
        self.indices = other.indices
        self.indptr = other.indptr
        self._shape = check_shape(other.shape)

    def check_format(self, full_check=True):
        """check whether the matrix format is valid

        Parameters
        ----------
        full_check : bool, optional
            If `True`, rigorous check, O(N) operations. Otherwise
            basic check, O(1) operations (default True).
        """
        # use _swap to determine proper bounds
        major_name, minor_name = self._swap(('row', 'column'))
        major_dim, minor_dim = self._swap(self.shape)

        # index arrays should have integer data types
        if self.indptr.dtype.kind != 'i':
            warn("indptr array has non-integer dtype ({})"
                 "".format(self.indptr.dtype.name), stacklevel=3)
        if self.indices.dtype.kind != 'i':
            warn("indices array has non-integer dtype ({})"
                 "".format(self.indices.dtype.name), stacklevel=3)

        idx_dtype = get_index_dtype((self.indptr, self.indices))
        self.indptr = np.asarray(self.indptr, dtype=idx_dtype)
        self.indices = np.asarray(self.indices, dtype=idx_dtype)
        self.data = to_native(self.data)

        # check array shapes
        for x in [self.data.ndim, self.indices.ndim, self.indptr.ndim]:
            if x != 1:
                raise ValueError('data, indices, and indptr should be 1-D')

        # check index pointer
        if (len(self.indptr) != major_dim + 1):
            raise ValueError("index pointer size ({}) should be ({})"
                             "".format(len(self.indptr), major_dim + 1))
        if (self.indptr[0] != 0):
            raise ValueError("index pointer should start with 0")

        # check index and data arrays
        if (len(self.indices) != len(self.data)):
            raise ValueError("indices and data should have the same size")
        if (self.indptr[-1] > len(self.indices)):
            raise ValueError("Last value of index pointer should be less than "
                             "the size of index and data arrays")

        self.prune()

        if full_check:
            # check format validity (more expensive)
            if self.nnz > 0:
                if self.indices.max() >= minor_dim:
                    raise ValueError("{} index values must be < {}"
                                     "".format(minor_name, minor_dim))
                if self.indices.min() < 0:
                    raise ValueError("{} index values must be >= 0"
                                     "".format(minor_name))
                if np.diff(self.indptr).min() < 0:
                    raise ValueError("index pointer values must form a "
                                     "non-decreasing sequence")

        # if not self.has_sorted_indices():
        #    warn('Indices were not in sorted order.  Sorting indices.')
        #    self.sort_indices()
        #    assert(self.has_sorted_indices())
        # TODO check for duplicates?

    #######################
    # Boolean comparisons #
    #######################

    def _scalar_binopt(self, other, op):
        """Scalar version of self._binopt, for cases in which no new nonzeros
        are added. Produces a new spmatrix in canonical form.
        """
        self.sum_duplicates()
        res = self._with_data(op(self.data, other), copy=True)
        res.eliminate_zeros()
        return res

    def __eq__(self, other):
        # Scalar other.
        if isscalarlike(other):
            if np.isnan(other):
                return self.__class__(self.shape, dtype=np.bool_)

            if other == 0:
                warn("Comparing a sparse matrix with 0 using == is inefficient"
                     ", try using != instead.", SparseEfficiencyWarning,
                     stacklevel=3)
                all_true = self.__class__(np.ones(self.shape, dtype=np.bool_))
                inv = self._scalar_binopt(other, operator.ne)
                return all_true - inv
            else:
                return self._scalar_binopt(other, operator.eq)
        # Dense other.
        elif isdense(other):
            return self.todense() == other
        # Pydata sparse other.
        elif is_pydata_spmatrix(other):
            return NotImplemented
        # Sparse other.
        elif isspmatrix(other):
            warn("Comparing sparse matrices using == is inefficient, try using"
                 " != instead.", SparseEfficiencyWarning, stacklevel=3)
            # TODO sparse broadcasting
            if self.shape != other.shape:
                return False
            elif self.format != other.format:
                other = other.asformat(self.format)
            res = self._binopt(other, '_ne_')
            all_true = self.__class__(np.ones(self.shape, dtype=np.bool_))
            return all_true - res
        else:
            return False

    def __ne__(self, other):
        # Scalar other.
        if isscalarlike(other):
            if np.isnan(other):
                warn("Comparing a sparse matrix with nan using != is"
                     " inefficient", SparseEfficiencyWarning, stacklevel=3)
                all_true = self.__class__(np.ones(self.shape, dtype=np.bool_))
                return all_true
            elif other != 0:
                warn("Comparing a sparse matrix with a nonzero scalar using !="
                     " is inefficient, try using == instead.",
                     SparseEfficiencyWarning, stacklevel=3)
                all_true = self.__class__(np.ones(self.shape), dtype=np.bool_)
                inv = self._scalar_binopt(other, operator.eq)
                return all_true - inv
            else:
                return self._scalar_binopt(other, operator.ne)
        # Dense other.
        elif isdense(other):
            return self.todense() != other
        # Pydata sparse other.
        elif is_pydata_spmatrix(other):
            return NotImplemented
        # Sparse other.
        elif isspmatrix(other):
            # TODO sparse broadcasting
            if self.shape != other.shape:
                return True
            elif self.format != other.format:
                other = other.asformat(self.format)
            return self._binopt(other, '_ne_')
        else:
            return True

    def _inequality(self, other, op, op_name, bad_scalar_msg):
        # Scalar other.
        if isscalarlike(other):
            if 0 == other and op_name in ('_le_', '_ge_'):
                raise NotImplementedError(" >= and <= don't work with 0.")
            elif op(0, other):
                warn(bad_scalar_msg, SparseEfficiencyWarning)
                other_arr = np.empty(self.shape, dtype=np.result_type(other))
                other_arr.fill(other)
                other_arr = self.__class__(other_arr)
                return self._binopt(other_arr, op_name)
            else:
                return self._scalar_binopt(other, op)
        # Dense other.
        elif isdense(other):
            return op(self.todense(), other)
        # Sparse other.
        elif isspmatrix(other):
            # TODO sparse broadcasting
            if self.shape != other.shape:
                raise ValueError("inconsistent shapes")
            elif self.format != other.format:
                other = other.asformat(self.format)
            if op_name not in ('_ge_', '_le_'):
                return self._binopt(other, op_name)

            warn("Comparing sparse matrices using >= and <= is inefficient, "
                 "using <, >, or !=, instead.", SparseEfficiencyWarning)
            all_true = self.__class__(np.ones(self.shape, dtype=np.bool_))
            res = self._binopt(other, '_gt_' if op_name == '_le_' else '_lt_')
            return all_true - res
        else:
            raise ValueError("Operands could not be compared.")

    def __lt__(self, other):
        return self._inequality(other, operator.lt, '_lt_',
                                "Comparing a sparse matrix with a scalar "
                                "greater than zero using < is inefficient, "
                                "try using >= instead.")

    def __gt__(self, other):
        return self._inequality(other, operator.gt, '_gt_',
                                "Comparing a sparse matrix with a scalar "
                                "less than zero using > is inefficient, "
                                "try using <= instead.")

    def __le__(self, other):
        return self._inequality(other, operator.le, '_le_',
                                "Comparing a sparse matrix with a scalar "
                                "greater than zero using <= is inefficient, "
                                "try using > instead.")

    def __ge__(self, other):
        return self._inequality(other, operator.ge, '_ge_',
                                "Comparing a sparse matrix with a scalar "
                                "less than zero using >= is inefficient, "
                                "try using < instead.")

    #################################
    # Arithmetic operator overrides #
    #################################

    def _add_dense(self, other):
        if other.shape != self.shape:
            raise ValueError('Incompatible shapes ({} and {})'
                             .format(self.shape, other.shape))
        dtype = upcast_char(self.dtype.char, other.dtype.char)
        order = self._swap('CF')[0]
        result = np.array(other, dtype=dtype, order=order, copy=True)
        M, N = self._swap(self.shape)
        y = result if result.flags.c_contiguous else result.T
        csr_todense(M, N, self.indptr, self.indices, self.data, y)
        return matrix(result, copy=False)

    def _add_sparse(self, other):
        return self._binopt(other, '_plus_')

    def _sub_sparse(self, other):
        return self._binopt(other, '_minus_')

    def multiply(self, other):
        """Point-wise multiplication by another matrix, vector, or
        scalar.
        """
        # Scalar multiplication.
        if isscalarlike(other):
            return self._mul_scalar(other)
        # Sparse matrix or vector.
        if isspmatrix(other):
            if self.shape == other.shape:
                other = self.__class__(other)
                return self._binopt(other, '_elmul_')
            # Single element.
            elif other.shape == (1, 1):
                return self._mul_scalar(other.toarray()[0, 0])
            elif self.shape == (1, 1):
                return other._mul_scalar(self.toarray()[0, 0])
            # A row times a column.
            elif self.shape[1] == 1 and other.shape[0] == 1:
                return self._mul_sparse_matrix(other.tocsc())
            elif self.shape[0] == 1 and other.shape[1] == 1:
                return other._mul_sparse_matrix(self.tocsc())
            # Row vector times matrix. other is a row.
            elif other.shape[0] == 1 and self.shape[1] == other.shape[1]:
                other = dia_matrix((other.toarray().ravel(), [0]),
                                   shape=(other.shape[1], other.shape[1]))
                return self._mul_sparse_matrix(other)
            # self is a row.
            elif self.shape[0] == 1 and self.shape[1] == other.shape[1]:
                copy = dia_matrix((self.toarray().ravel(), [0]),
                                  shape=(self.shape[1], self.shape[1]))
                return other._mul_sparse_matrix(copy)
            # Column vector times matrix. other is a column.
            elif other.shape[1] == 1 and self.shape[0] == other.shape[0]:
                other = dia_matrix((other.toarray().ravel(), [0]),
                                   shape=(other.shape[0], other.shape[0]))
                return other._mul_sparse_matrix(self)
            # self is a column.
            elif self.shape[1] == 1 and self.shape[0] == other.shape[0]:
                copy = dia_matrix((self.toarray().ravel(), [0]),
                                  shape=(self.shape[0], self.shape[0]))
                return copy._mul_sparse_matrix(other)
            else:
                raise ValueError("inconsistent shapes")

        # Assume other is a dense matrix/array, which produces a single-item
        # object array if other isn't convertible to ndarray.
        other = np.atleast_2d(other)

        if other.ndim != 2:
            return np.multiply(self.toarray(), other)
        # Single element / wrapped object.
        if other.size == 1:
            return self._mul_scalar(other.flat[0])
        # Fast case for trivial sparse matrix.
        elif self.shape == (1, 1):
            return np.multiply(self.toarray()[0, 0], other)

        from .coo import coo_matrix
        ret = self.tocoo()
        # Matching shapes.
        if self.shape == other.shape:
            data = np.multiply(ret.data, other[ret.row, ret.col])
        # Sparse row vector times...
        elif self.shape[0] == 1:
            if other.shape[1] == 1:  # Dense column vector.
                data = np.multiply(ret.data, other)
            elif other.shape[1] == self.shape[1]:  # Dense matrix.
                data = np.multiply(ret.data, other[:, ret.col])
            else:
                raise ValueError("inconsistent shapes")
            row = np.repeat(np.arange(other.shape[0]), len(ret.row))
            col = np.tile(ret.col, other.shape[0])
            return coo_matrix((data.view(np.ndarray).ravel(), (row, col)),
                              shape=(other.shape[0], self.shape[1]),
                              copy=False)
        # Sparse column vector times...
        elif self.shape[1] == 1:
            if other.shape[0] == 1:  # Dense row vector.
                data = np.multiply(ret.data[:, None], other)
            elif other.shape[0] == self.shape[0]:  # Dense matrix.
                data = np.multiply(ret.data[:, None], other[ret.row])
            else:
                raise ValueError("inconsistent shapes")
            row = np.repeat(ret.row, other.shape[1])
            col = np.tile(np.arange(other.shape[1]), len(ret.col))
            return coo_matrix((data.view(np.ndarray).ravel(), (row, col)),
                              shape=(self.shape[0], other.shape[1]),
                              copy=False)
        # Sparse matrix times dense row vector.
        elif other.shape[0] == 1 and self.shape[1] == other.shape[1]:
            data = np.multiply(ret.data, other[:, ret.col].ravel())
        # Sparse matrix times dense column vector.
        elif other.shape[1] == 1 and self.shape[0] == other.shape[0]:
            data = np.multiply(ret.data, other[ret.row].ravel())
        else:
            raise ValueError("inconsistent shapes")
        ret.data = data.view(np.ndarray).ravel()
        return ret

    ###########################
    # Multiplication handlers #
    ###########################

    def _mul_vector(self, other):
        M, N = self.shape

        # output array
        result = np.zeros(M, dtype=upcast_char(self.dtype.char,
                                               other.dtype.char))

        # csr_matvec or csc_matvec
        fn = getattr(_sparsetools, self.format + '_matvec')
        fn(M, N, self.indptr, self.indices, self.data, other, result)

        return result

    def _mul_multivector(self, other):
        M, N = self.shape
        n_vecs = other.shape[1]  # number of column vectors

        result = np.zeros((M, n_vecs),
                          dtype=upcast_char(self.dtype.char, other.dtype.char))

        # csr_matvecs or csc_matvecs
        fn = getattr(_sparsetools, self.format + '_matvecs')
        fn(M, N, n_vecs, self.indptr, self.indices, self.data,
           other.ravel(), result.ravel())

        return result

    def _mul_sparse_matrix(self, other):
        M, K1 = self.shape
        K2, N = other.shape

        major_axis = self._swap((M, N))[0]
        other = self.__class__(other)  # convert to this format

        idx_dtype = get_index_dtype((self.indptr, self.indices,
                                     other.indptr, other.indices))

        fn = getattr(_sparsetools, self.format + '_matmat_maxnnz')
        nnz = fn(M, N,
                 np.asarray(self.indptr, dtype=idx_dtype),
                 np.asarray(self.indices, dtype=idx_dtype),
                 np.asarray(other.indptr, dtype=idx_dtype),
                 np.asarray(other.indices, dtype=idx_dtype))

        idx_dtype = get_index_dtype((self.indptr, self.indices,
                                     other.indptr, other.indices),
                                    maxval=nnz)

        indptr = np.empty(major_axis + 1, dtype=idx_dtype)
        indices = np.empty(nnz, dtype=idx_dtype)
        data = np.empty(nnz, dtype=upcast(self.dtype, other.dtype))

        fn = getattr(_sparsetools, self.format + '_matmat')
        fn(M, N, np.asarray(self.indptr, dtype=idx_dtype),
           np.asarray(self.indices, dtype=idx_dtype),
           self.data,
           np.asarray(other.indptr, dtype=idx_dtype),
           np.asarray(other.indices, dtype=idx_dtype),
           other.data,
           indptr, indices, data)

        return self.__class__((data, indices, indptr), shape=(M, N))

    def diagonal(self, k=0):
        rows, cols = self.shape
        if k <= -rows or k >= cols:
            return np.empty(0, dtype=self.data.dtype)
        fn = getattr(_sparsetools, self.format + "_diagonal")
        y = np.empty(min(rows + min(k, 0), cols - max(k, 0)),
                     dtype=upcast(self.dtype))
        fn(k, self.shape[0], self.shape[1], self.indptr, self.indices,
           self.data, y)
        return y

    diagonal.__doc__ = spmatrix.diagonal.__doc__

    #####################
    # Other binary ops  #
    #####################

    def _maximum_minimum(self, other, npop, op_name, dense_check):
        if isscalarlike(other):
            if dense_check(other):
                warn("Taking maximum (minimum) with > 0 (< 0) number results"
                     " to a dense matrix.", SparseEfficiencyWarning,
                     stacklevel=3)
                other_arr = np.empty(self.shape, dtype=np.asarray(other).dtype)
                other_arr.fill(other)
                other_arr = self.__class__(other_arr)
                return self._binopt(other_arr, op_name)
            else:
                self.sum_duplicates()
                new_data = npop(self.data, np.asarray(other))
                mat = self.__class__((new_data, self.indices, self.indptr),
                                     dtype=new_data.dtype, shape=self.shape)
                return mat
        elif isdense(other):
            return npop(self.todense(), other)
        elif isspmatrix(other):
            return self._binopt(other, op_name)
        else:
            raise ValueError("Operands not compatible.")

    def maximum(self, other):
        return self._maximum_minimum(other, np.maximum,
                                     '_maximum_', lambda x: np.asarray(x) > 0)

    maximum.__doc__ = spmatrix.maximum.__doc__

    def minimum(self, other):
        return self._maximum_minimum(other, np.minimum,
                                     '_minimum_', lambda x: np.asarray(x) < 0)

    minimum.__doc__ = spmatrix.minimum.__doc__

    #####################
    # Reduce operations #
    #####################

    def sum(self, axis=None, dtype=None, out=None):
        """Sum the matrix over the given axis.  If the axis is None, sum
        over both rows and columns, returning a scalar.
        """
        # The spmatrix base class already does axis=0 and axis=1 efficiently
        # so we only do the case axis=None here
        if (not hasattr(self, 'blocksize') and
                axis in self._swap(((1, -1), (0, 2)))[0]):
            # faster than multiplication for large minor axis in CSC/CSR
            res_dtype = get_sum_dtype(self.dtype)
            ret = np.zeros(len(self.indptr) - 1, dtype=res_dtype)

            major_index, value = self._minor_reduce(np.add)
            ret[major_index] = value
            ret = asmatrix(ret)
            if axis % 2 == 1:
                ret = ret.T

            if out is not None and out.shape != ret.shape:
                raise ValueError('dimensions do not match')

            return ret.sum(axis=(), dtype=dtype, out=out)
        # spmatrix will handle the remaining situations when axis
        # is in {None, -1, 0, 1}
        else:
            return spmatrix.sum(self, axis=axis, dtype=dtype, out=out)

    sum.__doc__ = spmatrix.sum.__doc__

    def _minor_reduce(self, ufunc, data=None):
        """Reduce nonzeros with a ufunc over the minor axis when non-empty

        Can be applied to a function of self.data by supplying data parameter.

        Warning: this does not call sum_duplicates()

        Returns
        -------
        major_index : array of ints
            Major indices where nonzero

        value : array of self.dtype
            Reduce result for nonzeros in each major_index
        """
        if data is None:
            data = self.data
        major_index = np.flatnonzero(np.diff(self.indptr))
        value = ufunc.reduceat(data,
                               downcast_intp_index(self.indptr[major_index]))
        return major_index, value

    #######################
    # Getting and Setting #
    #######################

    def _get_intXint(self, row, col):
        M, N = self._swap(self.shape)
        major, minor = self._swap((row, col))
        indptr, indices, data = get_csr_submatrix(
            M, N, self.indptr, self.indices, self.data,
            major, major + 1, minor, minor + 1)
        return data.sum(dtype=self.dtype)

    def _get_sliceXslice(self, row, col):
        major, minor = self._swap((row, col))
        if major.step in (1, None) and minor.step in (1, None):
            return self._get_submatrix(major, minor, copy=True)
        return self._major_slice(major)._minor_slice(minor)

    def _get_arrayXarray(self, row, col):
        # inner indexing
        idx_dtype = self.indices.dtype
        M, N = self._swap(self.shape)
        major, minor = self._swap((row, col))
        major = np.asarray(major, dtype=idx_dtype)
        minor = np.asarray(minor, dtype=idx_dtype)

        val = np.empty(major.size, dtype=self.dtype)
        csr_sample_values(M, N, self.indptr, self.indices, self.data,
                          major.size, major.ravel(), minor.ravel(), val)
        if major.ndim == 1:
            return asmatrix(val)
        return self.__class__(val.reshape(major.shape))

    def _get_columnXarray(self, row, col):
        # outer indexing
        major, minor = self._swap((row, col))
        return self._major_index_fancy(major)._minor_index_fancy(minor)

    def _major_index_fancy(self, idx):
        """Index along the major axis where idx is an array of ints.
        """
        idx_dtype = self.indices.dtype
        indices = np.asarray(idx, dtype=idx_dtype).ravel()

        _, N = self._swap(self.shape)
        M = len(indices)
        new_shape = self._swap((M, N))
        if M == 0:
            return self.__class__(new_shape)

        row_nnz = np.diff(self.indptr)
        idx_dtype = self.indices.dtype
        res_indptr = np.zeros(M+1, dtype=idx_dtype)
        np.cumsum(row_nnz[idx], out=res_indptr[1:])

        nnz = res_indptr[-1]
        res_indices = np.empty(nnz, dtype=idx_dtype)
        res_data = np.empty(nnz, dtype=self.dtype)
        csr_row_index(M, indices, self.indptr, self.indices, self.data,
                      res_indices, res_data)

        return self.__class__((res_data, res_indices, res_indptr),
                              shape=new_shape, copy=False)

    def _major_slice(self, idx, copy=False):
        """Index along the major axis where idx is a slice object.
        """
        if idx == slice(None):
            return self.copy() if copy else self

        M, N = self._swap(self.shape)
        start, stop, step = idx.indices(M)
        M = len(range(start, stop, step))
        new_shape = self._swap((M, N))
        if M == 0:
            return self.__class__(new_shape)

        row_nnz = np.diff(self.indptr)
        idx_dtype = self.indices.dtype
        res_indptr = np.zeros(M+1, dtype=idx_dtype)
        np.cumsum(row_nnz[idx], out=res_indptr[1:])

        if step == 1:
            all_idx = slice(self.indptr[start], self.indptr[stop])
            res_indices = np.array(self.indices[all_idx], copy=copy)
            res_data = np.array(self.data[all_idx], copy=copy)
        else:
            nnz = res_indptr[-1]
            res_indices = np.empty(nnz, dtype=idx_dtype)
            res_data = np.empty(nnz, dtype=self.dtype)
            csr_row_slice(start, stop, step, self.indptr, self.indices,
                          self.data, res_indices, res_data)

        return self.__class__((res_data, res_indices, res_indptr),
                              shape=new_shape, copy=False)

    def _minor_index_fancy(self, idx):
        """Index along the minor axis where idx is an array of ints.
        """
        idx_dtype = self.indices.dtype
        idx = np.asarray(idx, dtype=idx_dtype).ravel()

        M, N = self._swap(self.shape)
        k = len(idx)
        new_shape = self._swap((M, k))
        if k == 0:
            return self.__class__(new_shape)

        # pass 1: count idx entries and compute new indptr
        col_offsets = np.zeros(N, dtype=idx_dtype)
        res_indptr = np.empty_like(self.indptr)
        csr_column_index1(k, idx, M, N, self.indptr, self.indices,
                          col_offsets, res_indptr)

        # pass 2: copy indices/data for selected idxs
        col_order = np.argsort(idx).astype(idx_dtype, copy=False)
        nnz = res_indptr[-1]
        res_indices = np.empty(nnz, dtype=idx_dtype)
        res_data = np.empty(nnz, dtype=self.dtype)
        csr_column_index2(col_order, col_offsets, len(self.indices),
                          self.indices, self.data, res_indices, res_data)
        return self.__class__((res_data, res_indices, res_indptr),
                              shape=new_shape, copy=False)

    def _minor_slice(self, idx, copy=False):
        """Index along the minor axis where idx is a slice object.
        """
        if idx == slice(None):
            return self.copy() if copy else self

        M, N = self._swap(self.shape)
        start, stop, step = idx.indices(N)
        N = len(range(start, stop, step))
        if N == 0:
            return self.__class__(self._swap((M, N)))
        if step == 1:
            return self._get_submatrix(minor=idx, copy=copy)
        # TODO: don't fall back to fancy indexing here
        return self._minor_index_fancy(np.arange(start, stop, step))

    def _get_submatrix(self, major=None, minor=None, copy=False):
        """Return a submatrix of this matrix.

        major, minor: None, int, or slice with step 1
        """
        M, N = self._swap(self.shape)
        i0, i1 = _process_slice(major, M)
        j0, j1 = _process_slice(minor, N)

        if i0 == 0 and j0 == 0 and i1 == M and j1 == N:
            return self.copy() if copy else self

        indptr, indices, data = get_csr_submatrix(
            M, N, self.indptr, self.indices, self.data, i0, i1, j0, j1)

        shape = self._swap((i1 - i0, j1 - j0))
        return self.__class__((data, indices, indptr), shape=shape,
                              dtype=self.dtype, copy=False)

    def _set_intXint(self, row, col, x):
        i, j = self._swap((row, col))
        self._set_many(i, j, x)

    def _set_arrayXarray(self, row, col, x):
        i, j = self._swap((row, col))
        self._set_many(i, j, x)

    def _set_arrayXarray_sparse(self, row, col, x):
        # clear entries that will be overwritten
        self._zero_many(*self._swap((row, col)))

        M, N = row.shape  # matches col.shape
        broadcast_row = M != 1 and x.shape[0] == 1
        broadcast_col = N != 1 and x.shape[1] == 1
        r, c = x.row, x.col
        x = np.asarray(x.data, dtype=self.dtype)
        if broadcast_row:
            r = np.repeat(np.arange(M), len(r))
            c = np.tile(c, M)
            x = np.tile(x, M)
        if broadcast_col:
            r = np.repeat(r, N)
            c = np.tile(np.arange(N), len(c))
            x = np.repeat(x, N)
        # only assign entries in the new sparsity structure
        i, j = self._swap((row[r, c], col[r, c]))
        self._set_many(i, j, x)

    def _setdiag(self, values, k):
        if 0 in self.shape:
            return

        M, N = self.shape
        broadcast = (values.ndim == 0)

        if k < 0:
            if broadcast:
                max_index = min(M + k, N)
            else:
                max_index = min(M + k, N, len(values))
            i = np.arange(max_index, dtype=self.indices.dtype)
            j = np.arange(max_index, dtype=self.indices.dtype)
            i -= k

        else:
            if broadcast:
                max_index = min(M, N - k)
            else:
                max_index = min(M, N - k, len(values))
            i = np.arange(max_index, dtype=self.indices.dtype)
            j = np.arange(max_index, dtype=self.indices.dtype)
            j += k

        if not broadcast:
            values = values[:len(i)]

        self[i, j] = values

    def _prepare_indices(self, i, j):
        M, N = self._swap(self.shape)

        def check_bounds(indices, bound):
            idx = indices.max()
            if idx >= bound:
                raise IndexError('index (%d) out of range (>= %d)' %
                                 (idx, bound))
            idx = indices.min()
            if idx < -bound:
                raise IndexError('index (%d) out of range (< -%d)' %
                                 (idx, bound))

        i = np.array(i, dtype=self.indices.dtype, copy=False, ndmin=1).ravel()
        j = np.array(j, dtype=self.indices.dtype, copy=False, ndmin=1).ravel()
        check_bounds(i, M)
        check_bounds(j, N)
        return i, j, M, N

    def _set_many(self, i, j, x):
        """Sets value at each (i, j) to x

        Here (i,j) index major and minor respectively, and must not contain
        duplicate entries.
        """
        i, j, M, N = self._prepare_indices(i, j)
        x = np.array(x, dtype=self.dtype, copy=False, ndmin=1).ravel()

        n_samples = x.size
        offsets = np.empty(n_samples, dtype=self.indices.dtype)
        ret = csr_sample_offsets(M, N, self.indptr, self.indices, n_samples,
                                 i, j, offsets)
        if ret == 1:
            # rinse and repeat
            self.sum_duplicates()
            csr_sample_offsets(M, N, self.indptr, self.indices, n_samples,
                               i, j, offsets)

        if -1 not in offsets:
            # only affects existing non-zero cells
            self.data[offsets] = x
            return

        else:
            warn("Changing the sparsity structure of a {}_matrix is expensive."
                 " lil_matrix is more efficient.".format(self.format),
                 SparseEfficiencyWarning, stacklevel=3)
            # replace where possible
            mask = offsets > -1
            self.data[offsets[mask]] = x[mask]
            # only insertions remain
            mask = ~mask
            i = i[mask]
            i[i < 0] += M
            j = j[mask]
            j[j < 0] += N
            self._insert_many(i, j, x[mask])

    def _zero_many(self, i, j):
        """Sets value at each (i, j) to zero, preserving sparsity structure.

        Here (i,j) index major and minor respectively.
        """
        i, j, M, N = self._prepare_indices(i, j)

        n_samples = len(i)
        offsets = np.empty(n_samples, dtype=self.indices.dtype)
        ret = csr_sample_offsets(M, N, self.indptr, self.indices, n_samples,
                                 i, j, offsets)
        if ret == 1:
            # rinse and repeat
            self.sum_duplicates()
            csr_sample_offsets(M, N, self.indptr, self.indices, n_samples,
                               i, j, offsets)

        # only assign zeros to the existing sparsity structure
        self.data[offsets[offsets > -1]] = 0

    def _insert_many(self, i, j, x):
        """Inserts new nonzero at each (i, j) with value x

        Here (i,j) index major and minor respectively.
        i, j and x must be non-empty, 1d arrays.
        Inserts each major group (e.g. all entries per row) at a time.
        Maintains has_sorted_indices property.
        Modifies i, j, x in place.
        """
        order = np.argsort(i, kind='mergesort')  # stable for duplicates
        i = i.take(order, mode='clip')
        j = j.take(order, mode='clip')
        x = x.take(order, mode='clip')

        do_sort = self.has_sorted_indices

        # Update index data type
        idx_dtype = get_index_dtype((self.indices, self.indptr),
                                    maxval=(self.indptr[-1] + x.size))
        self.indptr = np.asarray(self.indptr, dtype=idx_dtype)
        self.indices = np.asarray(self.indices, dtype=idx_dtype)
        i = np.asarray(i, dtype=idx_dtype)
        j = np.asarray(j, dtype=idx_dtype)

        # Collate old and new in chunks by major index
        indices_parts = []
        data_parts = []
        ui, ui_indptr = np.unique(i, return_index=True)
        ui_indptr = np.append(ui_indptr, len(j))
        new_nnzs = np.diff(ui_indptr)
        prev = 0
        for c, (ii, js, je) in enumerate(zip(ui, ui_indptr, ui_indptr[1:])):
            # old entries
            start = self.indptr[prev]
            stop = self.indptr[ii]
            indices_parts.append(self.indices[start:stop])
            data_parts.append(self.data[start:stop])

            # handle duplicate j: keep last setting
            uj, uj_indptr = np.unique(j[js:je][::-1], return_index=True)
            if len(uj) == je - js:
                indices_parts.append(j[js:je])
                data_parts.append(x[js:je])
            else:
                indices_parts.append(j[js:je][::-1][uj_indptr])
                data_parts.append(x[js:je][::-1][uj_indptr])
                new_nnzs[c] = len(uj)

            prev = ii

        # remaining old entries
        start = self.indptr[ii]
        indices_parts.append(self.indices[start:])
        data_parts.append(self.data[start:])

        # update attributes
        self.indices = np.concatenate(indices_parts)
        self.data = np.concatenate(data_parts)
        nnzs = np.empty(self.indptr.shape, dtype=idx_dtype)
        nnzs[0] = idx_dtype(0)
        indptr_diff = np.diff(self.indptr)
        indptr_diff[ui] += new_nnzs
        nnzs[1:] = indptr_diff
        self.indptr = np.cumsum(nnzs, out=nnzs)

        if do_sort:
            # TODO: only sort where necessary
            self.has_sorted_indices = False
            self.sort_indices()

        self.check_format(full_check=False)

    ######################
    # Conversion methods #
    ######################

    def tocoo(self, copy=True):
        major_dim, minor_dim = self._swap(self.shape)
        minor_indices = self.indices
        major_indices = np.empty(len(minor_indices), dtype=self.indices.dtype)
        _sparsetools.expandptr(major_dim, self.indptr, major_indices)
        row, col = self._swap((major_indices, minor_indices))

        from .coo import coo_matrix
        return coo_matrix((self.data, (row, col)), self.shape, copy=copy,
                          dtype=self.dtype)

    tocoo.__doc__ = spmatrix.tocoo.__doc__

    def toarray(self, order=None, out=None):
        if out is None and order is None:
            order = self._swap('cf')[0]
        out = self._process_toarray_args(order, out)
        if not (out.flags.c_contiguous or out.flags.f_contiguous):
            raise ValueError('Output array must be C or F contiguous')
        # align ideal order with output array order
        if out.flags.c_contiguous:
            x = self.tocsr()
            y = out
        else:
            x = self.tocsc()
            y = out.T
        M, N = x._swap(x.shape)
        csr_todense(M, N, x.indptr, x.indices, x.data, y)
        return out

    toarray.__doc__ = spmatrix.toarray.__doc__

    ##############################################################
    # methods that examine or modify the internal data structure #
    ##############################################################

    def eliminate_zeros(self):
        """Remove zero entries from the matrix

        This is an *in place* operation.
        """
        M, N = self._swap(self.shape)
        _sparsetools.csr_eliminate_zeros(M, N, self.indptr, self.indices,
                                         self.data)
        self.prune()  # nnz may have changed

    def __get_has_canonical_format(self):
        """Determine whether the matrix has sorted indices and no duplicates

        Returns
            - True: if the above applies
            - False: otherwise

        has_canonical_format implies has_sorted_indices, so if the latter flag
        is False, so will the former be; if the former is found True, the
        latter flag is also set.
        """

        # first check to see if result was cached
        if not getattr(self, '_has_sorted_indices', True):
            # not sorted => not canonical
            self._has_canonical_format = False
        elif not hasattr(self, '_has_canonical_format'):
            self.has_canonical_format = _sparsetools.csr_has_canonical_format(
                len(self.indptr) - 1, self.indptr, self.indices)
        return self._has_canonical_format

    def __set_has_canonical_format(self, val):
        self._has_canonical_format = bool(val)
        if val:
            self.has_sorted_indices = True

    has_canonical_format = property(fget=__get_has_canonical_format,
                                    fset=__set_has_canonical_format)

    def sum_duplicates(self):
        """Eliminate duplicate matrix entries by adding them together

        This is an *in place* operation.
        """
        if self.has_canonical_format:
            return
        self.sort_indices()

        M, N = self._swap(self.shape)
        _sparsetools.csr_sum_duplicates(M, N, self.indptr, self.indices,
                                        self.data)

        self.prune()  # nnz may have changed
        self.has_canonical_format = True

    def __get_sorted(self):
        """Determine whether the matrix has sorted indices

        Returns
            - True: if the indices of the matrix are in sorted order
            - False: otherwise

        """

        # first check to see if result was cached
        if not hasattr(self, '_has_sorted_indices'):
            self._has_sorted_indices = _sparsetools.csr_has_sorted_indices(
                len(self.indptr) - 1, self.indptr, self.indices)
        return self._has_sorted_indices

    def __set_sorted(self, val):
        self._has_sorted_indices = bool(val)

    has_sorted_indices = property(fget=__get_sorted, fset=__set_sorted)

    def sorted_indices(self):
        """Return a copy of this matrix with sorted indices
        """
        A = self.copy()
        A.sort_indices()
        return A

        # an alternative that has linear complexity is the following
        # although the previous option is typically faster
        # return self.toother().toother()

    def sort_indices(self):
        """Sort the indices of this matrix *in place*
        """

        if not self.has_sorted_indices:
            _sparsetools.csr_sort_indices(len(self.indptr) - 1, self.indptr,
                                          self.indices, self.data)
            self.has_sorted_indices = True

    def prune(self):
        """Remove empty space after all non-zero elements.
        """
        major_dim = self._swap(self.shape)[0]

        if len(self.indptr) != major_dim + 1:
            raise ValueError('index pointer has invalid length')
        if len(self.indices) < self.nnz:
            raise ValueError('indices array has fewer than nnz elements')
        if len(self.data) < self.nnz:
            raise ValueError('data array has fewer than nnz elements')

        self.indices = _prune_array(self.indices[:self.nnz])
        self.data = _prune_array(self.data[:self.nnz])

    def resize(self, *shape):
        shape = check_shape(shape)
        if hasattr(self, 'blocksize'):
            bm, bn = self.blocksize
            new_M, rm = divmod(shape[0], bm)
            new_N, rn = divmod(shape[1], bn)
            if rm or rn:
                raise ValueError("shape must be divisible into %s blocks. "
                                 "Got %s" % (self.blocksize, shape))
            M, N = self.shape[0] // bm, self.shape[1] // bn
        else:
            new_M, new_N = self._swap(shape)
            M, N = self._swap(self.shape)

        if new_M < M:
            self.indices = self.indices[:self.indptr[new_M]]
            self.data = self.data[:self.indptr[new_M]]
            self.indptr = self.indptr[:new_M + 1]
        elif new_M > M:
            self.indptr = np.resize(self.indptr, new_M + 1)
            self.indptr[M + 1:].fill(self.indptr[M])

        if new_N < N:
            mask = self.indices < new_N
            if not np.all(mask):
                self.indices = self.indices[mask]
                self.data = self.data[mask]
                major_index, val = self._minor_reduce(np.add, mask)
                self.indptr.fill(0)
                self.indptr[1:][major_index] = val
                np.cumsum(self.indptr, out=self.indptr)

        self._shape = shape

    resize.__doc__ = spmatrix.resize.__doc__

    ###################
    # utility methods #
    ###################

    # needed by _data_matrix
    def _with_data(self, data, copy=True):
        """Returns a matrix with the same sparsity structure as self,
        but with different data.  By default the structure arrays
        (i.e. .indptr and .indices) are copied.
        """
        if copy:
            return self.__class__((data, self.indices.copy(),
                                   self.indptr.copy()),
                                  shape=self.shape,
                                  dtype=data.dtype)
        else:
            return self.__class__((data, self.indices, self.indptr),
                                  shape=self.shape, dtype=data.dtype)

    def _binopt(self, other, op):
        """apply the binary operation fn to two sparse matrices."""
        other = self.__class__(other)

        # e.g. csr_plus_csr, csr_minus_csr, etc.
        fn = getattr(_sparsetools, self.format + op + self.format)

        maxnnz = self.nnz + other.nnz
        idx_dtype = get_index_dtype((self.indptr, self.indices,
                                     other.indptr, other.indices),
                                    maxval=maxnnz)
        indptr = np.empty(self.indptr.shape, dtype=idx_dtype)
        indices = np.empty(maxnnz, dtype=idx_dtype)

        bool_ops = ['_ne_', '_lt_', '_gt_', '_le_', '_ge_']
        if op in bool_ops:
            data = np.empty(maxnnz, dtype=np.bool_)
        else:
            data = np.empty(maxnnz, dtype=upcast(self.dtype, other.dtype))

        fn(self.shape[0], self.shape[1],
           np.asarray(self.indptr, dtype=idx_dtype),
           np.asarray(self.indices, dtype=idx_dtype),
           self.data,
           np.asarray(other.indptr, dtype=idx_dtype),
           np.asarray(other.indices, dtype=idx_dtype),
           other.data,
           indptr, indices, data)

        A = self.__class__((data, indices, indptr), shape=self.shape)
        A.prune()

        return A

    def _divide_sparse(self, other):
        """
        Divide this matrix by a second sparse matrix.
        """
        if other.shape != self.shape:
            raise ValueError('inconsistent shapes')

        r = self._binopt(other, '_eldiv_')

        if np.issubdtype(r.dtype, np.inexact):
            # Eldiv leaves entries outside the combined sparsity
            # pattern empty, so they must be filled manually.
            # Everything outside of other's sparsity is NaN, and everything
            # inside it is either zero or defined by eldiv.
            out = np.empty(self.shape, dtype=self.dtype)
            out.fill(np.nan)
            row, col = other.nonzero()
            out[row, col] = 0
            r = r.tocoo()
            out[r.row, r.col] = r.data
            out = matrix(out)
        else:
            # integers types go with nan <-> 0
            out = r

        return out

def _process_slice(sl, num):
    if sl is None:
        i0, i1 = 0, num
    elif isinstance(sl, slice):
        i0, i1, stride = sl.indices(num)
        if stride != 1:
            raise ValueError('slicing with step != 1 not supported')
        i0 = min(i0, i1)  # give an empty slice when i0 > i1
    elif isintlike(sl):
        if sl < 0:
            sl += num
        i0, i1 = sl, sl + 1
        if i0 < 0 or i1 > num:
            raise IndexError('index out of bounds: 0 <= %d < %d <= %d' %
                             (i0, i1, num))
    else:
        raise TypeError('expected slice or scalar')

    return i0, i1


from djitellopy import Tello
import cv2
import pygame
import numpy as np
import time
import expertFSM

# Speed of the drone

S = 60
# Frames per second of the pygame window display
# A low number also results in input lag, as input information is processed once per frame.

FPS = 120
SCREEN_WIDTH = 960
SCREEN_HEIGHT = 720
face_cascade = cv2.CascadeClassifier('./src/cascades/haarcascade_frontalface_default.xml')
# Calculate center of frame
center_x = int(SCREEN_WIDTH/2)
center_y = int(SCREEN_HEIGHT/2)

#function to take in list of faces and remove all but the one with the largest area
def targetFace(faces):
    faces=list(faces)
    if len(faces) != 0:
        max=-1
        maxIndex=0
        for i,face in enumerate(faces):
            (x,y,w,h) = face
            if w*h > max:
                max=w*h
                maxIndex=i
        return [np.array(faces[maxIndex])]
    return np.array([])

class FrontEnd(object):

    def __init__(self):
        # Init pygame

        pygame.init()

        # Creat pygame window
        pygame.display.set_caption("Automatic Drone Tracking")
        self.screen = pygame.display.set_mode([960, 720])

        # Init Tello object that interacts with the Tello drone

        self.drone = Tello()

        # Drone velocities between -100~100

        self.for_back_velocity = 0
        self.left_right_velocity = 0
        self.up_down_velocity = 0
        self.yaw_velocity = 0
        self.speed = 10

        self.send_rc_control = False

        # create update timer
        pygame.time.set_timer(pygame.USEREVENT + 1, 1000 // FPS)

    def run(self):

        self.drone.connect()
        self.drone.set_speed(self.speed)

        #Streaming refresh if was already on
        self.drone.streamoff()
        self.drone.streamon()

        frame_read = self.drone.get_frame_read()

        should_stop = False
        flag = False
        while not should_stop:

            for event in pygame.event.get():
                if event.type == pygame.USEREVENT + 1:
                    self.update()
                    if(flag is True):
                        pygame.time.wait(250)

                    self.left_right_velocity, self.for_back_velocity, self.up_down_velocity, self.yaw_velocity = (0,0,0,0)
                    self.update()
                    flag = False
                elif event.type == pygame.QUIT:
                    should_stop = True
                elif event.type == pygame.KEYDOWN:
                    if event.key == pygame.K_ESCAPE:
                        should_stop = True
                    # else:
                    #     self.keydown(event.key)
                elif event.type == pygame.KEYUP:
                    self.keyup(event.key)

            if frame_read.stopped:
                break

            #Face Detection Begins
            self.screen.fill([0, 0, 0])
            #----------------------------
            frame = frame_read.frame
            frame = cv2.resize(frame, (SCREEN_WIDTH, SCREEN_HEIGHT))

            # Draw circle at center of the frame
            cv2.circle(frame, (center_x, center_y), 10, (0, 255, 0))

            # Convert frame to grayscale in order to apply the haar cascade for face identification
            gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            faces = face_cascade.detectMultiScale(gray, 1.3, minNeighbors=5)

            # If a face is recognized, add to list of faces and draw indicators to frame around face
            face_center_x = center_x
            face_center_y = center_y
            z_area = 0

            faces=targetFace(faces)

            for face in faces:
                (x, y, w, h) = face
                cv2.rectangle(frame,(x, y),(x + w, y + h),(255, 255, 0), 2)

                face_center_x = x + int(h/2)
                face_center_y = y + int(w/2)
                z_area = w * h
                cv2.circle(frame, (face_center_x, face_center_y), 10, (0, 0, 255))

            # Calculate recognized face offset from center
            offset_x = face_center_x - center_x
            offset_y = face_center_y - center_y

            '''
            CALL FSM, adjustment made at top of loop
            = (10,10,0,0)
            FSM_TICK(update_X, update_Y, update_ROI):
            '''
            #pygame.time.wait(1000)
            if(flag is False):
                velocities = expertFSM.FSM_TICK(offset_x, offset_y, z_area)
                if(velocities != (0,0,0,0)):
                    self.left_right_velocity, self.for_back_velocity, self.up_down_velocity, self.yaw_velocity = velocities
                    flag = True

            #print(response)

            cv2.putText(frame, f'[{offset_x}, {offset_y}, {z_area}]', (10, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (255,255,255), 2, cv2.LINE_AA)

            text = "Battery: {}%".format(self.drone.get_battery())
            cv2.putText(frame, text, (5, 720 - 5),cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            frame = np.rot90(frame)
            frame = np.flipud(frame)

            frame = pygame.surfarray.make_surface(frame)
            self.screen.blit(frame, (0, 0))
            pygame.display.update()

            time.sleep(1 / FPS)

        # Call it always before finishing. To deallocate resources.
        self.drone.end()

    def keyup(self, key):
        if key == pygame.K_t:  # takeoff
            self.drone.takeoff()
            self.send_rc_control = True
        elif key == pygame.K_l:  # land
            not self.drone.land()
            self.send_rc_control = False

    def update(self):
        if self.send_rc_control:
            self.drone.send_rc_control(self.left_right_velocity, self.for_back_velocity,self.up_down_velocity, self.yaw_velocity)
            #print("Sending Velocities", self.left_right_velocity, self.for_back_velocity,self.up_down_velocity, self.yaw_velocity)

def main():
    frontend = FrontEnd()
    frontend.run()

if __name__ == '__main__':
    main()


import copy
import pickle
import numpy as np
import ray
import torch

from sklearn.neighbors import KDTree

from odium.utils.simple_utils import convert_to_list_of_np_arrays, multi_append
from odium.utils.agent_utils.save_load_agent import save_agent, load_agent
from odium.utils.simulation_utils import set_sim_state_and_goal, apply_dynamics_residual

import odium.utils.logger as logger

from odium.agents.fetch_agents.sampler import rts_sampler
from odium.agents.fetch_agents.memory import rts_memory
from odium.agents.fetch_agents.approximators import StateValueResidual, get_state_value_residual, DynamicsResidual, get_next_observation, KNNDynamicsResidual, GPDynamicsResidual
from odium.agents.fetch_agents.normalizer import FeatureNormalizer
from odium.agents.fetch_agents.worker import InternalRolloutWorker, set_workers_num_expansions
from odium.agents.fetch_agents.discrepancy_utils import get_discrepancy_neighbors, apply_discrepancy_penalty

class push_rts_agent:
    def __init__(self, args, env_params, env, planning_env, controller):
        '''
        args - arguments
        env_params - environment parameters
        env - real world
        planning_env - internal model
        controller - planner
        '''
        # Store all given arguments
        self.args, self.env_params, self.env = args, env_params, env
        self.planning_env, self.controller = planning_env, controller

        # Sampler
        # Initialize sampler to sample from buffers
        self.sampler = rts_sampler(args,
                                   env.compute_reward,
                                   controller.heuristic_obs_g,
                                   env.extract_features)

        # Memory
        # Initialize memory/buffers to store transition
        self.memory = rts_memory(args,
                                 env_params,
                                 self.sampler)

        # Approximators
        # Initialize all relevant approximators
        # STATE VALUE RESIDUAL
        self.state_value_residual = StateValueResidual(env_params)
        self.state_value_target_residual = StateValueResidual(env_params)
        self.state_value_target_residual.load_state_dict(
            self.state_value_residual.state_dict())

        # Dynamics model
        if self.args.agent == 'mbpo':
            self.residual_dynamics = DynamicsResidual(env_params)
        elif self.args.agent == 'mbpo_knn':
            self.residual_dynamics = [KNNDynamicsResidual(
                args, env_params) for _ in range(self.env_params['num_actions'])]
        else:
            self.residual_dynamics = [GPDynamicsResidual(
                args, env_params) for _ in range(self.env_params['num_actions'])]

        # KD Tree models
        self.kdtrees = [None for _ in range(self.env_params['num_actions'])]

        # Optimizers
        # Initialize all optimizers
        # STATE VALUE RESIDUAL
        self.state_value_residual_optim = torch.optim.Adam(
            self.state_value_residual.parameters(),
            lr=args.lr_value_residual,
            weight_decay=args.l2_reg_value_residual)
        # LEARNED DYNAMICS
        if self.args.agent == 'mbpo':
            self.residual_dynamics_optim = torch.optim.Adam(
                self.residual_dynamics.parameters(),
                lr=args.lr_dynamics,
                weight_decay=args.l2_reg_dynamics)

        # Normalizers
        # Initialize all normalizers
        # FEATURES
        self.features_normalizer = FeatureNormalizer(env_params)

        # Workers
        # Initialize all workers
        #self.internal_rollout_workers = [InternalRolloutWorker.remote(args,
        #                                                             env_params,
        #                                                              worker_id=i)
        #                                 for i in range(args.n_rts_workers)]

        # Trackers
        self.n_planning_steps = 0
        return

    '''
    ONLINE
    '''

    def save(self, epoch, success_rate):
        return save_agent(path=self.args.save_dir+'/fetch_rts_agent.pth',
                          network_state_dict=self.state_value_residual.state_dict(),
                          optimizer_state_dict=self.state_value_residual_optim.state_dict(),
                          normalizer_state_dict=self.features_normalizer.state_dict(),
                          epoch=epoch,
                          success_rate=success_rate)

    def load(self):
        load_dict, load_dict_keys = load_agent(
            self.args.load_dir+'/fetch_rts_agent.pth')
        self.state_value_residual.load_state_dict(
            load_dict['network_state_dict'])
        self.state_value_target_residual.load_state_dict(
            load_dict['network_state_dict'])
        if 'optimizer_state_dict' in load_dict_keys:
            self.state_value_residual_optim.load_state_dict(
                load_dict['optimizer_state_dict'])
        if 'normalizer_state_dict' in load_dict_keys:
            self.features_normalizer.load_state_dict(
                load_dict['normalizer_state_dict'])
        if 'knn_dynamics_residuals_serialized' in load_dict_keys:
            self.knn_dynamics_residuals = pickle.loads(
                load_dict['knn_dynamics_residuals_serialized'])
        return

    def _update_target_network(self, target, source):
        # Simply copy parameter values
        for target_param, param in zip(target.parameters(), source.parameters()):
            target_param.data.copy_(
                (1 - self.args.polyak) * param.data + self.args.polyak * target_param.data)

    def _update_batch_residual_dynamics(self):
        # Sample transitions
        if self.args.offline:
            # If offline
            transitions = self.memory.sample_internal_world_memory(
                self.args.batch_size)
            obs, ac_ind, obs_next = transitions['obs'], transitions['actions'], transitions['obs_next']
            raise Exception('Online mode is required')
        else:
            # If online
            transitions = self.memory.sample_real_world_memory(
                self.args.batch_size)
            obs, ac_ind, obs_next = transitions['obs'], transitions['actions'], transitions['obs_next']
            obs_sim_next = transitions['obs_sim_next']

        gripper_pos = obs[:, :2]
        obj_pos = obs[:, 3:5]
        s = np.concatenate([gripper_pos, obj_pos], axis=1)
        s_tensor = torch.as_tensor(s, dtype=torch.float32)
        a_tensor = torch.as_tensor(ac_ind, dtype=torch.long)

        # Get predicted next state
        predicted_s_next_tensor = self.residual_dynamics(
            s_tensor, a_tensor)

        # Get true next state
        gripper_pos_next = obs_next[:, :2]
        obj_pos_next = obs_next[:, 3:5]
        s_next = np.concatenate([gripper_pos_next, obj_pos_next], axis=1)
        s_next_tensor = torch.as_tensor(s_next, dtype=torch.float32)

        # Get sim next state
        gripper_pos_sim_next = obs_sim_next[:, :2]
        obj_pos_sim_next = obs_sim_next[:, 3:5]
        s_sim_next = np.concatenate(
            [gripper_pos_sim_next, obj_pos_sim_next], axis=1)
        s_sim_next_tensor = torch.as_tensor(s_sim_next, dtype=torch.float32)

        # Compute target
        target_tensor = s_next_tensor - s_sim_next_tensor

        # Compute MSE loss
        loss = (predicted_s_next_tensor - target_tensor).pow(2).mean()
        # Backprop and step
        self.residual_dynamics_optim.zero_grad()
        loss.backward()
        self.residual_dynamics_optim.step()

        assert self.args.agent == 'mbpo'
        self.controller.reconfigure_residual_dynamics(
            self.get_residual_dynamics)

        return loss

    def _update_residual_dynamics(self):
        # Sample all real world transitions
        transitions = self.memory.sample_real_world_memory()
        obs, ac_ind, obs_next = transitions['obs'], transitions['actions'], transitions['obs_next']
        obs_sim_next = transitions['obs_sim_next']

        gripper_pos = obs[:, :2]
        obj_pos = obs[:, 3:5]
        s = np.concatenate([gripper_pos, obj_pos], axis=1)

        # Get true next state
        gripper_pos_next = obs_next[:, :2]
        obj_pos_next = obs_next[:, 3:5]
        s_next = np.concatenate([gripper_pos_next, obj_pos_next], axis=1)

        # Get sim next state
        gripper_pos_sim_next = obs_sim_next[:, :2]
        obj_pos_sim_next = obs_sim_next[:, 3:5]
        s_sim_next = np.concatenate(
            [gripper_pos_sim_next, obj_pos_sim_next], axis=1)

        # Compute target
        target = s_next - s_sim_next

        loss = 0
        for i in range(self.env_params['num_actions']):
            ac_mask = ac_ind == i
            s_mask = s[ac_mask]
            target_mask = target[ac_mask]

            if s_mask.shape[0] == 0:
                # No data points for this action
                continue
            if self.args.agent == 'mbpo_knn' or self.args.agent == 'mbpo_gp':
                # Fit the KNN/GP model
                loss += self.residual_dynamics[i].fit(s_mask, target_mask)
            else:
                raise Exception(
                    'This agent does not support residual dynamics')

        assert self.args.agent == 'mbpo_knn' or self.args.agent == 'mbpo_gp'
        self.controller.reconfigure_residual_dynamics(
            self.get_residual_dynamics)

        return loss

    def _update_state_value_residual(self):
        # Sample transitions
        transitions = self.memory.sample_internal_world_memory(
            self.args.batch_size)
        qpos, qvel = transitions['qpos'], transitions['qvel']
        obs, g, ag = transitions['obs'], transitions['g'], transitions['ag']
        # features, heuristic = transitions['features'], transitions['heuristic']

        # Compute target by restarting search from the sampled states
        num_workers = self.args.n_rts_workers
        if self.args.batch_size < self.args.n_rts_workers:
            num_workers = self.args.batch_size
        num_per_worker = self.args.batch_size // num_workers
        # Put residual in object store
        value_target_residual_state_dict_id = ray.put(
            self.state_value_target_residual.state_dict())
        # Put normalizer in object store
        feature_norm_dict_id = ray.put(self.features_normalizer.state_dict())
        # Put knn dynamics residuals in object store
        if self.args.agent == 'rts':
            kdtrees_serialized_id = ray.put(pickle.dumps(
                self.kdtrees))
        else:
            kdtrees_serialized_id = None
        # Put residual dynamics in object store
        if self.args.agent == 'mbpo':
            residual_dynamics_state_dict_id = ray.put(
                self.residual_dynamics.state_dict())
        elif self.args.agent == 'mbpo_knn' or self.args.agent == 'mbpo_gp':
            residual_dynamics_state_dict_id = ray.put(
                pickle.dumps(self.residual_dynamics))
        else:
            residual_dynamics_state_dict_id = None
        results, count = [], 0
        # Set all workers num expansions
        set_workers_num_expansions(self.internal_rollout_workers,
                                   self.args.n_offline_expansions)
        for worker_id in range(num_workers):
            if worker_id == num_workers - 1:
                # last worker takes the remaining load
                num_per_worker = self.args.batch_size - count
            # Set parameters
            ray.get(self.internal_rollout_workers[worker_id].set_worker_params.remote(
                value_residual_state_dict=value_target_residual_state_dict_id,
                feature_norm_dict=feature_norm_dict_id,
                kdtrees_serialized=kdtrees_serialized_id,
                residual_dynamics_state_dict=residual_dynamics_state_dict_id))
            # Send Job
            results.append(
                self.internal_rollout_workers[worker_id].lookahead_batch.remote(
                    obs[count:count+num_per_worker],
                    g[count:count+num_per_worker],
                    ag[count:count+num_per_worker],
                    qpos[count:count+num_per_worker],
                    qvel[count:count+num_per_worker]))
            count += num_per_worker
        # Check if all transitions have targets
        assert count == self.args.batch_size
        # Get all targets
        results = ray.get(results)
        target_infos = [item for sublist in results for item in sublist]

        # Extract the states, their features and their corresponding targets
        obs_closed = [k.obs['observation'].copy()
                      for info in target_infos for k in info['closed']]
        goals_closed = [k.obs['desired_goal'].copy()
                        for info in target_infos for k in info['closed']]
        heuristic_closed = [self.controller.heuristic_obs_g(
            obs_closed[i], goals_closed[i]) for i in range(len(obs_closed))]
        features_closed = [self.env.extract_features(
            obs_closed[i], goals_closed[i]) for i in range(len(obs_closed))]
        targets_closed = [info['best_node_f'] -
                          k._g for info in target_infos for k in info['closed']]

        targets_closed = np.array(targets_closed).reshape(-1, 1)
        targets_tensor = torch.as_tensor(targets_closed, dtype=torch.float32)
        # Set all workers num expansions
        set_workers_num_expansions(self.internal_rollout_workers,
                                   self.args.n_expansions)
        # Normalize features
        inputs_norm = torch.as_tensor(self.features_normalizer.normalize(features_closed),
                                      dtype=torch.float32)
        heuristic_tensor = torch.as_tensor(
            heuristic_closed, dtype=torch.float32).view(-1, 1)

        # Compute target residuals
        target_residual_tensor = targets_tensor - heuristic_tensor
        # Clip target residual tenssor to avoid value function less than zero
        target_residual_tensor = torch.max(
            target_residual_tensor, -heuristic_tensor)
        # Clip target residual tensor to avoid value function greater than horizon
        if self.args.offline:
            target_residual_tensor = torch.min(target_residual_tensor,
                                               self.env_params['offline_max_timesteps'] - heuristic_tensor)

        # COmpute predictions
        residual_tensor = self.state_value_residual(inputs_norm)
        # COmpute loss
        state_value_residual_loss = (
            residual_tensor - target_residual_tensor).pow(2).mean()

        # Backprop and step
        self.state_value_residual_optim.zero_grad()
        state_value_residual_loss.backward()
        self.state_value_residual_optim.step()

        # Configure heuristic for controller
        self.controller.reconfigure_heuristic(
            lambda obs: get_state_value_residual(obs,
                                                 self.preproc_inputs,
                                                 self.state_value_residual))

        return state_value_residual_loss

    def _update_discrepancy_model(self):
        # For now updating the KDTrees in batches, which is not really efficient
        # Future TODO: is to make it incremental and efficient
        # Get all transitions with discrepancy in dynamics
        transitions = self.memory.sample_real_world_memory()
        # Extract relevant quantities
        obs, ac_ind = transitions['obs'], transitions['actions']

        # Construct 4D points
        # obs[0:2] is gripper 2D position
        # obs[3:5] is object 2D position
        real_pos = np.concatenate([obs[:, 0:2], obs[:, 3:5]], axis=1)

        # Add it to the respective KDTrees
        for i in range(self.env_params['num_actions']):
            # Get points corresponding to this action
            ac_mask = ac_ind == i
            points = real_pos[ac_mask]

            if points.shape[0] == 0:
                # No data points for this action
                continue

            # Fit the KDTree
            self.kdtrees[i] = KDTree(points)

        # Configure discrepancy model for controller
        assert self.args.agent == 'rts'
        self.controller.reconfigure_discrepancy(
            lambda obs, ac: get_discrepancy_neighbors(obs,
                                                      ac,
                                                      self.construct_4d_point,
                                                      self.kdtrees,
                                                      self.args.neighbor_radius)
        )

        return

    def _check_dynamics_transition(self, transition):
        obs, _, _, _, _, obs_next, obs_sim_next = transition
        real_next_pos = np.concatenate([obs_next[0:2], obs_next[3:5]])
        sim_next_pos = np.concatenate([obs_sim_next[0:2], obs_sim_next[3:5]])
        residual = real_next_pos - sim_next_pos
        if np.linalg.norm(residual) < self.args.dynamic_residual_threshold:
            return False
        return True

    def preproc_inputs(self, obs, g):
        '''
        Function to preprocess inputs
        '''
        features = self.env.extract_features(obs, g)
        features_norm = self.features_normalizer.normalize(features)
        inputs = torch.as_tensor(
            features_norm, dtype=torch.float32).view(1, -1)
        return inputs

    def construct_4d_point(self, obs, ac):
        # Concatenate 2D gripper pos and 2D object pos
        pos = np.concatenate([obs[0:2], obs[3:5]]).reshape(1, -1)
        ac_ind = self.env.discrete_actions[tuple(ac)]
        return pos, ac_ind

    def online_rollout(self, initial_observation):
        n_steps = 0
        mb_obs, mb_ag, mb_g, mb_actions, mb_heuristic = [], [], [], [], []
        mb_reward, mb_qpos, mb_qvel, mb_features = [], [], [], []
        mb_penetration = []
        # Set initial state
        observation = copy.deepcopy(initial_observation)
        # Data structures
        r_obs, r_ag, r_g, r_actions, r_heuristic = [], [], [], [], []
        r_reward, r_qpos, r_qvel, r_features = [], [], [], []
        r_penetration = []
        # Start
        obs = observation['observation']
        ag = observation['achieved_goal']
        g = observation['desired_goal']
        qpos = observation['qpos']
        qvel = observation['qvel']
        #set_sim_state_and_goal(
        #    self.env,
        #    qpos.copy(),
        #    qvel.copy(),
        #    g.copy())
        features = self.env.extract_features(obs, g)
        heuristic = self.controller.heuristic_obs_g(obs, g)
        for _ in range(self.env_params['max_timesteps']):
            ac, _ = self.controller.act(observation)
            ac_ind = self.env.discrete_actions[tuple(ac)]
            next_observation, rew, _, info = self.env.step(ac)
            penetration = info['penetration']
            if self.args.agent == 'rts':
                rew = apply_discrepancy_penalty(
                    observation, ac, rew, self.controller.discrepancy_fn)
            if self.args.agent == 'mbpo' or self.args.agent == 'mbpo_knn':
                next_observation, rew = apply_dynamics_residual(self.env,
                                                                self.get_residual_dynamics,
                                                                observation,
                                                                info,
                                                                ac,
                                                                next_observation)
            n_steps += 1
            # Add to data structures
            multi_append([r_obs, r_ag, r_g, r_actions, r_heuristic, r_reward, r_qpos, r_qvel, r_features, r_penetration],
                         [obs.copy(), ag.copy(), g.copy(), ac_ind, heuristic,
                          rew, qpos.copy(), qvel.copy(), features.copy(), penetration])
            # Move to next step
            observation = copy.deepcopy(next_observation)
            obs = observation['observation']
            ag = observation['achieved_goal']
            g = observation['desired_goal']
            qpos = observation['sim_state'].qpos
            qvel = observation['sim_state'].qvel
            features = self.env.extract_features(obs, g)
            heuristic = self.controller.heuristic_obs_g(obs, g)
        multi_append([r_obs, r_ag, r_heuristic, r_features],
                     [obs.copy(), ag.copy(), heuristic, features.copy()])
        multi_append([mb_obs, mb_ag, mb_g, mb_actions, mb_heuristic,
                      mb_reward, mb_qpos, mb_qvel, mb_features, mb_penetration],
                     [r_obs, r_ag, r_g, r_actions, r_heuristic,
                      r_reward, r_qpos, r_qvel, r_features, r_penetration])

        mb_obs, mb_ag, mb_g, mb_actions, mb_heuristic, mb_reward, mb_qpos, mb_qvel, mb_features, mb_penetration = convert_to_list_of_np_arrays(
            [mb_obs, mb_ag, mb_g, mb_actions, mb_heuristic,
             mb_reward, mb_qpos, mb_qvel, mb_features, mb_penetration]
        )

        # Store in memory
        self.memory.store_internal_model_rollout([mb_obs, mb_ag, mb_g,
                                                  mb_actions, mb_heuristic, mb_reward,
                                                  mb_qpos, mb_qvel, mb_features, mb_penetration])
        # Update normalizer
        self.features_normalizer.update_normalizer([mb_obs, mb_ag, mb_g,
                                                    mb_actions, mb_heuristic, mb_reward,
                                                    mb_qpos, mb_qvel, mb_features, mb_penetration], self.sampler)

        return n_steps

    def plan_online_in_model(self, n_planning_updates, initial_observation):
        # Clear memory
        # self.memory.clear(internal=True, real=False)
        n_updates = 0
        losses = []
        while n_updates < n_planning_updates:
            n_updates += 1
            self.online_rollout(initial_observation)
            state_value_residual_loss = self._update_state_value_residual()
            losses.append(state_value_residual_loss.item())
            self._update_target_network(self.state_value_target_residual,
                                        self.state_value_residual)

        return np.mean(losses)

    def get_residual_dynamics(self, obs, ac):
        if self.args.agent == 'mbpo':
            return get_next_observation(obs, ac,
                                        self.preproc_dynamics_inputs, self.residual_dynamics)
        elif self.args.agent == 'mbpo_knn' or self.args.agent == 'mbpo_gp':
            gripper_pos = obs['observation'][:2]
            obj_pos = obs['observation'][3:5]
            s = np.concatenate([gripper_pos, obj_pos])
            ac_ind = self.env.discrete_actions[tuple(ac)]
            residual = self.residual_dynamics[ac_ind].predict(
                s.reshape(1, -1)).squeeze()
            return residual
        else:
            raise NotImplementedError

    def preproc_dynamics_inputs(self, obs, ac):
        gripper_pos = obs[:2]
        obj_pos = obs[3:5]
        s = np.concatenate([gripper_pos, obj_pos])
        ac_ind = self.env.discrete_actions[tuple(ac)]

        s_tensor = torch.as_tensor(s, dtype=torch.float32).view(1, -1)
        a_tensor = torch.as_tensor(ac_ind, dtype=torch.long).view(1, -1)

        return s_tensor, a_tensor

    def learn_online_in_real_world(self, max_timesteps=None):
        # If any pre-existing model is given, load it
        if self.args.load_dir:
            self.load()

        # Reset the environment
        observation = self.env.reset()
        # Configure heuristic for controller
        self.controller.reconfigure_heuristic(
            lambda obs: get_state_value_residual(obs,
                                                 self.preproc_inputs,
                                                 self.state_value_residual))
        # Configure dynamics for controller
        if self.args.agent == 'rts':
            self.controller.reconfigure_discrepancy(
                lambda obs, ac: get_discrepancy_neighbors(obs,
                                                          ac,
                                                          self.construct_4d_point,
                                                          self.kdtrees,
                                                          self.args.neighbor_radius))

        # Configure dynamics for controller
        if self.args.agent == 'mbpo' or self.args.agent == 'mbpo_knn' or self.args.agent == 'mbpo_gp':
            self.controller.reconfigure_residual_dynamics(
                self.get_residual_dynamics)
        # Count of total number of steps
        total_n_steps = 0
        while True:
            obs = observation['observation']
            g = observation['desired_goal']
            qpos = observation['qpos']
            qvel = observation['qvel']
            # Get action from the controller
            simple_state = observation['observation'] #planner needs to pick this somehow
            ac, info = self.controller.act(simple_state)
            if self.args.agent == 'rts':
                assert self.controller.residual_dynamics_fn is None
            if self.args.agent == 'mbpo' or self.args.agent == 'mbpo_knn' or self.args.agent == 'mbpo_gp':
                assert self.controller.discrepancy_fn is None
            # Get discrete action index
            if ac is None:
                print("Debug")
            ac_ind = self.env.discrete_actions[tuple(ac)]
            # Get the next observation
            next_observation, rew, _, _ = self.env.step(ac)
            if self.args.render:
                self.env.render()
            total_n_steps += 1
            # Check if we reached the goal
            if self.env.env.is_success(next_observation['achieved_goal'], g):
                print('REACHED GOAL!')
                break
            # Get the next obs
            obs_next = next_observation['observation']
            # Get the sim next obs
            #set_sim_state_and_goal(self.env,
            #                       qpos.copy(),
            #                       qvel.copy(),
            #                       g.copy())
            next_observation_sim, _, _, _ = self.env.step(ac)
            obs_sim_next = next_observation_sim['observation']
            # Store transition
            transition = [obs.copy(),
                          g.copy(),
                          ac_ind,
                          qpos.copy(),
                          qvel.copy(),
                          obs_next.copy(),
                          obs_sim_next.copy()]
            dynamics_losses = []
            # RTS
            if self.args.agent == 'rts' and self._check_dynamics_transition(transition):
                # print('DISCREPANCY IN DYNAMICS')
                self.memory.store_real_world_transition(transition)
                # # Fit model
                self._update_discrepancy_model()
            # MBPO
            elif self.args.agent == 'mbpo' or self.args.agent == 'mbpo_knn' or self.args.agent == 'mbpo_gp':
                self.memory.store_real_world_transition(transition)
                # Update the dynamics
                if self.args.agent == 'mbpo':
                    for _ in range(self.args.n_online_planning_updates):
                        # Update dynamics
                        loss = self._update_batch_residual_dynamics()
                        dynamics_losses.append(loss.item())
                else:
                    loss = self._update_residual_dynamics()
                    dynamics_losses.append(loss)
            # # Plan in the model
            #value_loss = self.plan_online_in_model(n_planning_updates=self.args.n_online_planning_updates,
            #                                       initial_observation=copy.deepcopy(observation))
            #AL TODO replace this with something else
            # Log
            logger.record_tabular('n_steps', total_n_steps)
            if self.args.agent == 'mbpo' or self.args.agent == 'mbpo_knn' or self.args.agent == 'mbpo_gp':
                logger.record_tabular(
                    'dynamics loss', np.mean(dynamics_losses))
            #logger.record_tabular('residual_loss', value_loss)
            # logger.dump_tabular()
            # Move to next iteration
            observation = copy.deepcopy(next_observation)

            if max_timesteps and total_n_steps >= max_timesteps:
                break

        return total_n_steps


import numpy as np
import cv2
from matplotlib import pyplot as plt
from skimage import io
from datetime import datetime
import os

def yolo_detect(img, height, width, channels):
    # yolo
    net = cv2.dnn.readNet(os.path.join(
        os.getcwd(), "Lib/Detection/yolov3.weights"), os.path.join(os.getcwd(), "Lib/Detection/yolov3.cfg"))
    classes = []

    with open(os.path.join(
            os.getcwd(), "Lib/Detection/coco.names"), "r") as f:
        classes = [line.strip() for line in f.readlines()]
    layer_names = net.getLayerNames()
    output_layers = [layer_names[i[0] - 1]
                     for i in net.getUnconnectedOutLayers()]
    colors = (0, 255, 0)

    blob = cv2.dnn.blobFromImage(
        img, 0.00392, (416, 416), (0, 0, 0), True, crop=False)

    net.setInput(blob)
    outs = net.forward(output_layers)

    # Showing informations on the screen
    class_ids = []
    confidences = []
    boxes = []
    max_area = 0

    for out in outs:
        for detection in out:
            scores = detection[5:]
            class_id = np.argmax(scores)
            confidence = scores[class_id]
            if confidence > 0.75:

                # Object detected
                center_x = int(detection[0] * width)
                center_y = int(detection[1] * height)
                w = int(detection[2] * width)
                h = int(detection[3] * height)

                # Rectangle coordinates
                x = int(center_x - w / 2)
                y = int(center_y - h / 2)

                boxes.append([x, y, w, h])
                if max_area < h*w:
                    max_area = h*w
                    # print(max_area)
                confidences.append(float(confidence))
                class_ids.append(class_id)

    # Getting the biggest area Objects Alone and Discarding the others
    max_area_boxes = []
    new_confidence = []
    for i in range(len(boxes)):
        if boxes[i][2]*boxes[i][3] == max_area:
            max_area_boxes.append(boxes[i])
            new_confidence.append(confidences[i])

    # Marking the bounding Boxes for the largest Area Objects
    indexes = cv2.dnn.NMSBoxes(max_area_boxes, new_confidence, 0.5, 0.4)

    return max_area_boxes, indexes, classes, class_ids

def rgb_to_hsv(r, g, b):
    r, g, b = r/255.0, g/255.0, b/255.0
    mx = max(r, g, b)
    mn = min(r, g, b)
    df = mx-mn
    if mx == mn:
        h = 0
    elif mx == r:
        h = (60 * ((g-b)/df) + 360) % 360
    elif mx == g:
        h = (60 * ((b-r)/df) + 120) % 360
    elif mx == b:
        h = (60 * ((r-g)/df) + 240) % 360
    if mx == 0:
        s = 0
    else:
        s = (df/mx)*100
    v = mx*100
    return h, s, v

def visualize_detection(img, boxes, indexes, classes, class_ids):
    cropped_image = 0
    label = 0
    # print(indexes)
    font = cv2.FONT_HERSHEY_PLAIN
    for i in range(len(boxes)):
        if i in indexes:
            x, y, w, h = boxes[i]
            label = str(classes[class_ids[i]])
            color = (0, 0, 255)
            cv2.rectangle(img, (x, y), (x + w, y + h), color, 2)
            # Uncomment the below line to show label on image
            cv2.putText(img, label, (x, y-5), font, 1, color, 2)

            # Crop out unwanted region
            cropped_image = img[y+5:y+h-5, x+5:x+w-10]

    cropped_image = cv2.resize(cropped_image, None, fx=1, fy=1)
    return cropped_image, label

def yolo_detection(img):
    # Loading image
    img = cv2.resize(img, None, fx=1, fy=1)
    height, width, channels = img.shape

    max_area_boxes, indexes, classes, class_ids = yolo_detect(
        img, height, width, channels)

    cropped_image, label = visualize_detection(img,
                                               max_area_boxes, indexes, classes, class_ids)

    cv2.imwrite("../../Dataset/Cropped_Vehicle.jpg", cropped_image)

    # Getting the dominant Color
    data = np.reshape(cropped_image, (-1, 3))
    data = np.float32(data)
    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 10, 1.0)
    flags = cv2.KMEANS_RANDOM_CENTERS
    compactness, labels, centers = cv2.kmeans(
        data, 1, None, criteria, 10, flags)

    # Converting to HSV color space from RGB.
    r = (centers[0].astype(np.int32))[2]
    g = (centers[0].astype(np.int32))[1]
    b = (centers[0].astype(np.int32))[0]
    h, s, v = rgb_to_hsv(r, g, b)

    # Printing All the Details Analysed

    now = datetime.now()
    dt_string = now.strftime("%d/%m/%Y %H:%M:%S")
    return h, v, s, dt_string, label, cropped_image


# Code derived from https://github.com/openai/improved-gan/tree/master/inception_score
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os.path
import sys
import tarfile

import numpy as np
from six.moves import urllib
import tensorflow as tf
import glob
import scipy.misc
import math
import sys
import chainer
from chainer import functions as F

MODEL_DIR = '/mnt/cephfs_new_wj/bytetrans/xuminkai/DCD/model/inception'
DATA_URL = 'http://download.tensorflow.org/models/image/imagenet/inception-2015-12-05.tgz'

# if not os.path.exists(MODEL_DIR):
#     os.makedirs(MODEL_DIR)
filename = DATA_URL.split('/')[-1]
filepath = os.path.join(MODEL_DIR, filename)
# if not os.path.exists(filepath):
#     def _progress(count, block_size, total_size):
#         sys.stdout.write('\r>> Downloading %s %.1f%%' % (
#             filename, float(count * block_size) / float(total_size) * 100.0))
#         sys.stdout.flush()
#     filepath, _ = urllib.request.urlretrieve(DATA_URL, filepath, _progress)
#     print()
#     statinfo = os.stat(filepath)
#     print('Succesfully downloaded', filename, statinfo.st_size, 'bytes.')
tarfile.open(filepath, 'r:gz').extractall(MODEL_DIR)