#!/usr/bin/env python
from copy import deepcopy
from warnings import warn
from functools import partial
from multiprocessing import cpu_count

import numpy as np
from scipy.sparse import coo_matrix
from scipy.special import digamma, gammaln, psi
try:
    from scipy.misc import logsumexp
except ImportError:
    from scipy.special import logsumexp

from sklearn.base import BaseEstimator
import joblib
from joblib import Parallel, delayed

# TODO warn if can't import, and allow computation with slow
from schpf.hpf_numba import *
from schpf.util import minibatch_ix_generator
import schpf.loss as ls
import schpf

class HPF_Gamma(object):
    """Gamma variational distributions

    Parameters
    ----------
    vi_shape: np.ndarray
        Gamma shape parameter for the variational Gamma distributions.
        Ndarray.shape[0] must match `vi_rate`
    vi_rate: np.ndarray
        Gamma rate parameter for the variational Gamma distributions.
        Ndarray.shape[0] must match `vi_shape`

    Attributes
    ----------
    vi_shape : ndarray
    vi_rate : ndarray
    dims : ndarray
        The shape of vi_shape and vi_rate
    dtype : dtype
        dtype of both vi_shape and vi_rate
    """

    @staticmethod
    def random_gamma_factory(dims, shape_prior, rate_prior, dtype=np.float64):
        """Factory method to randomly initialize variational distributions

        Parameters
        ----------
        dims: list-like
            Numpy-style shape of the matrix of Gammas.
        shape_prior: float
            Prior for variational Gammas' shapes.  Must be greater than 0.
        rate_prior: float
            Prior for variational Gammas' rates.  Must be greater than 0.

        Returns
        -------
            A randomly initialized HPF_Gamma instance
        """
        vi_shape = np.random.uniform(0.5 * shape_prior, 1.5 * shape_prior,
                                     dims).astype(dtype)
        vi_rate  = np.random.uniform(0.5 * rate_prior, 1.5 * rate_prior,
                                     dims).astype(dtype)
        return HPF_Gamma(vi_shape,vi_rate)

    def __init__(self, vi_shape, vi_rate):
        """Initializes HPF_Gamma with variational shape and rates"""
        assert(vi_shape.shape == vi_rate.shape)
        assert(vi_shape.dtype == vi_rate.dtype)
        assert(np.all(vi_shape > 0))
        assert(np.all(vi_rate > 0))
        self.vi_shape = vi_shape
        self.vi_rate = vi_rate
        self.dtype = vi_shape.dtype

    def __eq__(self, other):
        if isinstance(other, self.__class__):
            shape_equal = np.array_equal(self.vi_shape, other.vi_shape)
            rate_equal = np.array_equal(self.vi_rate, other.vi_rate)
            dtype_equal = self.dtype == other.dtype
            return shape_equal and rate_equal and dtype_equal
        return False

    @property
    def dims(self):
        assert self.vi_shape.shape == self.vi_rate.shape
        return self.vi_shape.shape

    @property
    def e_x(self):
        """Expected value of the random variable(s) given variational
        distribution(s)
        """
        return self.vi_shape / self.vi_rate

    @property
    def e_logx(self):
        """Expectation of the log of random variable given variational
        distribution(s)"""
        return digamma(self.vi_shape) - np.log(self.vi_rate)

    @property
    def entropy(self):
        """Entropy of variational Gammas"""
        return  self.vi_shape - np.log(self.vi_rate) \
                + gammaln(self.vi_shape) \
                + (1 - self.vi_shape) * digamma(self.vi_shape)

    def sample(self, nsamples=1):
        """Sample from variational distributions

        Parameters
        ----------
        nsamples: int (optional, default 1)
            Number of samples to take.

        Returns
        -------
        X_rep : np.ndarray
            An ndarray of samples from the variational distributions, where
            the last dimension is the number of samples `nsamples`
        """
        samples = []
        for i in range(nsamples):
            samples.append(np.random.gamma(self.vi_shape, 1/self.vi_rate).T)
        return np.stack(samples).T

    def combine(self, other, other_ixs):
        """ Combine with another HPF_Gamma

        Useful for combining variational distributions from training data with
        variational distributions from cells that were projected onto frozen
        beta and eta

        Parameters
        ----------
        other : `HPF_Gamma`
            Other HPF_Gamma to merge with
        other_ixs : list or ndarray
            Ordered indices of other in the merged HPF_Gamma. Must have len
            equal to other.shape[0]. Must have a maximum value less than
            self.dims[0] + other.shape[0]. May not have any repeat values.

        Returns
        -------
        combined_model : `HPF_Gamma`
        """
        assert other.dims[0] == len(other_ixs)
        assert len(np.unique(other_ixs)) == len(other_ixs)
        assert self.dims[0] + other.dims[0] > np.max(other_ixs)

        new_dims = [self.dims[0]+other.dims[0], *self.dims[1:]]
        self_ixs = np.setdiff1d(np.arange(new_dims[0]),
                other_ixs)

        new_vi_shape = np.empty(new_dims, dtype=self.dtype)
        new_vi_shape[self_ixs] = self.vi_shape
        new_vi_shape[other_ixs] = other.vi_shape

        new_vi_rate = np.empty(new_dims, dtype=self.dtype)
        new_vi_rate[self_ixs] = self.vi_rate
        new_vi_rate[other_ixs] = other.vi_rate

        return HPF_Gamma(new_vi_shape, new_vi_rate)

class scHPF(BaseEstimator):
    """scHPF as described in Levitin et al., Molecular Systems Biology 2019

    Parameters
    ----------
    nfactors: int
        Number of factors (K)
    a: float, (optional, default 0.3)
        Hyperparameter a
    ap: float (optional, default 1.0)
        Hyperparameter a'
    bp: float (optional, default None)
        Hyperparameter b'. Set empirically from observed data if not
        given.
    c: float, (optional, default 0.3)
        Hyperparameter c
    cp: float (optional, default 1.0)
        Hyperparameter c'
    dp: float (optional, default None)
        Hyperparameter d'. Set empirically from observed data if not
        given.
    min_iter: int (optional, default 30):
        Minimum number of interations for training.
    max_iter: int (optional, default 1000):
        Maximum number of interations for training.
    check_freq: int (optional, default 10)
        Number of training iterations between calculating loss.
    epsilon: float (optional, default 0.001)
        Percent change of loss for convergence.
    better_than_n_ago: int (optional, default 5)
        Stop condition if loss is getting worse.  Stops training if loss
        is worse than `better_than_n_ago`*`check_freq` training steps
        ago and getting worse.
    xi: HPF_Gamma (optional, default None)
        Variational distributions for xi
    theta: HPF_Gamma (optional, default None)
        Variational distributions for theta
    eta: HPF_Gamma (optional, default None)
        Variational distributions for eta
    beta: HPF_Gamma (optional, default None)
        Variational distributions for beta
    verbose: bool (optional, default True)
            Print messages at each check_freq
    """
    def __init__(
            self,
            nfactors,
            a=0.3,
            ap=1,
            bp=None,
            c=0.3,
            cp=1,
            dp=None,
            min_iter=30,
            max_iter=1000,
            check_freq=10,
            epsilon=0.001,
            better_than_n_ago=5,
            dtype=np.float64,
            xi=None,
            theta=None,
            eta=None,
            beta=None,
            loss=[],
            verbose=True,
            ):
        """Initialize HPF instance"""
        self.version = schpf.__version__
        self.nfactors = nfactors
        self.a = a
        self.ap = ap
        self.bp = bp
        self.c = c
        self.cp = cp
        self.dp = dp
        self.min_iter = min_iter
        self.max_iter = max_iter
        self.check_freq = check_freq
        self.epsilon = epsilon
        self.better_than_n_ago = better_than_n_ago
        self.dtype = dtype
        self.verbose = verbose

        self.xi = xi
        self.eta = eta
        self.theta = theta
        self.beta = beta

        self.loss = []

    @property
    def a(self):
        try:
            return self._a
        except AttributeError:
            msg = 'Automatically using a=0.3. If you are loading a model'
            msg += ' generated with scHPF version < 0.5 and set a custom value'
            msg += ' for a, you must manually reset it and re-save the model.'
            warn(msg, RuntimeWarning)
            return 0.3

    @a.setter
    def a(self, val):
        if val == -2:
            if self.nfactors is None:
                raise ValueError('Can only set a as a function of nfactors when'
                        ' nfactors is not None')
            else:
                self._a = 1/np.sqrt(self.nfactors)
        else:
            assert val > 0
            self._a = val

    @property
    def c(self):
        try:
            return self._c
        except AttributeError:
            msg = 'Automatically using c=0.3. If you are loading a model '
            msg += ' generated with scHPF version < 0.5 and set a custom value'
            msg += ' for c, you must manually reset it and re-save the model.'
            warn(msg, RuntimeWarning)
            return 0.3

    @c.setter
    def c(self, val):
        if val == -2:
            if self.nfactors is None:
                raise ValueError('Can only set a as a function of nfactors when'
                        ' nfactors is not None')
            else:
                self._c = 1/np.sqrt(self.nfactors)
        else:
            assert val > 0
            self._c = val

    @property
    def ngenes(self):
        return self.eta.dims[0] if self.eta is not None else None

    @property
    def ncells(self):
        return self.xi.dims[0] if self.xi is not None else None

    def cell_score(self, xi=None, theta=None):
        """Get cell score from xi and theta

        Parameters
        ----------
        xi : HPF_Gamma, (optional, default self.xi)
            varitional distributions for xi
        theta : HPF_Gamma, (optional, default self.theta)
            varitional distributions for theta

        Returns
        -------
        cell_score : ndarray
            ncell x nfactor array of cell scores
        """
        xi = self.xi if xi is None else xi
        theta = self.theta if theta is None else theta
        return self._score(xi, theta)

    def gene_score(self, eta=None, beta=None):
        """Get cell score from eta and beta

        Parameters
        ----------
        eta : HPF_Gamma, (optional, default self.eta)
            varitional distributions for eta
        beta : HPF_Gamma, (optional, default self.beta)
            varitional distributions for beta

        Returns
        -------
        gene_score : ndarray
            ngene x nfactor array of cell scores
        """
        eta = self.eta if eta is None else eta
        beta = self.beta if beta is None else beta
        return self._score(eta, beta)

    def pois_llh_pointwise(self, X, theta=None, beta=None):
        """Poisson log-likelihood (for each nonzero data)

        Attempt to use numba/cffi/gsl, use numpy otherwise

        Parameters
        ----------
        X: coo_matrix
            Data to compute Poisson log likelihood of. Assumed to be nonzero.
        theta : HPF_Gamma, optional
            If given, use for theta instead of self.theta
        beta : HPF_Gamma, optional
            If given, use for beta instead of self.beta

        Returns
        -------
        llh: ndarray
        """
        theta = self.theta if theta is None else theta
        beta = self.beta if beta is None else beta
        return ls.pois_llh_pointwise(X=X, theta=theta, beta=beta)

    def cellmean_negative_pois_llh(self, X, theta=None, beta=None):
        """Convenience method for mean negative llh of nonzero entries,
           averaged by cell

        """
        theta = self.theta if theta is None else theta
        assert(theta.vi_shape.shape[0] == X.shape[0])
        beta = self.beta if beta is None else beta
        llh_pointwise= self.pois_llh_pointwise(X=X, theta=theta, beta=beta)

        llh_csr = coo_matrix((-llh_pointwise, (X.row,X.col)), shape=X.shape).tocsr()
        sums = llh_csr.sum(axis=1).A1
        counts = np.diff(llh_csr.indptr)
        averages = sums/counts

        assert(averages.shape[0] == theta.vi_shape.shape[0])
        return averages

    def mean_negative_pois_llh(self, X, theta=None, beta=None, **kwargs):
        """Convenience method for mean negative llh of nonzero entries

        """
        theta = self.theta if theta is None else theta
        beta = self.beta if beta is None else beta
        return ls.mean_negative_pois_llh(X=X, theta=theta, beta=beta)

    def fit(self, X, **kwargs):
        """Fit an scHPF model

        Parameters
        ----------
        X: coo_matrix
            Data to fit
        loss_function : function, optional (Default: None)
            loss function to use for fit. set to negative poisson likelihood
            of X if not given
        """
        (bp, dp, xi, eta, theta, beta, loss) = self._fit(
                X, **kwargs)
        self.bp = bp
        self.dp = dp
        self.xi = xi
        self.eta = eta
        self.theta = theta
        self.beta = beta
        self.loss = loss
        return self

    def project(self, X, recalc_bp=False, replace=False, min_iter=2, max_iter=50,
            check_freq=2, **kwargs):
        """Project new cells into latent space

        Gene distributions (beta and eta) are fixed.

        Parameters
        ----------
        X: coo_matrix
            Data to project.  Should have self.ngenes columns
        recalc_bp : bool, optional (Default: False)
            Recalculated value of empirical hyperparameter bp. Do not do this
            for withheld text data.
        replace: bool, optional (Default: False)
            Replace theta and xi with projected values in self. Note that
            loss will not be updated
        min_iter: int, (Default: 2)
            Replaces self.min_iter if not None. Few iterations are needed
            because beta and eta are fixed.
        max_iter: int, (Default: 10)
            Replaces self.max_iter if not None. Few iterations are needed
            because beta and eta are fixed.
        check_freq: int, optional (Default: 2)
            Number of training iterations between calculating loss.

        Returns
        -------
        result : scHPF or ndarray
            If replace=`False`, an  scHPF object with variational
            distributions theta and xi (for the new cells in `X`) and the
            same variational distributions as self for gene distributions
            beta and eta. If replace=`True`, then the loss for the projection
            (xi and theta will be updated in self but not returned). In both
            cases, bp will only be updated for the new data if self.bp==None or
            recalc_bp=`True`.

        """
        if replace and recalc_bp:
            msg = 'Cannot replace `bp` with recalculated value'
            raise ValueError(msg)

        model = self if replace else deepcopy(self)
        if recalc_bp: model.bp = None
        (bp, _, xi, _, theta, _, loss) = model._fit(X,
                min_iter=min_iter, max_iter=max_iter, check_freq=check_freq,
                freeze_genes=True, **kwargs)
        if replace:
            self.xi = xi
            self.theta = theta
            return loss
        else:
            model.bp = bp
            model.xi = xi
            model.theta = theta
            model.loss = loss
            return model

    def _score(self, capacity, loading):
        """Get the hierarchically normalized loadings which we call the cell
        or gene score in the scHPF paper

        Parameters
        ----------
        capacity : HPF_Gamma
            xi or eta
        loading : HPF_Gamma
            theta or beta

        Returns
        -------
        score : ndarray
        """
        assert(loading.dims[0] == capacity.dims[0])
        return loading.e_x * capacity.e_x[:,None]

    def _fit(self, X, freeze_genes=False, reinit=True, loss_function=None,
            min_iter=None, max_iter=None, epsilon=None, check_freq=None,
            single_process=False, checkstep_function=None, verbose=None,
            batchsize=None, beta_theta_simultaneous=False,
            loss_smoothing=1):
        """Combined internal fit/transform function

        Parameters
        ----------
        X: coo_matrix
            Data to fit
        freeze_genes: bool, (optional, default False)
            Should we update gene variational distributions eta and beta
        reinit: bool, (optional, default True)
            Randomly initialize variational distributions even if they
            already exist. Superseded by freeze_genes. Does not affect
            self.bp and self.dp which will only be set empirically if they
            are None
        loss_function : function, (optional, default None)
            Function to use for loss, which is assumed to be nonzero and
            decrease with improvement. Must accept hyperparameters a, ap,
            bp, c, cp, and dp and the variational distributions for xi, eta,
            theta, and beta even if only some of these values are used.
            Should have an internal reference to any data used (_fit will
            not pass it any data). If `loss_function` is not given or equal
            to None, the mean negative log likelihood of nonzero values in
            training data `X` is used.
        min_iter: int (optional, default None)
            Replaces self.min_iter if given.  Useful when projecting
            new data onto an existing scHPF model.
        max_iter: int (optional, default None)
            Replaces self.max_iter if given.  Useful when projecting
            new data onto an existing scHPF model.
        epsilon: float (optional, default None)
            Replaces self.epsilon if given. Percent change of loss for
            convergence.
        check_freq : int, optional (Default: None)
            Replaces self.check_freq if given.  Useful when projecting
            new data onto an existing scHPF model.
        single_process : bool, optional (Default: False)
            Use single-threaded versions of updates
        checkstep_function : function  (optional, default None)
            A function that takes arguments bp, dp, xi, eta, theta, beta,
            and t and, if given, is called at check_interval. Intended use
            is to check additional stats during training, potentially with
            hardcoded data, but is unrestricted.  Use at own risk.
        verbose: bool (optional, default None)
            If not None, overrides self.verbose
        batchsize: int, optional (Default 0)
            number of cells per batch. When 0, all cells are used
        beta_theta_simultaneous: bool, option (Default False)
            Should updates for beta and theta be computed simultaneously.
            If False, beta is updated first, and theta is updated using
            that beta
        loss_smoothing: int, optional (Default: 1)
            Smooth loss up to `loss_smoothing` check frequencies ago. 1 results
            in no smoothing. Intended to be used with batching when assessing
            convergence based on training loss, where a good value might be
            int(ncells/n_batches)

        Returns
        -------
        bp: float
            Empirically set value for bp
        dp: float
            Empirically set value for dp. Unchanged if freeze_genes.
        xi: HPF_Gamma
            Learned variational distributions for xi
        eta: HPF_Gamma
            Learned variational distributions for eta. Unchanged if
            freeze_genes.
        theta: HPF_Gamma
            Learned variational distributions for theta
        beta: HPF_Gamma
            Learned variational distributions for beta. Unchanged if
            freeze_genes.
        loss : list
            loss at each checkstep
        """
        assert loss_smoothing > 0

        # local (convenience) vars for model
        nfactors, (ncells, ngenes) = self.nfactors, X.shape
        a, ap, c, cp = self.a, self.ap, self.c, self.cp

        # get empirically set hyperparameters and variational distributions
        bp, dp, xi, eta, theta, beta = self._setup(X, freeze_genes, reinit)

        # Make first updates for hierarchical shape prior
        # (vi_shape is constant, but want to update full distribution)
        xi.vi_shape[:] = ap + nfactors * a
        if not freeze_genes:
            eta.vi_shape[:] = cp + nfactors * c

        # setup loss function as mean negative llh of nonzero training data
        # if the loss function is not given
        if loss_function is None:
            loss_function = ls.loss_function_for_data(
                    ls.mean_negative_pois_llh, X)

        # setup batch_ix iterator
        if batchsize is not None and batchsize > 1 and batchsize <= ncells:
            batched = True
            batch_ix_generator = minibatch_ix_generator(ncells, batchsize)
        else:
            batched = False
            batch_ix_generator = None

        ## init
        loss, unsmoothed_loss, pct_change = [], [], []
        # check variable overrides
        min_iter = self.min_iter if min_iter is None else min_iter
        max_iter = self.max_iter if max_iter is None else max_iter
        epsilon = self.epsilon if epsilon is None else epsilon
        check_freq = self.check_freq if check_freq is None else check_freq
        verbose = self.verbose if verbose is None else verbose
        for t in range(max_iter):
            # setup batching
            if batch_ix_generator is None:
                batch_ix = np.arange(X.shape[0])
                batchsize = ncells
                X_batch = X
            else:
                batch_ix = next(batch_ix_generator)
                X_batch = X.tocsr()[batch_ix,:].tocoo()

            if t==0 and reinit: #randomize phi for first iteration
                random_phi = np.random.dirichlet( np.ones(nfactors),
                        X_batch.data.shape[0])
                Xphi_data = X_batch.data[:,None] * random_phi
            else:
                if single_process:
                    Xphi_data = compute_Xphi_data_numpy(X_batch, theta, beta,
                            theta_ix=batch_ix)
                else:
                    Xphi_data = compute_Xphi_data(
                            X_batch.data, X_batch.row, X_batch.col,
                            theta.vi_shape[batch_ix], theta.vi_rate[batch_ix],
                            beta.vi_shape, beta.vi_rate)

            if beta_theta_simultaneous:
                # calculate gene updates but don't assign yet
                if not freeze_genes:
                    bvs = compute_loading_shape_update(Xphi_data,
                            X_batch.col, ngenes, c)
                    bvr = compute_loading_rate_update(eta.vi_shape,
                            eta.vi_rate, theta.vi_shape[batch_ix],
                            theta.vi_rate[batch_ix])
                # cell updates
                theta.vi_shape[batch_ix] = compute_loading_shape_update(
                        Xphi_data, X_batch.row, batchsize, a)
                theta.vi_rate[batch_ix] = compute_loading_rate_update(
                        xi.vi_shape[batch_ix], xi.vi_rate[batch_ix],
                        beta.vi_shape, beta.vi_rate)
                xi.vi_rate[batch_ix] = bp + theta.e_x[batch_ix].sum(1)
                # make gene updates
                if not freeze_genes:
                    beta.vi_shape = bvs
                    beta.vi_rate = bvr
                    eta.vi_rate = dp + beta.e_x.sum(1)

            else:
                if batched:
                    # cell updates, must do first for batching
                    theta.vi_shape[batch_ix] = compute_loading_shape_update(
                            Xphi_data, X_batch.row, batchsize, a)
                    theta.vi_rate[batch_ix] = compute_loading_rate_update(
                            xi.vi_shape[batch_ix], xi.vi_rate[batch_ix],
                            beta.vi_shape, beta.vi_rate)
                    xi.vi_rate[batch_ix] = bp + theta.e_x[batch_ix].sum(1)

                if not freeze_genes:
                    #gene updates
                    beta.vi_shape = compute_loading_shape_update(Xphi_data,
                            X_batch.col, ngenes, c)
                    beta.vi_rate = compute_loading_rate_update(eta.vi_shape,
                            eta.vi_rate, theta.vi_shape[batch_ix],
                            theta.vi_rate[batch_ix])
                    eta.vi_rate = dp + beta.e_x.sum(1)

                if not batched:
                    # cell updates, doing after gene updates when not batched
                    # for legacy consistency
                    theta.vi_shape[batch_ix] = compute_loading_shape_update(
                            Xphi_data, X_batch.row, batchsize, a)
                    theta.vi_rate[batch_ix] = compute_loading_rate_update(
                            xi.vi_shape[batch_ix], xi.vi_rate[batch_ix],
                            beta.vi_shape, beta.vi_rate)
                    xi.vi_rate[batch_ix] = bp + theta.e_x[batch_ix].sum(1)

            # record llh/percent change and check for convergence
            if t % check_freq == 0:

                # chech llh
                # vX = validation_data if validation_data is not None else X
                try :
                    curr = loss_function(
                                a=a, ap=ap, bp=bp, c=c, cp=cp, dp=dp,
                                xi=xi, eta=eta, theta=theta, beta=beta)
                    unsmoothed_loss.append(curr)
                    if len(unsmoothed_loss) > loss_smoothing:
                        unsmoothed_loss = unsmoothed_loss[1:]
                    # normally this is just curr as loss_smoothing=1 by default
                    loss.append(np.mean(unsmoothed_loss))
                except NameError as e:
                    print('Invalid loss function')
                    raise e

                # calculate percent change
                try:
                    curr, prev = loss[-1], loss[-2]
                    pct_change.append(100 * (curr - prev) / np.abs(prev))
                except IndexError:
                    pct_change.append(100)
                if verbose:
                    msg = '[Iter. {0: >4}]  loss:{1:.6f}  pct:{2:.9f}'.format(
                            t, curr, pct_change[-1])
                    print(msg)
                if checkstep_function is not None:
                    checkstep_function(bp=bp, dp=dp, xi=xi, eta=eta, theta=theta,
                            beta=beta, t=t)

                # check convergence
                if len(loss) > 3 and t >= min_iter:
                    # convergence conditions (all must be met)
                    current_small = np.abs(pct_change[-1]) < self.epsilon
                    prev_small = np.abs(pct_change[-2]) < self.epsilon
                    not_inflection = not (
                            (np.abs(loss[-3]) < np.abs(prev)) \
                            and (np.abs(prev) > np.abs(curr)))
                    converged = current_small and prev_small and not_inflection
                    if converged:
                        if verbose:
                            print('converged')
                        break

                    # getting worse, and has been for better_than_n_ago checks
                    # (don't waste time on a bad run)
                    if len(loss) > self.better_than_n_ago \
                            and self.better_than_n_ago:
                        nprev = loss[-self.better_than_n_ago] \
                                if len(loss)>self.better_than_n_ago else loss[0]
                        worse_than_n_ago = np.abs(nprev) < np.abs(curr)
                        getting_worse = np.abs(prev) < np.abs(curr)
                        if worse_than_n_ago and getting_worse:
                            if verbose:
                                print('getting worse break')
                            break

            # TODO message or warning or something
            if t >= self.max_iter:
                break

        return (bp, dp, xi, eta, theta, beta, loss)

    def _setup(self, X, freeze_genes=False, reinit=True, clip=True):
        """Setup variational distributions

        Parameters
        ----------
        X: coo_matrix
            Data to fit
        freeze_genes: bool, optional (Default: False)
            Should we update gene variational distributions eta and beta
        reinit: bool, optional (Default: True)
            Randomly initialize variational distributions even if they
            already exist. Superseded by freeze_genes. Does not affect
            self.bp and self.dp (which will only be set empirically if
            they are None)
        clip : bool, optional (Default: True)
            If empirically calculating dp and bp > 1000 * dp, clip dp to
            bp / 1000.

        Returns
        -------
        bp : float
        dp : float
        xi : HPF_Gamma
        eta : HPF_Gamma
        theta : HPF_Gamma
        beta : HPF_Gamma

        """
        # locals for convenience
        nfactors, (ncells, ngenes) = self.nfactors, X.shape
        a, ap, c, cp = self.a, self.ap, self.c, self.cp
        bp, dp = self.bp, self.dp

        xi, eta, theta, beta = (self.xi, self.eta, self.theta, self.beta)

        # empirically set bp and dp
        bp, dp = self._get_empirical_hypers(X, freeze_genes, clip)

        if reinit or (xi is None):
            xi = HPF_Gamma.random_gamma_factory((ncells,), ap, bp,
                    dtype=self.dtype)
        if reinit or (theta is None):
            theta = HPF_Gamma.random_gamma_factory((ncells,nfactors), a, bp,
                    dtype=self.dtype)

        # Check if variational distributions for genes exist, create if not
        # Error if freeze_genes and eta and beta don't exists
        if freeze_genes:
            if eta is None or beta is None:
                msg = 'To fit with frozen gene variational distributions ' \
                    + '(`freeze_genes`==True), eta and beta must be set to ' \
                    + 'valid HPF_Gamma instances.'
                raise ValueError(msg)
        else:
            if reinit or (eta is None):
                eta = HPF_Gamma.random_gamma_factory((ngenes,), cp, dp,
                        dtype=self.dtype)
            if reinit or (beta is None):
                beta = HPF_Gamma.random_gamma_factory((ngenes,nfactors),
                        c, dp, dtype=self.dtype)

        return (bp, dp, xi, eta, theta, beta)

    def _get_empirical_hypers(self, X, freeze_genes=False, clip=True):
        """Get empirical values for bp, dp

        Parameters
        ----------
        X : coo_matrix
            Data to fit

        Returns
        -------
        bp : float
        dp : float
        """
        bp, dp = self.bp, self.dp
        # empirically set bp and dp
        def mean_var_ratio(X, axis):
            axis_sum = X.sum(axis=axis)
            return np.mean(axis_sum) / np.var(axis_sum)
        if bp is None:
            bp = self.ap * mean_var_ratio(X, axis=1)
        if dp is None: # dp first in case of error
            if freeze_genes:
                msg = 'dp is None and cannot be set'
                msg += ' when freeze_genes is True.'
                raise ValueError(msg)
            else:
                dp = self.cp *  mean_var_ratio(X, axis=0)
                if clip and bp > 1000 * dp:
                    old_val = dp
                    dp = bp / 1000
                    print('Clipping dp: was {} now {}'.format(old_val, dp))

        return bp, dp

    def _initialize(self, X, freeze_genes=False):
        """Shortcut to setup random distributions & set variables
        """
        bp, dp, xi, eta, theta, beta = self._setup(X, freeze_genes,
                reinit=True)
        self.bp = bp
        self.dp = dp
        self.xi = xi
        self.eta = eta
        self.theta = theta
        self.beta = beta

def load_model(file_name):
    """Load a model from a joblib file

    Parameters
    ----------
    file_name : str
        Joblib file containing a saved scHPF model

    Returns
    -------
    model : scHPF
        The scHPF model in the file
    """
    return joblib.load(file_name)

def save_model(model, file_name):
    """Save model to (joblib) file

    Serialize scHPF model as a joblib file.  Joblib is simillar to pickle,
    but preferable for objects with many numpy arrays

    Parameters
    ----------
    model : scHPF
        The scHPF model object to save
    file_name : str
        Name of file to save model to
    """
    joblib.dump(model, file_name)

def combine_across_cells(x, y, y_ixs):
    """Combine theta & xi from two scHPF instance with the same beta & eta

    Intended to be used combining variational distributions for local
    variables (theta,xi) from training data with variational distributions
    for local variables from validation or other data that was projected
    onto the same global variational distributions (beta,eta)

    If `x.bp` != `y.bp`, returned model `xy.bp` is set to None. All other
    attributes (except for the merged xi and eta) are inherited from `x`.

    Parameters
    ----------
    x : `scHPF`
    y : `scHPF`
        The scHPF instance whose rows in the output should be at the
        corresponding indices `y_ixs`
    y_ixs : ndarray
        Row indices of `y` in the returned distributions. Must be 1-d and
        have same number of rows as `y`, have no repeats, and have no index
        greater than or equal to x.ncells + y.ncells.

    Returns
    -------
    ab : `scHPF`

    """
    assert x.dp == y.dp
    assert x.eta == y.eta
    assert x.beta == y.beta

    xy = deepcopy(x)
    if y.bp != x.bp:
        xy.bp = None
    xy.xi = x.xi.combine(y.xi, y_ixs)
    xy.theta = x.theta.combine(y.theta, y_ixs)
    return xy

def run_trials(X, nfactors,
        ntrials=5,
        min_iter=30,
        max_iter=1000,
        check_freq=10,
        epsilon=0.001,
        better_than_n_ago=5,
        dtype=np.float64,
        verbose=True,
        vcells = None,
        vX = None,
        loss_function=None,
        model_kwargs = {},
        return_all = False,
        reproject = False,
        reproject_kwargs = {},
        batchsize=0,
        beta_theta_simultaneous=False,
        loss_smoothing=1
        ):
    """
    Train with multiple random initializations, selecting model with best loss

    As scHPF uses non-convex optimization, it benefits from training with
    multiple random initializations to avoid local minima.

    Parameters
    ----------
    X: coo_matrix
        Data to fit
    nfactors: int
        Number of factors (K)
    ntrials : int,  optional (Default 5)
        Number of random initializations for training
    min_iter: int, optional (Default 30)
        Minimum number of interations for training.
    max_iter: int, optional (Default 1000):
        Maximum number of interations for training.
    check_freq: int, optional (Default 10)
        Number of training iterations between calculating loss.
    epsilon: float, optional (Default 0.001)
        Percent change of loss for convergence.
    better_than_n_ago: int, optional (Default 5)
        Stop condition if loss is getting worse.  Stops training if loss
        is worse than `better_than_n_ago`*`check_freq` training steps
        ago and getting worse.
    dtype : datatype, optional (Default np.float64)
        np.float64 or np.float32
    verbose: bool, optional (Default True)
        verbose
    vcells : coo_matrix, optional (Default None)
        cells to use in a validation loss function
    vX : coo_matrix, optional (Default None)
        nonzero entries from the cells in vX
    loss_function : function, optional (Default None)
        A loss function to asses convergence that accepts data, model
        variational parameters, and model hyperparameters.  Note this is
        distinct from the `loss_function` argument in scHPF._fit (called by
        scHPF.fit and scHPF.project), which assumes a fixed reference to data
        is included in the function and *does not* accept data as an argument.
    model_kwargs: dict, optional (Default {})
        dictionary of additional keyword arguments for model
        initialization
    return_all: bool, optional (Default False)
        return all models
    reproject: bool, optional (Default False)
        Reproject the data onto the frozen gene variables before calculating
        loss. The reprojected loss will be added to the end of loss as a
        sublist. Note that this reprojection will *not* use the `loss_function`
        argument, and instead use the default provided log likelihood
    reproject_kwargs: dict, optional (Default {'replace':True})
        Only used if `reproject` is True. Keyword args for scHPF.project.
        'replace':True cannot be changed, and will be overwritten if given
    batchsize: int, optional (Defualt 0)
        Number of cells to use per training round. All cells used if 0.
    loss_smoothing: int, optional (Default: 1)
        Smooth loss up to `loss_smoothing` check frequencies ago. 1 results in
        no smoothing. Intended to be used with batching when assessing
        convergence based on training loss, where a good value might be
        int(ncells/n_batches)

    Returns
    -------
    best_model: scHPF
        The model with the best loss facter `ntrials` random initializations
        and training runs
    rejected_models: list, optional
        Rejected models, ordered by decreasing loss . Only returned if
        return_all is True
    """
    ncells, ngenes = X.shape
    if ngenes >= 20000:
        msg = 'WARNING: you are running scHPF with {} genes,'.format(ngenes)
        msg += ' which is more than the ~20k protein coding genes in the'
        msg += ' human genome. We suggest running scHPF on protein-coding'
        msg += ' genes only.'
        print(msg)

    # get the loss function for any data
    if loss_function is None:
        loss_function = partial(ls.mean_negative_pois_llh,
                single_process=False)

    # check data we're using for loss
    if vcells is not None:
        assert X.shape[1] == vcells.shape[1]
    if vX is not None:
        assert vX.shape == X.shape
    else:
        vX = X
    # setup loss fnc w/data (will be overridden if vcells is not None)
    data_loss_function = ls.loss_function_for_data(loss_function, vX)
    # setup smoothed_loss if using batches

    # run trials
    best_loss, best_model, best_t = np.finfo(np.float64).max, None, None
    models, losses = [], [] # only used if return_all
    for t in range(ntrials):
        # make a new model
        model = scHPF(nfactors=nfactors,
                    min_iter=min_iter, max_iter=max_iter,
                    check_freq=check_freq, epsilon=epsilon,
                    better_than_n_ago=better_than_n_ago,
                    verbose=verbose, dtype=dtype,
                    **model_kwargs
                    )

        # override the loss function data if we have vcells
        # (must be redone for each new model)
        if vcells is not None:
            proj_kwargs = dict(reinit=False,
                               min_iter=1,
                               max_iter=min(10, check_freq),
                               check_freq= check_freq+1,
                               verbose=False
                               )
            data_loss_function = ls.projection_loss_function(
                    loss_function, vcells, nfactors,
                    proj_kwargs=proj_kwargs)
            def checkstep_function(**kwargs):
                loss = ls.loss_function_for_data(loss_function, X)
                print('\ttrain:', '{0:.6f}'.format(loss(**kwargs)))
        else:
            checkstep_function = None

        # fit the model
        model.fit(X, loss_function=data_loss_function,
                  checkstep_function=checkstep_function,
                  batchsize=batchsize, loss_smoothing=loss_smoothing,
                  beta_theta_simultaneous=beta_theta_simultaneous)
        if reproject:
            print('Reprojecting data...')
            reproject_kwargs['replace'] = True
            reproject_kwargs['reinit'] = False
            proj_loss = model.project(X, **reproject_kwargs)
            model.loss.append(proj_loss)
            loss = proj_loss[-1]
        else:
            loss = model.loss[-1]

        if loss < best_loss:
            best_model = model
            best_loss = loss
            best_t = t
            if verbose:
                print('New best!'.format(t))
        if return_all:
            models.append(model)
            losses.append(loss)
        if verbose:
            print('Trial {0} loss: {1:.6f}'.format(t, loss))
            print('Best loss: {0:.6f} (trial {1})'.format(best_loss, best_t))

    if return_all:
        return_order = np.argsort(losses)
        ordered_models = [models[i] for i in return_order]
        assert ordered_models[0] == best_model
        return best_model, ordered_models[1:]
    else:
        return best_model

# TODO deal with verbosity
def run_trials_pool(X, nfactors,
        ntrials=5,
        njobs=0,
        max_threads=None,
        min_iter=30,
        max_iter=1000,
        check_freq=10,
        epsilon=0.001,
        better_than_n_ago=5,
        dtype=np.float64,
        verbose=True,
        vcells = None,
        vX = None,
        loss_function=None,
        model_kwargs = {},
        return_all = False,
        reproject = False,
        reproject_kwargs = {},
        batchsize=0,
        beta_theta_simultaneous=False,
        loss_smoothing=1
        ):
    """
    Train with multiple random initializations, selecting model with best loss.
    Parallelization is done at the trial level rather than within computations

    As scHPF uses non-convex optimization, it benefits from training with
    multiple random initializations to avoid local minima.

    Parameters
    ----------
    X: coo_matrix
        Data to fit
    nfactors: int or list of ints
        Number of factors (K), may be a list for multiple k
    ntrials : int,  optional (Default 5)
        Number of random initializations for training
    njobs : int, optional (Default 0)
        Maximum number of threads in the threadpool.  0 will use all available.
    min_iter: int, optional (Default 30)
        Minimum number of interations for training.
    max_iter: int, optional (Default 1000):
        Maximum number of interations for training.
    check_freq: int, optional (Default 10)
        Number of training iterations between calculating loss.
    epsilon: float, optional (Default 0.001)
        Percent change of loss for convergence.
    better_than_n_ago: int, optional (Default 5)
        Stop condition if loss is getting worse.  Stops training if loss
        is worse than `better_than_n_ago`*`check_freq` training steps
        ago and getting worse.
    dtype : datatype, optional (Default np.float64)
        np.float64 or np.float32
    verbose: bool, optional (Default True)
        verbose
    vcells : coo_matrix, optional (Default None)
        cells to use in a validation loss function
    vX : coo_matrix, optional (Default None)
        nonzero entries from the cells in vX
    loss_function : function, optional (Default None)
        A loss function that accepts data, model variational parameters,
        and model hyperparameters.  Note this is distinct from the
        `loss_function` argument in scHPF._fit (called by scHPF.fit and
        scHPF.project), which assumes a fixed reference to data is included
        in the function and *does not* accept data as an argument.
    model_kwargs: dict, optional (Default {})
        dictionary of additional keyword arguments for model
        initialization
    return_all: bool, optional (Default False)
        return all models
    reproject: bool, optional (Default False)
        Reproject the data onto the frozen gene variables before calculating
        loss. The reprojected loss will be added to the end of loss as a
        sublist. Note that this reprojection will *not* use the `loss_function`
        argument, and instead use the default provided log likelihood
    reproject_kwargs: dict, optional (Default {'replace':True})
        Only used if `reproject` is True. Keyword args for scHPF.project.
        'replace':True cannot be changed, and will be overwritten if given
    batchsize: int, optional (Defualt 0)
            Number of cells to use per training round. All cells used if 0.
    loss_smoothing: int, optional (Default: 1)
        Smooth loss up to `loss_smoothing` check frequencies ago. 1 results in
        no smoothing. Intended to be used with batching when assessing
        convergence based on training loss, where a good value might be
        int(ncells/n_batches)

    Returns
    -------
    best_models: list(scHPF)
        The model with the best loss facter `ntrials` random initializations
        and training runs for each value in nfactors
    rejected_models: list(list(scHPF)), optional
        Rejected models, ordered by corresponding nfactors and then by
        decreasing loss . Only returned if return_all is True
    """
    ngenes = X.shape[1]
    if ngenes >= 20000:
        msg = 'WARNING: you are running scHPF with {} genes,'.format(ngenes)
        msg += ' which is more than the ~20k protein coding genes in the'
        msg += ' human genome. We suggest running scHPF on protein-coding'
        msg += ' genes only.'
        print(msg)

    # get the loss function for any data
    if loss_function is None:
        loss_function = partial(ls.mean_negative_pois_llh,
                single_process=True)

    # check data we're using for loss
    if vcells is not None:
        assert X.shape[1] == vcells.shape[1]
    if vX is not None:
        assert vX.shape == X.shape
    else:
        vX = X
    # setup loss fnc w/data (will be overridden if vcells is not None)
    data_loss_function = ls.loss_function_for_data(loss_function, vX)

    # only need to create once because will be copied to processes
    # override the loss function data if we have vcells
    # (must be redone for each new model)
    if vcells is not None:
        proj_kwargs = dict(reinit=False,
                            min_iter=1,
                            max_iter=min(10, check_freq),
                            check_freq= check_freq+1,
                            verbose=False
                            )
        data_loss_function = ls.projection_loss_function(
                loss_function, vcells, nfactors,
                proj_kwargs=proj_kwargs)

    # function to fit model
    def fit_model(nfactors):
        model = scHPF(nfactors=nfactors,
                    min_iter=min_iter, max_iter=max_iter,
                    check_freq=check_freq, epsilon=epsilon,
                    better_than_n_ago=better_than_n_ago,
                    verbose=False, dtype=dtype,
                    **model_kwargs
                    )
        # fit the model
        model.fit(X, loss_function=data_loss_function,
                  checkstep_function=None, single_process=True,
                  batchsize=batchsize, loss_smoothing=loss_smoothing)
        if reproject:
            # print('Reprojecting data...')
            reproject_kwargs['replace'] = True
            proj_loss = model.project(X, loss_function=data_loss_function,
                    **reproject_kwargs)
            model.loss.append(proj_loss)
        return model

    # get nfactors for every trial
    if isinstance(nfactors, int):
        nfactors = [nfactors]
    trial_nfactors = [t for trial_set in [[K]*ntrials for K in nfactors] \
            for t in trial_set]

    # set max processes if not given
    if njobs == 0: njobs = min(cpu_count(), len(trial_nfactors))

    # training
    with Parallel(n_jobs=njobs, verbose=10) as pool: # make the pool
        candidates = pool( delayed(fit_model)(K)  for K in trial_nfactors)

    # get the best model for every K
    ordered_best, ordered_reject = [], []
    for i,K in enumerate(nfactors):
        my_candidates = candidates[i*ntrials : (i+1)*ntrials]
        loss = [m.loss[-1][-1] if reproject else m.loss[-1] for m in
                my_candidates]
        # print(list(zip([m.nfactors for m in my_candidates],loss)))
        best_ix = np.argmin(loss)
        ordered_best.append(my_candidates[best_ix])
        ordered_reject.append([my_candidates[i] for i in np.argsort(loss)[1:]])
    if return_all:
        return ordered_best, ordered_reject
    else:
        return ordered_best


import json

class Character:
    DEFAULT_CONFIG = {
        "id": "default_id",
        "name": "Default Character",
        "race": "Human",
        "character_class": "Warrior",
        "strength": 10,
        "dexterity": 10,
        "constitution": 10,
        "intelligence": 10,
        "wisdom": 10,
        "charisma": 10,
    }

    def __init__(self, config=None):
        self.config = {**self.DEFAULT_CONFIG, **(config or {})}
        if self.config["id"] is None:
            raise ValueError("Character must have an ID")
        if self.config["name"] is None:
            raise ValueError("Character must have a Name")
        # Validate other attributes as needed
        self.state = {
          "hp": 10,
          "mp": 10,
          "xp": 0,
          "level": 1,
          "inventory": [{"id": "journal", "entries": []}, {"gold": 0,}],
          "body": {"head": [], "torso": [], "left_hand": [], "right_hand": [], "belt": [], "legs": [], "feet": []},
          "spells": [],
          "status_effects": [],
          "position": (0, 0),
          "current_map": None,
        }

    def __str__(self):
        return f"{self.config['name']} the {self.config['race']} {self.config['character_class']}"

    # Add methods to update individual attributes if needed

    def set_race(self, race):
        if race not in ['Human', 'Elf', 'Dwarf', 'Orc']:  # Add all valid races
            raise ValueError("Invalid race chosen.")
        self.config['race'] = race

    def set_character_class(self, character_class):
        if character_class not in ['Warrior', 'Mage', 'Rogue']:  # Add all valid classes
            raise ValueError("Invalid class chosen.")
        self.config['character_class'] = character_class

    def set_attribute(self, attribute, value):
        if attribute in ['strength', 'dexterity', 'constitution', 'intelligence', 'wisdom', 'charisma']:
            if not (1 <= value <= 20):  # Assuming DnD-like attribute range
                raise ValueError(f"{attribute} must be between 1 and 20.")
            self.config[attribute] = value
        else:
            raise ValueError(f"{attribute} is not a valid character attribute.")

    def save(self, file_path):
        with open(file_path, 'w') as file:
            json.dump(self.config, file)

    @classmethod
    def load(cls, file_path):
        with open(file_path, 'r') as file:
            config = json.load(file)
        return cls(config)

    def move(self, direction):
        # Check if the direction is valid in the current location
        current_location = self.state['current_map'].get_location(self.state['position'])
        if direction in current_location.connections:
            new_position_key = current_location.connections[direction]
            new_position = self.state['current_map'].get_location(new_position_key)
            self.state['position'] = new_position_key
            print(f"You move {direction} to {new_position.description}.")
        else:
            print("You can't go that way.")

    def to_dict(self):
        return {
            "config": self.config,
            "state": self.state,
            # Include any other necessary data
        }

    @classmethod
    def from_dict(cls, data):
        character = cls(data['config'])
        character.state = data['state']
        return character


import os
import torch
import torch.utils.data
import torchvision
import numpy as np
import cv2
from math import*

from pix_co_to_cartesian_co import pix_orientation
# from data.apple_dataset import AppleDataset
from data.apple_dataset_pre import AppleDataset
# from data.apple_dataset_data1 import AppleDataset_data1
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor

import utility.utils as utils
import utility.transforms as T

######################################################
# Predict with either a Faster-RCNN or Mask-RCNN predictor
######################################################
def get_transform(train):
	transforms = []
	transforms.append(T.ToTensor())
	if train:
		transforms.append(T.RandomHorizontalFlip(0.5))
	return T.Compose(transforms)

def get_maskrcnn_model_instance(num_classes):
	# load an instance segmentation model pre-trained pre-trained on COCO
	model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=False)
	# get number of input features for the classifier
	in_features = model.roi_heads.box_predictor.cls_score.in_features
	# replace the pre-trained head with a new one
	model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

	# now get the number of input features for the mask classifier
	in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
	hidden_layer = 256
	# and replace the mask predictor with a new one
	model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask, hidden_layer, num_classes)
	return model

def get_frcnn_model_instance(num_classes):
	# load an instance segmentation model pre-trained pre-trained on COCO
	model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=False)

	# get number of input features for the classifier
	in_features = model.roi_heads.box_predictor.cls_score.in_features
	# replace the pre-trained head with a new one
	model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
	return model

def drawBoundingBox(stacked,centers,roll_yaw_pitch,args,f_x,f_y,depth_array):#,imgcv,result):
#	imgcv= cv2.imread('/home/nalin/Desktop/DataSet/detection/original_data_sets/test/images/{}'.format(stacked[0][0]))
	imgcv= cv2.imread(os.path.join(args.data_path, 'images/{}'.format(stacked[0][0])))
	h,w,_ = imgcv.shape
	for box in stacked:
		print('========')
		print(box)
		print('========')
		x1 = int(box[-5])
		y1 = int(box[-4])
		x2 = int(box[-3])
		y2 = int(box[-2])
		# x.astype(int)
		conf = box[-1]
		print('x1 = ',x1)
		print('y1 = ',y1)
		print('x2 = ',x2)
		print('y2 = ',y2)
		# print(conf)
		# label = box['label']
		if float(conf) > 0.75:
			cv2.rectangle(imgcv,(x1,y1),(x2,y2),(0,255,0),1)
			centers.append([0.5*(x1+x2),0.5*(y1+y2)])
			print('=============================================')
			# img_to_array = np.asarray(imgcv)
			# print('img_to_array shape = ',img_to_array.shape)
			# print('img_to_array shape = ',img_to_array[:,:,0].shape)
			# np.savetxt('npB.csv', imgcv[:,:,0], delimiter=',')
			np.savetxt('npG.csv', imgcv[:,:,1], delimiter=',')
			# np.savetxt('npR.csv', imgcv[:,:,2], delimiter=',')
			print('Image shape = ',imgcv.shape)
			# print('=================================================')
			# for i in range [y1:y2]:
			# 	for j in range [x1:x2]:
			# 		print('pixel values of each [i][j] = ',imgcv[i,j,1])
			theta_ij_total = 0
			phi_ij_total = 0
			length_data = 0
			i = y1
			j = x1
			while(i >= y1 and i < y2):
				while(j >= x1 and j < x2):
					if( imgcv[i,j,1]<100 ):
						# print('i = ',i)
						# print('j = ',j)
						# print('pixel values of each [i][j] = ',imgcv[i,j,1])
						theta_ij,phi_ij = pix_orientation(j,i,h,w,f_x,f_y,depth_array)
						theta_ij_total = theta_ij_total+theta_ij
						phi_ij_total = phi_ij_total+phi_ij
						length_data = length_data+1
					j = j+1
				i = i+1
				j = x1

			theta = theta_ij_total/length_data
			phi = phi_ij_total/length_data
			# theta,phi = pix_orientation(int(0.5*(y1+y2)),int(0.5*(x1+x2)),h,w,f_x,f_y,depth_array)
			r11 = cos(theta)*cos(phi)
			r21 = sin(theta)*cos(phi)
			r31 = -sin(phi)
			r32 = 0
			r33 = cos(phi)
			yaw = atan2(r21,r11)
			pitch = atan2(-r31,sqrt(r32**2+r33**2))
			roll = atan2(r32,r33)
			print('roll,yaw,pitch = ',degrees(roll),degrees(yaw),degrees(pitch))
			roll_yaw_pitch.append([roll,yaw,pitch])
			cv2.putText(imgcv,conf,(x1,y1),cv2.FONT_HERSHEY_COMPLEX,0.5,(0,0,0),1)
			cv2.imwrite(os.path.join(args.data_path, 'images_with_bbox/{}'.format(stacked[0][0])), imgcv)
	box_image = cv2.imread(os.path.join(args.data_path, 'images_with_bbox/{}'.format(stacked[0][0])))
	#print('Path = ',)
	return (h,w,centers,roll_yaw_pitch)

def main_pixel_co(args,f_x,f_y,depth_array): #_pixel_co
	num_classes = 2
	device = args.device

	# Load the model from
	print("Loading model")
	# Create the correct model type
	if args.mrcnn:
		model = get_maskrcnn_model_instance(num_classes)
	else:
		model = get_frcnn_model_instance(num_classes)

	# Load model parameters and keep on CPU
	checkpoint = torch.load(args.weight_file, map_location=device)
	# print(checkpoint)
	# for key, value in checkpoint.items() :
	# 	print(key)
	model.load_state_dict(checkpoint, strict=False)
	model.eval()

	print("Creating data loaders")
	dataset_test = AppleDataset(args.data_path, get_transform(train=False))
	# dataset_test = AppleDataset_data1(args.data_path, get_transform(train=False))
	data_loader_test = torch.utils.data.DataLoader(dataset_test, batch_size=1,
												   shuffle=False, num_workers=1,
												   collate_fn=utils.collate_fn)

	# Create output directory
	base_path = os.path.dirname(args.output_file)
	if not os.path.exists(base_path):
		os.makedirs(base_path)

	# Predict on bboxes on each image
	f = open(args.output_file, 'a')
	for image, targets in data_loader_test:    # for different images one after another
		image = list(img.to(device) for img in image)
		outputs = model(image)
		print('Output = \n',outputs)
		centers = []
		roll_yaw_pitch = []
		for ii, output in enumerate(outputs):  # bbox of one instance(apple) in one image
			# print('ii,output = ', ii, output)
			img_id = targets[ii]['image_id']
			img_name = data_loader_test.dataset.get_img_name(img_id)
			print("Predicting on image: {}".format(img_name))
			boxes = output['boxes'].detach().numpy()
			scores = output['scores'].detach().numpy()

			im_names = np.repeat(img_name, len(boxes), axis=0)
			stacked = np.hstack((im_names.reshape(len(scores), 1), boxes.astype(int), scores.reshape(len(scores), 1)))

			# File to write predictions to
			print('stacked = ',stacked)
			h,w,centers,roll_yaw_pitch = drawBoundingBox(stacked,centers,roll_yaw_pitch,args,f_x,f_y,depth_array)
			print('centers = ',centers)
			# print('h,w,centers = ',h,w,centers)
			np.savetxt(f, stacked, fmt='%s', delimiter=',', newline='\n')
	return(h,w,centers,roll_yaw_pitch)


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""

@author: Adil Zouitine <adilzouitinegm@gmail.com

"""
import numpy as np
from nltk import word_tokenize
import os.path

__all__ = ['Feel']

class Feel():
    '''
    Calculates the emotion of a string with a'bag of word' method.
    object.dict to see the bag of word.

    >>> from pyFeel import Feel
    >>> test = Feel("Ma classe marche bien, c'est génial non ?")
    >>> test.emotions()
    Out : {'positivity': 1.0, 'joy': 0.25, 'fear': 0.0, 'sadness': 0.25,
          'angry': 0.0, 'surprise': 0.0, 'disgust': 0.0}
    '''

    def __init__(self,text):
        '''
        Input : string Output : /
        '''

        scriptpath = os.path.dirname(__file__)
        dict_name = os.path.join(scriptpath, 'feel.npy')

        self.dict = np.load(dict_name).item()
        self.token = [word for word in word_tokenize(text.lower())]
        self.sentiment = {}
        self.number_emotion = 7
        self.vector = np.zeros(self.number_emotion)
        self.list_emotions = ['positivity',
                              'joy',
                              'fear',
                              'sadness',
                              'angry',
                              'surprise',
                              'disgust']

    def emotions(self):
        '''
        Calculates emotion by averaging the emotion vectors of each word in
        the input string and the bag of word.

        Input : (self) Output : dict
        '''

        retain_word = 0
        for word in self.token:
            if word.lower() in self.dict:
                self.vector = np.add(self.vector,
                                     self.dict[word],
                                     out = self.vector)
                retain_word += 1
        if retain_word != 0:
            self.vector /= retain_word
        for emotion,emotion_value in zip(self.list_emotions,self.vector):
            self.sentiment.update({emotion : float(emotion_value)})
        return self.sentiment

if __name__ == '__main__':
    test = Feel("Ma classe fonctionne bien, c'est sympathique non ?")
    print(test.emotions())


import argparse
import glob
import logging
import os
import sys
from functools import partial

import numpy as np
import pandas as pd

sys.path += [os.path.abspath('.'), os.path.abspath('..')]  # Add path to root
from birl.utilities.data_io import create_folder, load_landmarks_csv, save_landmarks_csv
from birl.utilities.dataset import compute_bounding_polygon, inside_polygon, list_sub_folders, parse_path_scale
from birl.utilities.experiments import is_iterable, iterate_mproc_map, parse_arg_params
from birl.utilities.registration import estimate_affine_transform, transform_points
from bm_dataset.rescale_tissue_images import DEFAULT_SCALES, FOLDER_TEMPLATE, NB_WORKERS

def arg_parse_params():
    """ argument parser from cmd

    :return dict:
    """
    # SEE: https://docs.python.org/3/library/argparse.html
    parser = argparse.ArgumentParser()
    parser.add_argument('-a', '--path_annots', type=str, required=False, help='path to folder with annotations')
    parser.add_argument('-d', '--path_dataset', type=str, required=False, help='path to the output directory - dataset')
    parser.add_argument(
        '--scales',
        type=int,
        required=False,
        nargs='*',
        help='generated scales for the dataset',
        default=DEFAULT_SCALES
    )
    parser.add_argument(
        '--nb_selected', type=float, required=False, default=None, help='number ot ration of selected landmarks'
    )
    parser.add_argument(
        '--nb_total', type=int, required=False, default=None, help='total number of generated landmarks'
    )
    parser.add_argument(
        '--nb_workers', type=int, required=False, default=NB_WORKERS, help='number of processes in parallel'
    )
    args = parse_arg_params(parser)
    if not is_iterable(args['scales']):
        args['scales'] = [args['scales']]
    return args

def load_largest_scale(path_set):
    """ in given set find the largest scale and load all landmarks in full size

    :param str path_set: path to image/landmark set
    :return dict: dictionary of ndarray with loaded landmarks in full scale
    """
    scales_folders = [(parse_path_scale(p), os.path.basename(p)) for p in list_sub_folders(path_set)]
    if not scales_folders:
        return
    scale, folder = sorted(scales_folders, reverse=True)[0]

    paths_csv = glob.glob(os.path.join(path_set, folder, '*.csv'))
    scaling = 100. / scale
    names_lnds = {os.path.basename(p): load_landmarks_csv(p) * scaling for p in paths_csv}
    return names_lnds

def generate_random_points_inside(ref_points, nb_extras):
    """ generate some extra points inside the tissue boundary polygon

    :param ref_points: point of the tissue
    :param int nb_extras: number of point to be added
    :return [(int, int)]: extra points
    """
    # tighter approximation, not all tissue is really convex
    convex_polygon = compute_bounding_polygon(ref_points)
    poly_mins = np.min(convex_polygon, axis=0)
    poly_size = np.max(convex_polygon, axis=0) - poly_mins
    # generate sample points inside polygon
    points_extra = []
    for i in range(int(1e5)):
        point = (np.random.random(2) * poly_size + poly_mins).astype(int)
        if inside_polygon(convex_polygon, point):
            points_extra.append(point)
            if len(points_extra) > nb_extras:
                logging.debug('particular polygon generated %f inside', nb_extras / float(i))
                break
    else:  # in case the loop ended regularly without break
        logging.warning('something went wrong with ')
    return points_extra

def expand_random_warped_landmarks(names_lnds, names_lnds_new, nb_total):
    """ add some extra point which are randomly sampled in the first sample
    and warped to the other images using estimated affine transform

    :param dict names_lnds: mapping to ndarray of the original landmarks
    :param dict names_lnds_new: mapping to ndarray of the generated landmarks
    :param int nb_total: total number of point - landmarks
    :return dict: mapping to ndarray
    """
    # estimate then number of required points
    nb_min_new = min(map(len, names_lnds_new.values()))
    nb_extras = nb_total - nb_min_new
    if nb_extras <= 0:
        return names_lnds_new

    ref_name = sorted(names_lnds)[0]
    ref_points = names_lnds[ref_name]
    points_extra = generate_random_points_inside(ref_points, nb_extras)

    for name in filter(lambda n: n != ref_name, names_lnds):
        # prepare the points
        nb_common = min([len(names_lnds[ref_name]), len(names_lnds[name])])
        pts1 = names_lnds[ref_name][:nb_common]
        pts2 = names_lnds[name][:nb_common]
        # estimate the internal affine transformation
        matrix, _, _, _ = estimate_affine_transform(pts1, pts2)
        points_warp = transform_points(points_extra, matrix)
        # insert the warped points
        names_lnds_new[name] = np.vstack([names_lnds_new[name][:nb_min_new], points_warp])
    # insert also the reference sample
    names_lnds_new[ref_name] = np.vstack([names_lnds_new[ref_name][:nb_min_new], points_extra])

    # reorder landmarks but equally in all sets
    reorder = list(range(nb_total))
    np.random.shuffle(reorder)
    names_lnds_new = {n: names_lnds_new[n][reorder] for n in names_lnds_new}
    return names_lnds_new

def extend_landmarks(path_set, path_dataset, nb_selected=None, nb_total=None):
    """ select and extend the original set of landmarks

    :param str path_set: path to the particular set if images/landmarks
    :param str path_dataset: root path to generated dataset
    :param float|int|None nb_selected: portion of selected points,
        if None use all original landmarks
    :param int|None nb_total: add extra points up to total number,
        if None, no adding extra points
    :return:
    """
    logging.debug('> processing: %s', path_set)

    # search form mas scale in set and load all related landmarks
    names_lnds = load_largest_scale(path_set)
    if not names_lnds:
        logging.warning('no landmarks was loaded for "%s"', path_set)
        return

    # select subset of selected landmarks
    names_lnds_new = {}
    if nb_selected is not None:
        if nb_selected < 0:
            raise ValueError('number of selected has to be positive')
        lens = list(map(len, names_lnds.values()))
        # transform the relative count to absolute number
        if nb_selected < 1:
            nb_selected = np.ceil(nb_selected * max(lens)).astype(int)
        # perform the selection
        indexes = list(range(min(lens)))
        np.random.shuffle(indexes)
        # just a required subset
        indexes = indexes[:nb_selected]
        for name in names_lnds:
            names_lnds_new[name] = names_lnds[name][indexes]
    else:
        names_lnds_new = names_lnds

    if nb_total is not None:
        names_lnds_new = expand_random_warped_landmarks(names_lnds, names_lnds_new, nb_total)

    # export the landmarks
    path_set_scale = os.path.join(path_dataset, os.path.basename(path_set), FOLDER_TEMPLATE % 100)
    create_folder(path_set_scale)
    for name, val in names_lnds_new.items():
        save_landmarks_csv(os.path.join(path_set_scale, name), val)

def dataset_expand_landmarks(path_annots, path_dataset, nb_selected=None, nb_total=None, nb_workers=NB_WORKERS):
    """ select and expand over whole dataset

    :param str path_annots: root path to original dataset
    :param str path_dataset: root path to generated dataset
    :param float|int|None nb_selected: portion of selected points
    :param int|None nb_total: add extra points up to total number
    :param int nb_workers: number of jobs running in parallel
    :return list(int):
    """
    list_sets = list_sub_folders(path_annots)
    logging.info('Found sets: %i', len(list_sets))

    _wrap_extend = partial(extend_landmarks, path_dataset=path_dataset, nb_selected=nb_selected, nb_total=nb_total)
    counts = list(iterate_mproc_map(_wrap_extend, sorted(list_sets), nb_workers=nb_workers, desc='expand landmarks'))
    return counts

def scale_set_landmarks(path_set, scales=DEFAULT_SCALES):
    """ scale the updated (generated) landmarks

    the scales are created within the same path set

    :param str path_set: path to the image/landmark set
    :param list(int) scales: created scales
    :return dict:
    """
    logging.debug('> processing: %s', path_set)
    path_scale100 = os.path.join(path_set, FOLDER_TEMPLATE % 100)
    if not os.path.isdir(path_scale100):
        logging.error('missing base scale 100pc in "%s"', path_scale100)
        return
    list_csv = glob.glob(os.path.join(path_scale100, '*.csv'))
    logging.debug('>> found landmarks: %i', len(list_csv))
    dict_lnds = {os.path.basename(p): pd.read_csv(p, index_col=0) for p in list_csv}
    set_scales = {}
    for sc in (sc for sc in scales if sc not in [100]):  # drop the base scale
        folder_name = FOLDER_TEMPLATE % sc
        path_scale = create_folder(os.path.join(path_set, folder_name))
        for name in dict_lnds:
            df_scale = dict_lnds[name] * (sc / 100.)
            df_scale.to_csv(os.path.join(path_scale, name))
        set_scales[sc] = len(dict_lnds)
    dict_lens = {os.path.basename(path_set): set_scales}
    return dict_lens

def dataset_scale_landmarks(path_dataset, scales=DEFAULT_SCALES, nb_workers=NB_WORKERS):
    """ generate several scales within the same dataset

    :param str path_dataset: path to the source/generated dataset
    :param [inr] scales: created scales
    :param int nb_workers: number of jobs running in parallel
    :return:
    """
    list_sets = list_sub_folders(path_dataset)
    logging.info('Found sets: %i', len(list_sets))

    _wrap_scale = partial(scale_set_landmarks, scales=scales)
    counts = list(iterate_mproc_map(_wrap_scale, sorted(list_sets), nb_workers=nb_workers, desc='scaling sets'))
    return counts

def main(path_annots, path_dataset, scales, nb_selected=None, nb_total=None, nb_workers=NB_WORKERS):
    """ main entry point

    :param str path_annots: root path to original dataset
    :param str path_dataset: root path to generated dataset
    :param list(int) scales: generated scales
    :param float|int|None nb_selected: portion of selected points
    :param int|None nb_total: add extra points up to total number
    :param int nb_workers: number of jobs running in parallel
    :return tuple(int,int):
    """
    count_gene = dataset_expand_landmarks(path_annots, path_dataset, nb_selected, nb_total, nb_workers=nb_workers)
    count_scale = dataset_scale_landmarks(path_dataset, scales=scales, nb_workers=nb_workers)
    return count_gene, count_scale

if __name__ == '__main__':
    logging.basicConfig(level=logging.INFO)

    arg_params = arg_parse_params()
    logging.info('running...')
    main(**arg_params)
    logging.info('DONE')


# input.txtに記述した内容を標準入力として受け取り実行する
# 提出する時にはこの部分は含めない
import io,sys
with open("input.txt") as txt_opend:
    TXT_INPUT = txt_opend.read()
sys.stdin=io.StringIO(TXT_INPUT)
# --------------------------------------------------------
K = int(input())  # nは入力回数
int_list = list(map(int, input().split()))  # n個の数字の入力を受け取る

sum = 0
ans = 0
for l in int_list:
    sum+=l
    ans = sum%10
print(int(ans))


# -*- coding: utf-8 -*-
"""
Created on Thu Sep 28 16:12:30 2017

@author: tkoller
"""
import os.path

import numpy as np
import pytest
from casadi import MX, Function
from casadi import reshape as cas_reshape

try:
    from safe_exploration.ssm_gpy.gaussian_process import SimpleGPModel
    from GPy.kern import RBF
    _has_ssm_gpy = True
except:
    _has_ssm_gpy = False

from .. import gp_reachability as reach_num
from .. import gp_reachability_casadi as reach_cas
from ..utils import array_of_vec_to_array_of_mat

a_tol = 1e-5
r_tol = 1e-4

# @pytest.fixture(params = [("InvPend",True,True)])
@pytest.fixture(params=[("InvPend", True, True), ("InvPend", False, True),
                        ("InvPend", True, True), ("InvPend", False, True)])
def before_test_onestep_reachability(request):
    if not _has_ssm_gpy:
        pytest.skip("Test requires optional dependencies 'ssm_gpy'")

    np.random.seed(125)

    env, init_uncertainty, lin_model = request.param
    n_s = 2
    n_u = 1
    c_safety = 2
    a = None
    b = None
    if lin_model:
        a = np.random.rand(n_s, n_s)
        b = np.random.rand(n_s, n_u)

    train_data = np.load(os.path.join(os.path.dirname(__file__), 'invpend_data.npz'))
    X = train_data["X"]
    y = train_data["y"]
    m = 50
    #kerns = [RBF(n_s+n_u)]*n_s
    gp = SimpleGPModel(n_s, n_s, n_u, X, y, m, train=False)
    gp.train(X, y, m, opt_hyp=True, choose_data=False)
    L_mu = np.array([0.001] * n_s)
    L_sigm = np.array([0.001] * n_s)
    k_fb = np.random.rand(n_u, n_s)  # need to choose this appropriately later
    k_ff = np.random.rand(n_u, 1)

    p = .1 * np.random.randn(n_s, 1)
    if init_uncertainty:
        q = .2 * np.array(
            [[.5, .2], [.2, .65]])  # reachability based on previous uncertainty
    else:
        q = None  # no initial uncertainty

    return p, q, gp, k_fb, k_ff, L_mu, L_sigm, c_safety, a, b

def test_onestep_reachability(before_test_onestep_reachability):
    """ do we get the same results as the numpy equivalent?"""

    p, q, gp, k_fb, k_ff, L_mu, L_sigm, c_safety, a, b = before_test_onestep_reachability

    n_u, n_s = np.shape(k_fb)

    k_fb_cas = MX.sym("k_fb", (n_u, n_s))
    k_ff_cas = MX.sym("k_ff", (n_u, 1))

    p_new_cas, q_new_cas, _ = reach_cas.onestep_reachability(p, gp, k_ff_cas, L_mu,
                                                             L_sigm, q, k_fb_cas,
                                                             c_safety, a=a, b=b)
    f = Function("f", [k_fb_cas, k_ff_cas], [p_new_cas, q_new_cas])

    f_out_cas = f(k_fb, k_ff)

    f_out_num = reach_num.onestep_reachability(p, gp, k_ff, L_mu, L_sigm, q, k_fb,
                                               c_safety, 0, a=a, b=b)

    assert np.allclose(f_out_cas[0],
                       f_out_num[0]), "Are the centers of the next state the same?"
    assert np.allclose(f_out_cas[1], f_out_num[
        1]), "Are the shape matrices of the next state the same?"

# @pytest.mark.xfail
def test_multistep_reachability(before_test_onestep_reachability):
    """ """
    p, _, gp, k_fb, _, L_mu, L_sigm, c_safety, a, b = before_test_onestep_reachability
    T = 3

    n_u, n_s = np.shape(k_fb)

    u_0 = .2 * np.random.randn(n_u, 1)
    k_fb_0 = np.random.randn(T - 1,
                             n_s * n_u)  # np.zeros((T-1,n_s*n_u))# np.random.randn(T-1,n_s*n_u)
    k_ff = np.random.randn(T - 1, n_u)
    # k_fb_ctrl = np.zeros((n_u,n_s))#np.random.randn(n_u,n_s)

    u_0_cas = MX.sym("u_0", (n_u, 1))
    k_fb_cas_0 = MX.sym("k_fb", (T - 1, n_u * n_s))
    k_ff_cas = MX.sym("k_ff", (T - 1, n_u))

    p_new_cas, q_new_cas, _ = reach_cas.multi_step_reachability(p, u_0, k_fb_cas_0,
                                                                k_ff_cas, gp, L_mu,
                                                                L_sigm, c_safety, a, b)
    f = Function("f", [u_0_cas, k_fb_cas_0, k_ff_cas], [p_new_cas, q_new_cas])

    k_fb_0_cas = np.copy(k_fb_0)  # np.copy(k_fb_0)

    for i in range(T - 1):
        k_fb_0_cas[i, None, :] = k_fb_0_cas[i, None, :] + cas_reshape(k_fb,
                                                                      (1, n_u * n_s))
    p_all_cas, q_all_cas = f(u_0, k_fb_0_cas, k_ff)

    k_ff_all = np.vstack((u_0.T, k_ff))

    k_fb_apply = array_of_vec_to_array_of_mat(k_fb_0, n_u, n_s)

    for i in range(T - 1):
        k_fb_apply[i, :, :] += k_fb

    _, _, p_all_num, q_all_num = reach_num.multistep_reachability(p, gp, k_fb_apply,
                                                                  k_ff_all, L_mu,
                                                                  L_sigm, None,
                                                                  c_safety, 0, a, b,
                                                                  None)

    assert np.allclose(p_all_cas, p_all_num, r_tol,
                       a_tol), "Are the centers of the ellipsoids same?"

    assert np.allclose(q_all_cas[0, :], q_all_num[0, :, :].reshape((-1, n_s * n_s)),
                       r_tol, a_tol), "Are the first shape matrices the same?"
    assert np.allclose(q_all_cas[1, :], q_all_num[1, :, :].reshape((-1, n_s * n_s)),
                       r_tol, a_tol), "Are the second shape matrices the same?"
    # assert np.allclose(q_all_cas[1,:],q_all_num[1,:].reshape((-1,n_s*n_s))), "Are the second shape matrices the same?"
    assert np.allclose(q_all_cas[-1, :], q_all_num[-1, :, :].reshape((-1, n_s * n_s)),
                       r_tol, a_tol), "Are the last shape matrices the same?"
    assert np.allclose(q_all_cas, q_all_num.reshape((T, n_s * n_s)), r_tol,
                       a_tol), "Are the shape matrices the same?"