#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
A Metropolis-Hastings and Gibbs samplers for a state-ful model.

"""

import numpy as np
from emcee import autocorr
from emcee.sampler import Sampler
import logging
import h5py
from Starfish import constants as C

class StateSampler(Sampler):
    """
    The most basic possible Metropolis-Hastings style MCMC sampler, but with
    optional callbacks for acceptance and rejection.

    :param cov:
        The covariance matrix to use for the proposal distribution.

    :param dim:
        Number of dimensions in the parameter space.

    :param lnpostfn:
        A function that takes a vector in the parameter space as input and
        returns the natural logarithm of the posterior probability for that
        position.

    :param query_lnprob:
        A function which queries the model for the current lnprob value. This is
        because an intermediate sampler might have intervened.

    :param rejectfn: (optional)
        A function to be executed if the proposal is rejected. Generally this
        should be a function which reverts the model to the previous parameter
        values. Might need to be a closure.

    :param acceptfn: (optional)
        A function to be executed if the proposal is accepted.

    :param args: (optional)
        A list of extra positional arguments for ``lnpostfn``. ``lnpostfn``
        will be called with the sequence ``lnpostfn(p, *args, **kwargs)``.

    :param kwargs: (optional)
        A list of extra keyword arguments for ``lnpostfn``. ``lnpostfn``
        will be called with the sequence ``lnpostfn(p, *args, **kwargs)``.

    """
    def __init__(self, lnprob, p0, cov, query_lnprob=None, rejectfn=None,
        acceptfn=None, debug=False, outdir="", *args, **kwargs):
        dim = len(p0)
        super().__init__(dim, lnprob, *args, **kwargs)
        self.cov = cov
        self.p0 = p0
        self.query_lnprob = query_lnprob
        self.rejectfn = rejectfn
        self.acceptfn = acceptfn
        self.logger = logging.getLogger(self.__class__.__name__)
        self.outdir = outdir
        self.debug = debug
        if self.debug:
            self.logger.setLevel(logging.DEBUG)
        else:
            self.logger.setLevel(logging.INFO)

    def reset(self):
        super().reset()
        self._chain = np.empty((0, self.dim))
        self._lnprob = np.empty(0)

    def sample(self, p0, lnprob0=None, randomstate=None, thin=1,
               storechain=True, iterations=1, incremental_save=0, **kwargs):
        """
        Advances the chain ``iterations`` steps as an iterator

        :param p0:
            The initial position vector.

        :param lnprob0: (optional)
            The log posterior probability at position ``p0``. If ``lnprob``
            is not provided, the initial value is calculated.

        :param rstate0: (optional)
            The state of the random number generator. See the
            :func:`random_state` property for details.

        :param iterations: (optional)
            The number of steps to run.

        :param thin: (optional)
            If you only want to store and yield every ``thin`` samples in the
            chain, set thin to an integer greater than 1.

        :param storechain: (optional)
            By default, the sampler stores (in memory) the positions and
            log-probabilities of the samples in the chain. If you are
            using another method to store the samples to a file or if you
            don't need to analyse the samples after the fact (for burn-in
            for example) set ``storechain`` to ``False``.

        At each iteration, this generator yields:

        * ``pos`` - The current positions of the chain in the parameter
          space.

        * ``lnprob`` - The value of the log posterior at ``pos`` .

        * ``rstate`` - The current state of the random number generator.

        """

        self.random_state = randomstate

        p = np.array(p0)
        if lnprob0 is None:
            # See if there's something there
            lnprob0 = self.query_lnprob()

            # If not, we're on the first iteration
            if lnprob0 is None:
                lnprob0 = self.get_lnprob(p)

        # Resize the chain in advance.
        if storechain:
            N = int(iterations / thin)
            self._chain = np.concatenate((self._chain,
                                          np.zeros((N, self.dim))), axis=0)
            self._lnprob = np.append(self._lnprob, np.zeros(N))

        i0 = self.iterations

        # Use range instead of xrange for python 3 compatability
        for i in range(int(iterations)):
            self.iterations += 1

            # Since the independent nuisance sampling may have changed parameters,
            # query each process for the current lnprob
            lnprob0 = self.query_lnprob()
            self.logger.debug("Queried lnprob: {}".format(lnprob0))

            # Calculate the proposal distribution.
            if self.dim == 1:
                q = self._random.normal(loc=p[0], scale=self.cov[0], size=(1,))
            else:
                q = self._random.multivariate_normal(p, self.cov)

            newlnprob = self.get_lnprob(q)
            diff = newlnprob - lnprob0
            self.logger.debug("old lnprob: {}".format(lnprob0))
            self.logger.debug("proposed lnprob: {}".format(newlnprob))

            # M-H acceptance ratio
            if diff < 0:
                diff = np.exp(diff) - self._random.rand()
                if diff < 0:
                    #Reject the proposal and revert the state of the model
                    self.logger.debug("Proposal rejected")
                    if self.rejectfn is not None:
                        self.rejectfn()

            if diff > 0:
                #Accept the new proposal
                self.logger.debug("Proposal accepted")
                p = q
                lnprob0 = newlnprob
                self.naccepted += 1
                if self.acceptfn is not None:
                    self.acceptfn()

            if storechain and i % thin == 0:
                ind = i0 + int(i / thin)
                self._chain[ind, :] = p
                self._lnprob[ind] = lnprob0

            # The default of 0 evaluates to False
            if incremental_save:
                if (((i+1) % incremental_save) == 0) & (i > 0):
                    np.save('chain_backup.npy', self._chain)

            # Heavy duty iterator action going on right here...
            yield p, lnprob0, self.random_state

    @property
    def acor(self):
        """
        An estimate of the autocorrelation time for each parameter (length:
        ``dim``).

        """
        return self.get_autocorr_time()

    def get_autocorr_time(self, window=50):
        """
        Compute an estimate of the autocorrelation time for each parameter
        (length: ``dim``).

        :param window: (optional)
            The size of the windowing function. This is equivalent to the
            maximum number of lags to use. (default: 50)

        """
        return autocorr.integrated_time(self.chain, axis=0, window=window)

    def write(self, fname="mc.hdf5"):
        '''
        Write the samples to an HDF file.

        flatchain
        acceptance fraction

        Everything can be saved in the dataset self.fname

        '''

        filename = self.outdir + fname
        self.logger.debug("Opening {} for writing HDF5 flatchains".format(filename))
        hdf5 = h5py.File(filename, "w")
        samples = self.flatchain

        dset = hdf5.create_dataset("samples", samples.shape, compression='gzip', compression_opts=9)
        dset[:] = samples
        dset.attrs["acceptance"] = "{}".format(self.acceptance_fraction)
        dset.attrs["commit"] = "{}".format(C.get_git_commit())
        hdf5.close()

# class PSampler(ParallelSampler):
#     '''
#     Subclasses the GibbsSampler in emcee
#
#     :param cov:
#     :param starting_param_dict: the dictionary of starting parameters
#     :param cov: the MH proposal
#     :param revertfn:
#     :param acceptfn:
#     :param debug:
#
#     '''
#
#     def __init__(self, **kwargs):
#         self.dim = len(self.param_tuple)
#         #p0 = np.empty((self.dim,))
#         #starting_param_dict = kwargs.get("starting_param_dict")
#         #for i,param in enumerate(self.param_tuple):
#         #    p0[i] = starting_param_dict[param]
#
#         kwargs.update({"dim":self.dim})
#         #self.spectra_list = kwargs.get("spectra_list", [0])
#
#         super(PSampler, self).__init__(**kwargs)
#
#         #Each subclass will have to overwrite how it parses the param_dict into the correct order
#         #and sets the param_tuple
#
#         #SUBCLASS here and define self.param_tuple
#         #SUBCLASS here and define self.lnprob
#         #SUBCLASS here and do self.revertfn
#         #then do super().__init__() to call the following code
#
#         self.outdir = kwargs.get("outdir", "")
#
#     def startdict_to_tuple(self, startdict):
#         raise NotImplementedError("To be implemented by a subclass!")
#
#     def zip_p(self, p):
#         return dict(zip(self.param_tuple, p))
#
#     def lnprob(self):
#         raise NotImplementedError("To be implemented by a subclass!")
#
#     def revertfn(self):
#         raise NotImplementedError("To be implemented by a subclass!")
#
#     def acceptfn(self):
#         raise NotImplementedError("To be implemented by a subclass!")
#
#     def write(self):
#         '''
#         Write all of the relevant sample output to an HDF file.
#
#         Write the lnprobability to an HDF file.
#
#         flatchain
#         acceptance fraction
#         tuple parameters as an attribute in the header from self.param_tuple
#
#         The actual HDF5 file is structured as follows
#
#         /
#         stellar parameters.flatchain
#         00/
#         ...
#         22/
#         23/
#             global_cov.flatchain
#             regions/
#                 region1.flatchain
#
#         Everything can be saved in the dataset self.fname
#
#         '''
#
#         filename = self.outdir + "flatchains.hdf5"
#         self.logger.debug("Opening {} for writing HDF5 flatchains".format(filename))
#         hdf5 = h5py.File(filename, "w")
#         samples = self.flatchain
#
#         self.logger.debug("Creating dataset with fname:{}".format(self.fname))
#         dset = hdf5.create_dataset(self.fname, samples.shape, compression='gzip', compression_opts=9)
#         self.logger.debug("Storing samples and header attributes.")
#         dset[:] = samples
#         dset.attrs["parameters"] = "{}".format(self.param_tuple)
#         dset.attrs["acceptance"] = "{}".format(self.acceptance_fraction)
#         dset.attrs["acor"] = "{}".format(self.acor)
#         dset.attrs["commit"] = "{}".format(C.get_git_commit())
#         hdf5.close()
#
#         #lnprobability is the lnprob at each sample
#         filename = self.outdir + "lnprobs.hdf5"
#         self.logger.debug("Opening {} for writing HDF5 lnprobs".format(filename))
#         hdf5 = h5py.File(filename, "w") #creates if doesn't exist, otherwise read/write
#         lnprobs = self.lnprobability
#
#         dset = hdf5.create_dataset(self.fname, samples.shape[:1], compression='gzip', compression_opts=9)
#         dset[:] = lnprobs
#         dset.attrs["commit"] = "{}".format(C.get_git_commit())
#         hdf5.close()
#
#     def plot(self, triangle_plot=False):
#         '''
#         Generate the relevant plots once the sampling is done.
#         '''
#         samples = self.flatchain
#
#         plot_walkers(self.outdir + self.fname + "_chain_pos.png", samples, labels=self.param_tuple)
#
#         if triangle_plot:
#             import triangle
#             figure = triangle.corner(samples, labels=self.param_tuple, quantiles=[0.16, 0.5, 0.84],
#                                      show_titles=True, title_args={"fontsize": 12})
#             figure.savefig(self.outdir + self.fname + "_triangle.png")
#
#             plt.close(figure)
#
# class StellarSampler(PSampler):
#     """
#     Subclasses the Sampler specifically for stellar parameters
#
#     """
#     def __init__(self, **kwargs):
#         '''
#         :param pconns: Collection of parent ends of the PIPEs
#         :type pconns: dict
#
#         :param starting_param_dict:
#             the dictionary of starting parameters
#
#         :param cov:
#             the MH proposal
#
#         :param fix_logg:
#             fix logg? If so, to what value?
#
#         :param debug:
#
#         :param args: []
#         '''
#
#         self.fix_logg = kwargs.get("fix_logg", False)
#         starting_pram_dict = kwargs.get("starting_param_dict")
#         self.param_tuple = self.startdict_to_tuple(starting_pram_dict)
#         print("param_tuple is {}".format(self.param_tuple))
#         self.p0 = np.array([starting_pram_dict[key] for key in self.param_tuple])
#
#         kwargs.update({"p0":self.p0, "revertfn":self.revertfn, "acceptfn": self.acceptfn, "lnprobfn":self.lnprob})
#         super(StellarSampler, self).__init__(**kwargs)
#
#         #self.pconns is a dictionary of parent connections to each PIPE connecting to the child processes.
#         self.spectrum_ids = sorted(self.pconns.keys())
#         self.fname = "stellar"
#
#     def startdict_to_tuple(self, startdict):
#         tup = ()
#         for param in C.stellar_parameters:
#             #check if param is in keys, if so, add to the tuple
#             if param in startdict:
#                 tup += (param,)
#         return tup
#
#     def reset(self):
#         super(StellarSampler, self).reset()
#
#     def revertfn(self):
#         '''
#         Revert the model to the previous state of parameters, in the case of a rejected MH proposal.
#         '''
#         self.logger.debug("reverting stellar parameters")
#         self.prior = self.prior_last
#
#         #Decide we don't want these stellar params. Tell the children to reject the proposal.
#         for pconn in self.pconns.values():
#             pconn.send(("DECIDE", False))
#
#     def acceptfn(self):
#         '''
#         Execute this if the MH proposal is accepted.
#         '''
#         self.logger.debug("accepting stellar parameters")
#         #Decide we do want to keep these stellar params. Tell the children to accept the proposal.
#         for pconn in self.pconns.values():
#             pconn.send(("DECIDE", True))
#
#     def lnprob(self, p):
#         # We want to send the same stellar parameters to each model,
#         # but also send the different vz and logOmega parameters
#         # to the separate spectra, based upon spectrum_id.
#         #self.logger.debug("StellarSampler lnprob p is {}".format(p))
#
#         #Extract only the temp, logg, Z, vsini parameters
#         if not self.fix_logg:
#             params = self.zip_p(p[:4])
#             others = p[4:]
#         else:
#             #Coming in as temp, Z, vsini, vz, logOmega...
#             params = self.zip_p(p[:3])
#             others = p[3:]
#             params.update({"logg": self.fix_logg})
#
#         # Prior
#         self.prior_last = self.prior
#
#         logg = params["logg"]
#         self.prior = -0.5 * (logg - 5.0)**2/(0.05)**2
#
#         #others should now be either [vz, logOmega] or [vz0, logOmega0, vz1, logOmega1, ...] etc. Always div by 2.
#         #split p up into [vz, logOmega], [vz, logOmega] pairs that update the other parameters.
#         #mparams is now a list of parameter dictionaries
#
#         #Now, pack up mparams into a dictionary to send the right stellar parameters to the right subprocesses
#         mparams = {}
#         for (spectrum_id, order_id), (vz, logOmega) in zip(self.spectrum_ids, grouper(others, 2)):
#             p = params.copy()
#             p.update({"vz":vz, "logOmega":logOmega})
#             mparams[spectrum_id] = p
#
#         self.logger.debug("updated lnprob params: {}".format(mparams))
#
#         lnps = np.empty((self.nprocs,))
#
#         #Distribute the calculation to each process
#         self.logger.debug("Distributing params to children")
#         for ((spectrum_id, order_id), pconn) in self.pconns.items():
#             #Parse the parameters into what needs to be sent to each Model here.
#             pconn.send(("LNPROB", mparams[spectrum_id]))
#
#         #Collect the answer from each process
#         self.logger.debug("Collecting params from children")
#         for i, pconn in enumerate(self.pconns.values()):
#             lnps[i] = pconn.recv()
#
#         self.logger.debug("lnps : {}".format(lnps))
#         s = np.sum(lnps)
#         self.logger.debug("sum lnps {}".format(s))
#         return s + self.prior
#
#
#
#
# # From emcee
# class GibbsSampler(Sampler):
#     """
#     The most basic possible Metropolis-Hastings style MCMC sampler, designed to be
#     encapsulated as part of a state-ful Gibbs sampler.
#
#     :param cov:
#         The covariance matrix to use for the proposal distribution.
#
#     :param dim:
#         Number of dimensions in the parameter space.
#
#     :param lnpostfn:
#         A function that takes a vector in the parameter space as input and
#         returns the natural logarithm of the posterior probability for that
#         position.
#
#     :param revertfn:
#         A function which reverts the model to the previous parameter values, in the
#         case that the proposal parameters are rejected.
#
#     :param acceptfn:
#         A function to be executed if the proposal is accepted.
#
#     :param resample:
#         Redraw the previous lnprob at these parameters? Designed to work with the fact that an
#         Emulator is non-deterministic and prevent it from getting caught at a maximum.
#
#     :param args: (optional)
#         A list of extra positional arguments for ``lnpostfn``. ``lnpostfn``
#         will be called with the sequence ``lnpostfn(p, *args, **kwargs)``.
#
#     :param kwargs: (optional)
#         A list of extra keyword arguments for ``lnpostfn``. ``lnpostfn``
#         will be called with the sequence ``lnpostfn(p, *args, **kwargs)``.
#
#     """
#     def __init__(self, cov, p0, *args, **kwargs):
#         super(GibbsSampler, self).__init__(*args, **kwargs)
#         self.cov = cov
#         self.p0 = p0
#         self.revertfn = kwargs.get("revertfn", None)
#         self.acceptfn = kwargs.get("acceptfn", None)
#         self.logger = logging.getLogger(self.__class__.__name__)
#         self.debug = kwargs.get("debug", False)
#         if self.debug:
#             self.logger.setLevel(logging.DEBUG)
#         else:
#             self.logger.setLevel(logging.INFO)
#
#     def reset(self):
#         super(GibbsSampler, self).reset()
#         self._chain = np.empty((0, self.dim))
#         self._lnprob = np.empty(0)
#
#     def sample(self, p0, lnprob0=None, randomstate=None, thin=1,
#                storechain=True, iterations=1, **kwargs):
#         """
#         Advances the chain ``iterations`` steps as an iterator
#
#         :param p0:
#             The initial position vector.
#
#         :param lnprob0: (optional)
#             The log posterior probability at position ``p0``. If ``lnprob``
#             is not provided, the initial value is calculated.
#
#         :param rstate0: (optional)
#             The state of the random number generator. See the
#             :func:`random_state` property for details.
#
#         :param iterations: (optional)
#             The number of steps to run.
#
#         :param thin: (optional)
#             If you only want to store and yield every ``thin`` samples in the
#             chain, set thin to an integer greater than 1.
#
#         :param storechain: (optional)
#             By default, the sampler stores (in memory) the positions and
#             log-probabilities of the samples in the chain. If you are
#             using another method to store the samples to a file or if you
#             don't need to analyse the samples after the fact (for burn-in
#             for example) set ``storechain`` to ``False``.
#
#         At each iteration, this generator yields:
#
#         * ``pos`` - The current positions of the chain in the parameter
#           space.
#
#         * ``lnprob`` - The value of the log posterior at ``pos`` .
#
#         * ``rstate`` - The current state of the random number generator.
#
#         """
#
#         self.random_state = randomstate
#
#         p = np.array(p0)
#         if lnprob0 is None:
#             lnprob0 = self.get_lnprob(p)
#
#         # Resize the chain in advance.
#         if storechain:
#             N = int(iterations / thin)
#             self._chain = np.concatenate((self._chain,
#                                           np.zeros((N, self.dim))), axis=0)
#             self._lnprob = np.append(self._lnprob, np.zeros(N))
#
#         i0 = self.iterations
#         # Use range instead of xrange for python 3 compatability
#         for i in range(int(iterations)):
#             self.iterations += 1
#
#             # Calculate the proposal distribution.
#             if self.dim == 1:
#                 q = self._random.normal(loc=p[0], scale=self.cov[0], size=(1,))
#             else:
#                 q = self._random.multivariate_normal(p, self.cov)
#
#             newlnprob = self.get_lnprob(q)
#             diff = newlnprob - lnprob0
#             self.logger.debug("old lnprob: {}".format(lnprob0))
#             self.logger.debug("proposed lnprob: {}".format(newlnprob))
#
#             # M-H acceptance ratio
#             if diff < 0:
#                 diff = np.exp(diff) - self._random.rand()
#                 if diff < 0:
#                     #Reject the proposal and revert the state of the model
#                     self.logger.debug("Proposal rejected")
#                     if self.revertfn is not None:
#                         self.revertfn()
#
#             if diff > 0:
#                 #Accept the new proposal
#                 self.logger.debug("Proposal accepted")
#                 p = q
#                 lnprob0 = newlnprob
#                 self.naccepted += 1
#                 if self.acceptfn is not None:
#                     self.acceptfn()
#
#             if storechain and i % thin == 0:
#                 ind = i0 + int(i / thin)
#                 self._chain[ind, :] = p
#                 self._lnprob[ind] = lnprob0
#
#             # Heavy duty iterator action going on right here...
#             yield p, lnprob0, self.random_state
#
#     @property
#     def acor(self):
#         """
#         An estimate of the autocorrelation time for each parameter (length:
#         ``dim``).
#
#         """
#         return self.get_autocorr_time()
#
#     def get_autocorr_time(self, window=50):
#         """
#         Compute an estimate of the autocorrelation time for each parameter
#         (length: ``dim``).
#
#         :param window: (optional)
#             The size of the windowing function. This is equivalent to the
#             maximum number of lags to use. (default: 50)
#
#         """
#         return autocorr.integrated_time(self.chain, axis=0, window=window)
#
# class SSampler(GibbsSampler):
#     '''
#     Subclasses the GibbsSampler in emcee
#
#     :param cov:
#     :param starting_param_dict: the dictionary of starting parameters
#     :param cov: the MH proposal
#     :param revertfn:
#     :param acceptfn:
#     :param debug:
#
#     '''
#
#     def __init__(self, **kwargs):
#         self.dim = len(self.param_tuple)
#         #p0 = np.empty((self.dim,))
#         #starting_param_dict = kwargs.get("starting_param_dict")
#         #for i,param in enumerate(self.param_tuple):
#         #    p0[i] = starting_param_dict[param]
#
#         kwargs.update({"dim":self.dim})
#         #self.spectra_list = kwargs.get("spectra_list", [0])
#
#         super(Sampler, self).__init__(**kwargs)
#
#         #Each subclass will have to overwrite how it parses the param_dict into the correct order
#         #and sets the param_tuple
#
#         #SUBCLASS here and define self.param_tuple
#         #SUBCLASS here and define self.lnprob
#         #SUBCLASS here and do self.revertfn
#         #then do super().__init__() to call the following code
#
#         self.outdir = kwargs.get("outdir", "")
#
#     def startdict_to_tuple(self, startdict):
#         raise NotImplementedError("To be implemented by a subclass!")
#
#     def zip_p(self, p):
#         return dict(zip(self.param_tuple, p))
#
#     def lnprob(self):
#         raise NotImplementedError("To be implemented by a subclass!")
#
#     def revertfn(self):
#         raise NotImplementedError("To be implemented by a subclass!")
#
#     def acceptfn(self):
#         raise NotImplementedError("To be implemented by a subclass!")
#
#     def write(self):
#         '''
#         Write all of the relevant sample output to an HDF file.
#
#         Write the lnprobability to an HDF file.
#
#         flatchain
#         acceptance fraction
#         tuple parameters as an attribute in the header from self.param_tuple
#
#         The actual HDF5 file is structured as follows
#
#         /
#         stellar parameters.flatchain
#         00/
#         ...
#         22/
#         23/
#             global_cov.flatchain
#             regions/
#                 region1.flatchain
#
#         Everything can be saved in the dataset self.fname
#
#         '''
#
#         filename = self.outdir + "flatchains.hdf5"
#         self.logger.debug("Opening {} for writing HDF5 flatchains".format(filename))
#         hdf5 = h5py.File(filename, "w")
#         samples = self.flatchain
#
#         self.logger.debug("Creating dataset with fname:{}".format(self.fname))
#         dset = hdf5.create_dataset(self.fname, samples.shape, compression='gzip', compression_opts=9)
#         self.logger.debug("Storing samples and header attributes.")
#         dset[:] = samples
#         dset.attrs["parameters"] = "{}".format(self.param_tuple)
#         dset.attrs["acceptance"] = "{}".format(self.acceptance_fraction)
#         dset.attrs["acor"] = "{}".format(self.acor)
#         dset.attrs["commit"] = "{}".format(C.get_git_commit())
#         hdf5.close()
#
#         #lnprobability is the lnprob at each sample
#         filename = self.outdir + "lnprobs.hdf5"
#         self.logger.debug("Opening {} for writing HDF5 lnprobs".format(filename))
#         hdf5 = h5py.File(filename, "w")
#         lnprobs = self.lnprobability
#
#         dset = hdf5.create_dataset(self.fname, samples.shape[:1], compression='gzip', compression_opts=9)
#         dset[:] = lnprobs
#         dset.attrs["commit"] = "{}".format(C.get_git_commit())
#         hdf5.close()
#
#     def plot(self, triangle_plot=False):
#         '''
#         Generate the relevant plots once the sampling is done.
#         '''
#         samples = self.flatchain
#
#         plot_walkers(self.outdir + self.fname + "_chain_pos.png", samples, labels=self.param_tuple)
#
#         if triangle_plot:
#             import triangle
#             figure = triangle.corner(samples, labels=self.param_tuple, quantiles=[0.16, 0.5, 0.84],
#                                      show_titles=True, title_args={"fontsize": 12})
#             figure.savefig(self.outdir + self.fname + "_triangle.png")
#
#             plt.close(figure)
#
# class NuisanceSampler(SSampler):
#     def __init__(self, **kwargs):
#         '''
#
#         :param OrderModel: the parallel.OrderModel instance
#
#         :param starting_param_dict: the dictionary of starting parameters
#
#         :param cov:
#             the MH proposal
#
#         :param debug:
#
#         :param args: []
#
#         '''
#
#         starting_param_dict = kwargs.get("starting_param_dict")
#         self.param_tuple = self.startdict_to_tuple(starting_param_dict)
#         print("param_tuple is {}".format(self.param_tuple))
#         #print("param_tuple length {}".format(len(self.param_tuple)))
#
#         chebs = [starting_param_dict["cheb"][key] for key in self.cheb_tup]
#         covs = [starting_param_dict["cov"][key] for key in self.cov_tup]
#         regions = starting_param_dict["regions"]
#         #print("initializing {}".format(regions))
#         regs = [regions[id][kk] for id in sorted(regions) for kk in C.cov_region_parameters]
#         #print("regs {}".format(regs))
#
#         self.p0 = np.array(chebs + covs + regs)
#
#         kwargs.update({"p0":self.p0, "revertfn":self.revertfn, "lnprobfn":self.lnprob})
#         super(NuisanceSampler, self).__init__(**kwargs)
#
#         self.model = kwargs.get("OrderModel")
#         spectrum_id, order_id = self.model.id
#         order = kwargs.get("order", order_id)
#         #self.fname = "{}/{}/{}".format(spectrum_id, order, "nuisance")
#         self.fname = "nuisance"
#         self.params = None
#         self.prior_params = kwargs.get("prior_params", None)
#         if self.prior_params:
#             self.sigma0 = self.prior_params["regions"]["sigma0"]
#             self.mus = self.prior_params["regions"]["mus"]
#             self.mu_width = self.prior_params["regions"]["mu_width"]
#             self.sigma_knee = self.prior_params["regions"]["sigma_knee"]
#             self.frac_global = self.prior_params["regions"]["frac_global"]
#
#     def startdict_to_tuple(self, startdict):
#         #This is a little more tricky than the stellar parameters.
#         #How are the keys stored and passed in the dictionary?
#         #{"cheb": [c0, c1, c2, ..., cn], "cov": [sigAmp, logAmp, l],
#         #        "regions":{0: [logAmp, ], 1: [], N:[] }}
#
#         #Serialize the cheb parameters
#         self.ncheb = len(startdict["cheb"])
#         self.cheb_tup = ("logc0",) + tuple(["c{}".format(i) for i in range(1, self.ncheb)])
#
#         #Serialize the covariance parameters
#         self.ncov = 3
#         cov_tup = ()
#         for param in C.cov_global_parameters:
#             #check if param is in keys, if so, add to the tuple
#             if param in startdict["cov"]:
#                 cov_tup += (param,)
#         self.cov_tup = cov_tup
#
#         regions_tup = ()
#         self.regions = startdict.get("regions", None)
#         if self.regions:
#             self.nregions = len(self.regions)
#             for key in sorted(self.regions.keys()):
#                 for kk in C.cov_region_parameters:
#                     regions_tup += ("r{:0>2}-{}".format(key,kk),)
#             self.regions_tup = regions_tup
#         else:
#             self.nregions = 0
#             self.regions_tup = ()
#
#
#         tup = self.cheb_tup + self.cov_tup + self.regions_tup
#         #This should look like
#         #tup = ("c0", "c1", ..., "cn", "sigAmp", "logAmp", "l", "r00_logAmp", "r00_mu", "r00_sigma",
#         # "r01_logAmp", ..., "rNN_sigma")
#         return tup
#
#     def zip_p(self, p):
#         '''
#         Convert the vector to a dictionary
#         '''
#         cheb = dict(zip(self.cheb_tup, p[:self.ncheb]))
#         cov = dict(zip(self.cov_tup, p[self.ncheb:self.ncheb+self.ncov]))
#         regions = p[-self.nregions*3:]
#         rdict = {}
#         for i in range(self.nregions):
#             rdict[i] = dict(zip(("logAmp", "mu", "sigma"), regions[i*3:3*(i+1)]))
#
#         params = {"cheb":cheb, "cov":cov, "regions":rdict}
#         return params
#
#     def revertfn(self):
#         self.logger.debug("reverting model")
#         self.model.prior = self.prior_last
#         self.params = self.params_last
#         self.model.revert_nuisance()
#
#     def lnprob(self, p):
#         self.params_last = self.params
#         params = self.zip_p(p)
#         self.params = params
#         self.logger.debug("Updating nuisance params {}".format(params))
#
#         # Nuisance parameter priors implemented here
#         self.prior_last = self.model.prior
#         # Region parameter priors implemented here
#         if self.nregions > 0:
#             regions = params["regions"]
#             keys = sorted(regions)
#
#             #Unpack the region parameters into a vector of mus, amps, and sigmas
#             amps = 10**np.array([regions[key]["logAmp"] for key in keys])
#             cov_amp = 10**params["cov"]["logAmp"]
#
#             #First check to make sure that amplitude can't be some factor less than the global covariance
#             if np.any(amps < (cov_amp * self.frac_global)):
#                 return -np.inf
#
#             mus = np.array([regions[key]["mu"] for key in keys])
#             sigmas = np.array([regions[key]["sigma"] for key in keys])
#
#             #Make sure the region hasn't strayed too far from the original specification
#             if np.any(np.abs(mus - self.mus) > self.sigma0):
#                 # The region has strayed too far from the original specification
#                 return -np.inf
#
#             #Use a Gaussian prior on mu, that it keeps the region within the original setting.
#             # 1/(sqrt(2pi) * sigma) exp(-0.5 (mu-x)^2/sigma^2)
#             #-ln(sigma * sqrt(2 pi)) - 0.5 (mu - x)^2 / sigma^2
#             #width = 0.04
#             lnGauss = -0.5 * np.sum(np.abs(mus - self.mus)**2/self.mu_width**2 -
#                                     np.log(self.mu_width * np.sqrt(2. * np.pi)))
#
#             # Use a ln(logistic) function on sigma, that is flat before the knee and dies off for anything
#             # greater, to prevent dilution into global cov kernel
#             lnLogistic = np.sum(np.log(-1./(1. + np.exp(self.sigma_knee - sigmas)) + 1.))
#
#             self.model.prior = lnLogistic + lnGauss
#
#         try:
#             self.model.update_nuisance(params)
#             lnp = self.model.evaluate() # also sets OrderModel.lnprob to proposed value. Includes self.model.prior
#             return lnp
#         except C.ModelError:
#             return -np.inf


from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import sys

sys.path.remove("/opt/ros/kinetic/lib/python2.7/dist-packages")

import pickle
import json
import numpy as np
import cv2
from PIL import Image

DATA_PATH = '/media/andy/andy-2TB/2d_labeling/2d_detection_data/'
DEBUG = False

'''
# Values Name Description
1 type  'Car', 'Bus', 'Truck', 'Pedestrian', 'Bicycle', 'Traffic Cone', 'barricade', 'moto', 'other'

'''

def _center_to_coco_bbox(bbox):
    """
    from center format
    to (tlx, tly, w, h)
    :param bbox:
    :return:
    """
    tlx = bbox[0] - bbox[2] / 2
    tly = bbox[1] - bbox[3] / 2
    return [(tlx), (tly), (bbox[2]), (bbox[3])]

def _corner_to_coco_bbox(bbox):
    """
    from corner format
    to (tlx, tly, w, h)
    :param bbox:
    :return:
    """
    return [(bbox[0]), (bbox[1]), (bbox[2] - bbox[0]), (bbox[3] - bbox[1])]

def load_json_data(fpath):
    with open(fpath, "r") as f:
        json_data = json.load(f)
    return json_data

def dump_json(json_data, fpath):
    with open(fpath, "w") as f:
        json.dump(json_data, f)

cats = ['car', 'bus', 'truck', 'moto', 'pedestrian', 'bike', 'traffic_cone', 'barricade', 'other']
cat_ids = {cat: i + 1 for i, cat in enumerate(cats)}
cat_info = []
for i, cat in enumerate(cats):
    cat_info.append({'name': cat, 'id': i + 1})

ret = {"images": [], "annotations": [], "categories": cat_info}

labeling_path = DATA_PATH + "train/train.json"
# labeling_path = DATA_PATH + "test/test.json"
json_data = load_json_data(labeling_path)

image_to_id = {}
max_image_id = 0

for item in json_data:
    img_abs_path = DATA_PATH + "images/" + item['file']
    try:
        image = cv2.imread(img_abs_path)
        image_height, image_width, _ = image.shape
    except Exception as e:
        print(e)
        print("skip image, path ", img_abs_path)
        continue
    print(item['file'])
    image_id = image_to_id.get(item['file'], -1)
    if image_id == -1:
        image_id = max_image_id
        image_to_id.setdefault(item['file'], image_id)
        max_image_id += 1

    image_info = {'file_name': item['file'],
                  'id': image_id}

    ret['images'].append(image_info)

    for box_id, box in enumerate(item['boxes']):
        cat_id = cat_ids[box['class']]
        truncated = 1 if 'truncated_vehicle' in box and box['truncated_vehicle'] == 1 else 0
        if 'occluded' not in box or box['occluded'] is None:
            occluded = 0
        elif 'occluded' in box and box['occluded'] is 'partially':
            occluded = 1
        elif 'occluded' in box and box['occluded'] is 'mostly' or box['occluded'] is 'barely':
            occluded = 2
        bbox = [float(box['coord'][0]) * image_width, float(box['coord'][1]) * image_height,
                float(box['coord'][2]) * image_width, float(box['coord'][3]) * image_height]

        ann = {
            'image_id': image_id,
            'id': int(len(ret['annotations']) + 1),
            'category_id': cat_id,
            'bbox': _center_to_coco_bbox(bbox),
            'truncated': truncated,
            'occluded': occluded
        }
        ret['annotations'].append(ann)

print("# images: ", len(ret['images']))
print("# annotations: ", len(ret['annotations']))

out_path = "{}/annotations/plusai_{}.json".format(DATA_PATH, "train")
dump_json(ret, out_path)


#!/usr/bin/env python

"""
    nearest_cloud.py - Version 1.0 2013-07-28

    Compute the COG of the nearest object in x-y-z space and publish as a PoseStamped message.

    Relies on PCL ROS nodelets in the launch file to pre-filter the
    cloud on the x, y and z dimensions.

    Based on the follower application by Tony Pratkanis at:

    http://ros.org/wiki/turtlebot_follower

    Created for the Pi Robot Project: http://www.pirobot.org
    Copyright (c) 2013 Patrick Goebel.  All rights reserved.

    This program is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 2 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details at:

    http://www.gnu.org/licenses/gpl.html
"""

import rospy
from roslib import message
from sensor_msgs.msg import PointCloud2
from sensor_msgs import point_cloud2
from geometry_msgs.msg import Point, PoseStamped, Quaternion
from tf.transformations import quaternion_from_euler
import numpy as np
import cv2
from math import pi, radians

class NearestCloud():
    def __init__(self):
        rospy.init_node("nearest_cloud")

        self.min_points = rospy.get_param("~min_points", 25)
        self.z_percentile = rospy.get_param("~z_percentile", 100)

        # Define the target publisher
        self.target_pub = rospy.Publisher('target_pose', PoseStamped)

        rospy.Subscriber('point_cloud', PointCloud2, self.get_nearest_cloud)

        # Wait for the pointcloud topic to become available
        rospy.wait_for_message('point_cloud', PointCloud2)

    def get_nearest_cloud(self, msg):
        points = list()
        points_xy = list()

        # Get all the points in the visible cloud (may be prefiltered by other nodes)
        for point in point_cloud2.read_points(msg, skip_nans=True):
            points.append(point[:3])
            points_xy.append(point[:2])

        # Convert to a numpy array
        points_arr = np.float32([p for p in points]).reshape(-1, 1, 3)

        # Compute the COG
        cog = np.mean(points_arr, 0)

        # Convert to a Point
        cog_point = Point()
        cog_point.x = cog[0][0]
        cog_point.y = cog[0][1]
        cog_point.z = cog[0][2]
        #cog_point.z = 0.35

        # Abort if we get an NaN in any component
        if np.isnan(np.sum(cog)):
            return

        # If we have enough points, find the best fit ellipse around them
        try:
            if len(points_xy) > 6:
                points_xy_arr = np.float32([p for p in points_xy]).reshape(-1, 1, 2)
                track_box = cv2.fitEllipse(points_xy_arr)
            else:
                # Otherwise, find the best fitting rectangle
                track_box = cv2.boundingRect(points_xy_arr)

            angle = pi - radians(track_box[2])
        except:
            return

        #print angle

        # Convert the rotation angle to a quaternion
        q_angle = quaternion_from_euler(0, angle, 0, axes='sxyz')
        q = Quaternion(*q_angle)

        q.x = 0.707
        q.y = 0
        q.z = 0.707
        q.w = 0

        # Publish the COG and orientation
        target = PoseStamped()
        target.header.stamp = rospy.Time.now()
        target.header.frame_id = msg.header.frame_id
        target.pose.position = cog_point
        target.pose.orientation = q

        # Publish the movement command
        self.target_pub.publish(target)

if __name__ == '__main__':
    try:
        NearestCloud()
        rospy.spin()
    except rospy.ROSInterruptException:
        rospy.loginfo("Nearest cloud node terminated.")


import rospy
import smach
import moveit_lib.srv
import numpy as np
import baxter_interface
from std_msgs.msg import String
from std_msgs.msg import Float32, Int32
from geometry_msgs.msg import PoseStamped
from geometry_msgs.msg import PoseArray
from geometry_msgs.msg import Pose
from copy import deepcopy

import tf

class CheckHandPoseInShelf(smach.State):
    # send either a prefix to be added to (usually) the BIN name from the
    def __init__(self, relative_pose=None, pose_frame_id='',
                 movegroup=None, velocity_scale=1.0):
        smach.State.__init__(self, outcomes=['in_shelf', 'not_in_shelf', 'failed'])

        self.tf_listener = tf.TransformListener()

    # ==========================================================
    def execute(self, userdata):
        try:
            # painful tf lookup!
            # rospy.loginfo('Looking up transform from %s into %s frames....' % (frame_id, hand_frame))
            self.tf_listener.waitForTransform('/shelf', '/left_endpoint', rospy.Time(0), rospy.Duration(3.0))

            common_time = self.tf_listener.getLatestCommonTime('/shelf', '/left_endpoint')

            # this creates the pose of the hand, expressed in the goal frame
            hand_in_shelf_frame = self.tf_listener.lookupTransform('/shelf', '/left_endpoint', common_time)

            print 'hand_in_shelf_frame = '
            print hand_in_shelf_frame

            if hand_in_shelf_frame[0][2] < 0:
                rospy.loginfo("We are still in the shelf!")
                return 'in_shelf'
            else:
                rospy.loginfo("All good, we are out of the shelf!")
                return 'not_in_shelf'

        except Exception, e:
            rospy.logerr('Cannot lookup a transform from %s to %s now, but required to perform relative motion of the %s move group.'  % (hand_frame, frame_id, self.movegroup_))
            print e
            return 'failed'

        return 'failed'


from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import tensorflow.keras
from tensorflow.keras.preprocessing.image import Iterator
import numpy as np
import re
import os
import threading
import warnings
import multiprocessing.pool
from functools import partial

from tensorflow.keras import backend as K

class MultiLabelDirectoryIterator(Iterator):
    """Iterator capable of reading images from a directory on disk.

    # Arguments
        directory: Path to the directory to read images from.
            Each subdirectory in this directory will be
            considered to contain images from one class,
            or alternatively you could specify class subdirectories
            via the `classes` argument.
        image_data_generator: Instance of `ImageDataGenerator`
            to use for random transformations and normalization.
        target_size: tuple of integers, dimensions to resize input images to.
        color_mode: One of `"rgb"`, `"grayscale"`. Color mode to read images.
        classes: Optional list of strings, names of subdirectories
            containing images from each class (e.g. `["dogs", "cats"]`).
            It will be computed automatically if not set.
        class_mode: Mode for yielding the targets:
            `"binary"`: binary targets (if there are only two classes),
            `"categorical"`: categorical targets,
            `"sparse"`: integer targets,
            `"input"`: targets are images identical to input images (mainly
                used to work with autoencoders),
            `None`: no targets get yielded (only input images are yielded).
        batch_size: Integer, size of a batch.
        shuffle: Boolean, whether to shuffle the data between epochs.
        seed: Random seed for data shuffling.
        data_format: String, one of `channels_first`, `channels_last`.
        save_to_dir: Optional directory where to save the pictures
            being yielded, in a viewable format. This is useful
            for visualizing the random transformations being
            applied, for debugging purposes.
        save_prefix: String prefix to use for saving sample
            images (if `save_to_dir` is set).
        save_format: Format to use for saving sample images
            (if `save_to_dir` is set).
        interpolation: Interpolation method used to resample the image if the
            target size is different from that of the loaded image.
            Supported methods are "nearest", "bilinear", and "bicubic".
            If PIL version 1.1.3 or newer is installed, "lanczos" is also
            supported. If PIL version 3.4.0 or newer is installed, "box" and
            "hamming" are also supported. By default, "nearest" is used.
    """

    def __init__(self, x, y, image_data_generator,
                 target_size=(256, 256), color_mode='rgb',
                 class_mode='categorical',
                 batch_size=32, shuffle=True, seed=None,
                 data_format=None, save_to_dir=None,
                 save_prefix='', save_format='png',
                 follow_links=False, interpolation='nearest'):
        if data_format is None:
            data_format = K.image_data_format()

        self.x = x

        if y is not None:
            self.y = np.asarray(y)
        else:
            self.y = None

        self.image_data_generator = image_data_generator
        self.target_size = tuple(target_size)
        if color_mode not in {'rgb', 'grayscale'}:
            raise ValueError('Invalid color mode:', color_mode,
                             '; expected "rgb" or "grayscale".')
        self.color_mode = color_mode
        self.data_format = data_format
        if self.color_mode == 'rgb':
            if self.data_format == 'channels_last':
                self.image_shape = self.target_size + (3,)
            else:
                self.image_shape = (3,) + self.target_size
        else:
            if self.data_format == 'channels_last':
                self.image_shape = self.target_size + (1,)
            else:
                self.image_shape = (1,) + self.target_size

        if class_mode not in {'categorical', 'binary', 'sparse',
                              'input', None}:
            raise ValueError('Invalid class_mode:', class_mode,
                             '; expected one of "categorical", '
                             '"binary", "sparse", "input"'
                             ' or None.')
        self.class_mode = class_mode
        self.save_to_dir = save_to_dir
        self.save_prefix = save_prefix
        self.save_format = save_format
        self.interpolation = interpolation

        white_list_formats = {'png', 'jpg', 'jpeg', 'bmp', 'ppm'}

        # first, count the number of samples and classes
        self.samples = 0

        # second, build an index of the images in the different class subfolders
        results = []

        super(MultiLabelDirectoryIterator, self).__init__(len(x), batch_size, shuffle, seed)

    def _get_batches_of_transformed_samples(self, index_array):
        batch_x = np.zeros((len(index_array),) + self.image_shape, dtype=K.floatx())
        batch_x2 = np.zeros((len(index_array), 3) , dtype=K.floatx())
        #batch_x3 = np.zeros((len(index_array), len(self.x3[0])), dtype=K.floatx())

        grayscale = self.color_mode == 'grayscale'
        # build batch of image data
        for i, j in enumerate(index_array):
            fname = self.x[j]
            img = tensorflow.keras.preprocessing.image.load_img(fname,
                           grayscale=grayscale,
                           target_size=self.target_size,
                           interpolation=self.interpolation)
            x = tensorflow.keras.preprocessing.image.img_to_array(img, data_format=self.data_format)
            x = self.image_data_generator.random_transform(x)
            x = self.image_data_generator.standardize(x)
            batch_x[i] = x

            '''            arr = self.x2[j]
            lat = arr[0]
            lon = arr[1]
            x = np.cos(lat) * np.cos(lon)
            y = np.cos(lat) * np.sin(lon)
            z = np.sin(lat)
            '''

            #print(x, y, z)
            #arr = [arr[0]/200.0, arr[1]/200.0]
            #print(arr)
            #batch_x2[i] = np.array([x, y, z])
            #batch_x3[i] = self.x3[i]

        # optionally save augmented images to disk for debugging purposes
        if self.save_to_dir:
            for i, j in enumerate(index_array):
                img = tensorflow.keras.preprocessing.image.array_to_img(batch_x[i], self.data_format, scale=True)
                fname = '{prefix}_{index}_{hash}.{format}'.format(prefix=self.save_prefix,
                                                                  index=j,
                                                                  hash=np.random.randint(1e7),
                                                                  format=self.save_format)
                img.save(os.path.join(self.save_to_dir, fname))

        if self.y is None:
            return batch_x

        batch_y = self.y[index_array]
        #return [batch_x, batch_x2], batch_y
        return batch_x, batch_y
        #return [batch_x, batch_x3], batch_y

    def next(self):
        """For python 2.x.

        # Returns
            The next batch.
        """
        with self.lock:
            index_array = next(self.index_generator)
        # The transformation of images is not under thread lock
        # so it can be done in parallel
        return self._get_batches_of_transformed_samples(index_array)


import gym
from railrl.data_management.awr_env_replay_buffer import AWREnvReplayBuffer
from railrl.data_management.env_replay_buffer import EnvReplayBuffer
from railrl.envs.wrappers import NormalizedBoxEnv, StackObservationEnv, RewardWrapperEnv
import railrl.torch.pytorch_util as ptu
from railrl.samplers.data_collector import MdpPathCollector
from railrl.samplers.data_collector.step_collector import MdpStepCollector
from railrl.torch.networks import FlattenMlp
from railrl.torch.sac.policies import TanhGaussianPolicy, MakeDeterministic
from railrl.torch.sac.awr_sac import AWRSACTrainer
from railrl.torch.torch_rl_algorithm import (
    TorchBatchRLAlgorithm,
    TorchOnlineRLAlgorithm,
)

from railrl.demos.source.mdp_path_loader import MDPPathLoader
from railrl.visualization.video import save_paths

from multiworld.core.flat_goal_env import FlatGoalEnv
from multiworld.core.image_env import ImageEnv

from railrl.launchers.experiments.ashvin.rfeatures.encoder_wrapped_env import EncoderWrappedEnv
from railrl.launchers.experiments.ashvin.rfeatures.rfeatures_model import TimestepPredictionModel

import torch
import numpy as np
from torchvision.utils import save_image

from railrl.exploration_strategies.base import \
    PolicyWrappedWithExplorationStrategy
from railrl.exploration_strategies.gaussian_and_epislon import GaussianAndEpislonStrategy
from railrl.exploration_strategies.ou_strategy import OUStrategy

import os.path as osp
from railrl.core import logger
from railrl.misc.asset_loader import load_local_or_remote_file

ENV_PARAMS = {
    'half-cheetah': {  # 6 DoF
        'num_expl_steps_per_train_loop': 1000,
        'max_path_length': 1000,
        'env_id':'HalfCheetah-v2'
    },
    'hopper': {  # 6 DoF
        'num_expl_steps_per_train_loop': 1000,
        'max_path_length': 1000,
        'env_id':'Hopper-v2'
    },
    'humanoid': {  # 6 DoF
        'num_expl_steps_per_train_loop': 1000,
        'max_path_length': 1000,
        'env_id':'Humanoid-v2'
    },
    'inv-double-pendulum': {  # 2 DoF
        'num_expl_steps_per_train_loop': 1000,
        'max_path_length': 1000,
        'env_id':'InvertedDoublePendulum-v2'
    },
    'pendulum': {  # 2 DoF
        'num_expl_steps_per_train_loop': 200,
        'max_path_length': 200,
        'min_num_steps_before_training': 2000,
        'target_update_period': 200,
        'env_id':'Pendulum-v2'
    },
    'ant': {  # 6 DoF
        'num_expl_steps_per_train_loop': 1000,
        'max_path_length': 1000,
        'env_id':'Ant-v2'
    },
    'walker': {  # 6 DoF
        'num_expl_steps_per_train_loop': 1000,
        'max_path_length': 1000,
        'env_id':'Walker2d-v2'
    },
    'swimmer': {  # 6 DoF
        'num_expl_steps_per_train_loop': 1000,
        'max_path_length': 1000,
        'env_id':'Swimmer-v2'
    },

    'pen-v0': {
        'env_id': 'pen-v0',
        # 'num_expl_steps_per_train_loop': 1000,
        'max_path_length': 200,
        # 'num_epochs': 1000,
    },
    'door-v0': {
        'env_id': 'door-v0',
        # 'num_expl_steps_per_train_loop': 1000,
        'max_path_length': 200,
        # 'num_epochs': 1000,
    },
    'relocate-v0': {
        'env_id': 'relocate-v0',
        # 'num_expl_steps_per_train_loop': 1000,
        'max_path_length': 200,
        # 'num_epochs': 1000,
    },
    'hammer-v0': {
        'env_id': 'hammer-v0',
        # 'num_expl_steps_per_train_loop': 1000,
        'max_path_length': 200,
        # 'num_epochs': 1000,
    },

    'pen-sparse-v0': {
        'env_id': 'pen-v0',
        'max_path_length': 200,
        'sparse_reward': True,
        'env_demo_path': dict(
            path="demos/icml2020/hand/pen2_sparse.npy",
            obs_dict=True,
            is_demo=True,
        ),
        'env_offpolicy_data_path': dict(
            # path="demos/icml2020/hand/pen_bc_sparse1.npy",
            # path="demos/icml2020/hand/pen_bc_sparse2.npy",
            # path="demos/icml2020/hand/pen_bc_sparse3.npy",
            path="demos/icml2020/hand/pen_bc_sparse4.npy",
            obs_dict=False,
            is_demo=False,
            train_split=0.9,
        ),
    },
    'door-sparse-v0': {
        'env_id': 'door-v0',
        'max_path_length': 200,
        'sparse_reward': True,
        'env_demo_path': dict(
            path="demos/icml2020/hand/door2_sparse.npy",
            obs_dict=True,
            is_demo=True,
        ),
        'env_offpolicy_data_path': dict(
            # path="demos/icml2020/hand/door_bc_sparse1.npy",
            # path="demos/icml2020/hand/door_bc_sparse3.npy",
            path="demos/icml2020/hand/door_bc_sparse4.npy",
            obs_dict=False,
            is_demo=False,
            train_split=0.9,
        ),
    },
    'relocate-sparse-v0': {
        'env_id': 'relocate-v0',
        'max_path_length': 200,
        'sparse_reward': True,
        'env_demo_path': dict(
            path="demos/icml2020/hand/relocate2_sparse.npy",
            obs_dict=True,
            is_demo=True,
        ),
        'env_offpolicy_data_path': dict(
            # path="demos/icml2020/hand/relocate_bc_sparse1.npy",
            path="demos/icml2020/hand/relocate_bc_sparse4.npy",
            obs_dict=False,
            is_demo=False,
            train_split=0.9,
        ),
    },
    'hammer-sparse-v0': {
        'env_id': 'hammer-v0',
        'max_path_length': 200,
        'sparse_reward': True,
        'env_demo_path': dict(
            path="demos/icml2020/hand/hammer2_sparse.npy",
            obs_dict=True,
            is_demo=True,
        ),
        'env_offpolicy_data_path': dict(
            path="demos/icml2020/hand/hammer_bc_sparse1.npy",
            obs_dict=False,
            is_demo=False,
            train_split=0.9,
        ),
    },
}

def compute_hand_sparse_reward(next_obs, reward, done, info):
    return info['goal_achieved'] - 1

def encoder_wrapped_env(variant):
    representation_size = 128
    output_classes = 20

    model_class = variant.get('model_class', TimestepPredictionModel)
    model = model_class(
        representation_size,
        # decoder_output_activation=decoder_activation,
        output_classes=output_classes,
        **variant['model_kwargs'],
    )
    # model = torch.nn.DataParallel(model)

    model_path = variant.get("model_path")
    # model = load_local_or_remote_file(model_path)
    state_dict = torch.load(model_path)
    model.load_state_dict(state_dict)
    model.to(ptu.device)
    model.eval()

    traj = np.load(variant.get("desired_trajectory"), allow_pickle=True)[0]

    goal_image = traj["observations"][-1]["image_observation"]
    goal_image = goal_image.reshape(1, 3, 500, 300).transpose([0, 1, 3, 2]) / 255.0
    # goal_image = goal_image.reshape(1, 300, 500, 3).transpose([0, 3, 1, 2]) / 255.0 # BECAUSE RLBENCH DEMOS ARENT IMAGE_ENV WRAPPED
    # goal_image = goal_image[:, :, :240, 60:500]
    goal_image = goal_image[:, :, 60:, 60:500]
    goal_image_pt = ptu.from_numpy(goal_image)
    save_image(goal_image_pt.data.cpu(), 'gitignore/goal.png', nrow=1)
    goal_latent = model.encode(goal_image_pt).detach().cpu().numpy().flatten()

    initial_image = traj["observations"][0]["image_observation"]
    initial_image = initial_image.reshape(1, 3, 500, 300).transpose([0, 1, 3, 2]) / 255.0
    # initial_image = initial_image.reshape(1, 300, 500, 3).transpose([0, 3, 1, 2]) / 255.0
    # initial_image = initial_image[:, :, :240, 60:500]
    initial_image = initial_image[:, :, 60:, 60:500]
    initial_image_pt = ptu.from_numpy(initial_image)
    save_image(initial_image_pt.data.cpu(), 'gitignore/initial.png', nrow=1)
    initial_latent = model.encode(initial_image_pt).detach().cpu().numpy().flatten()

    # Move these to td3_bc and bc_v3 (or at least type for reward_params)
    reward_params = dict(
        goal_latent=goal_latent,
        initial_latent=initial_latent,
        type=variant["reward_params_type"],
    )

    config_params = variant.get("config_params")

    env = variant['env_class'](**variant['env_kwargs'])
    env = ImageEnv(env,
        recompute_reward=False,
        transpose=True,
        image_length=450000,
        reward_type="image_distance",
        # init_camera=sawyer_pusher_camera_upright_v2,
    )
    env = EncoderWrappedEnv(
        env,
        model,
        reward_params,
        config_params,
        **variant.get("encoder_wrapped_env_kwargs", dict())
    )
    env = FlatGoalEnv(env, obs_keys=["state_observation", ])

    return env

def resume(variant):
    data = load_local_or_remote_file(variant.get("pretrained_algorithm_path"), map_location="cuda")
    algo = data['algorithm']

    algo.num_epochs = variant['num_epochs']

    post_pretrain_hyperparams = variant["trainer_kwargs"].get("post_pretrain_hyperparams", {})
    algo.trainer.set_algorithm_weights(**post_pretrain_hyperparams)

    algo.train()

def experiment(variant):
    if variant.get("pretrained_algorithm_path", False):
        resume(variant)
        return

    if 'env' in variant:
        env_params = ENV_PARAMS[variant['env']]
        variant.update(env_params)

        if 'env_id' in env_params:
            if env_params['env_id'] in ['pen-v0', 'pen-sparse-v0', 'door-v0', 'relocate-v0', 'hammer-v0',
                                        'pen-sparse-v0', 'door-sparse-v0', 'relocate-sparse-v0', 'hammer-sparse-v0']:
                import mj_envs
            expl_env = gym.make(env_params['env_id'])
            eval_env = gym.make(env_params['env_id'])
        else:
            expl_env = NormalizedBoxEnv(variant['env_class']())
            eval_env = NormalizedBoxEnv(variant['env_class']())

        if variant.get('sparse_reward', False):
            expl_env = RewardWrapperEnv(expl_env, compute_hand_sparse_reward)
            eval_env = RewardWrapperEnv(eval_env, compute_hand_sparse_reward)

        if variant.get('add_env_demos', False):
            variant["path_loader_kwargs"]["demo_paths"].append(variant["env_demo_path"])

        if variant.get('add_env_offpolicy_data', False):
            variant["path_loader_kwargs"]["demo_paths"].append(variant["env_offpolicy_data_path"])
    else:
        expl_env = encoder_wrapped_env(variant)
        eval_env = encoder_wrapped_env(variant)

    path_loader_kwargs = variant.get("path_loader_kwargs", {})
    stack_obs = path_loader_kwargs.get("stack_obs", 1)
    if stack_obs > 1:
        expl_env = StackObservationEnv(expl_env, stack_obs=stack_obs)
        eval_env = StackObservationEnv(eval_env, stack_obs=stack_obs)

    obs_dim = expl_env.observation_space.low.size
    action_dim = eval_env.action_space.low.size
    if hasattr(expl_env, 'info_sizes'):
        env_info_sizes = expl_env.info_sizes
    else:
        env_info_sizes = dict()

    M = variant['layer_size']
    qf1 = FlattenMlp(
        input_size=obs_dim + action_dim,
        output_size=1,
        hidden_sizes=[M, M],
    )
    qf2 = FlattenMlp(
        input_size=obs_dim + action_dim,
        output_size=1,
        hidden_sizes=[M, M],
    )
    target_qf1 = FlattenMlp(
        input_size=obs_dim + action_dim,
        output_size=1,
        hidden_sizes=[M, M],
    )
    target_qf2 = FlattenMlp(
        input_size=obs_dim + action_dim,
        output_size=1,
        hidden_sizes=[M, M],
    )
    policy_class = variant.get("policy_class", TanhGaussianPolicy)
    policy = policy_class(
        obs_dim=obs_dim,
        action_dim=action_dim,
        **variant['policy_kwargs'],
    )

    buffer_policy = policy_class(
        obs_dim=obs_dim,
        action_dim=action_dim,
        **variant['policy_kwargs'],
    )

    eval_policy = MakeDeterministic(policy)
    eval_path_collector = MdpPathCollector(
        eval_env,
        eval_policy,
    )

    expl_policy = policy
    exploration_kwargs =  variant.get('exploration_kwargs', {})
    if exploration_kwargs:
        if exploration_kwargs.get("deterministic_exploration", False):
            expl_policy = MakeDeterministic(policy)

        exploration_strategy = exploration_kwargs.get("strategy", None)
        if exploration_strategy is None:
            pass
        elif exploration_strategy == 'ou':
            es = OUStrategy(
                action_space=expl_env.action_space,
                max_sigma=exploration_kwargs['noise'],
                min_sigma=exploration_kwargs['noise'],
            )
            expl_policy = PolicyWrappedWithExplorationStrategy(
                exploration_strategy=es,
                policy=expl_policy,
            )
        elif exploration_strategy == 'gauss_eps':
            es = GaussianAndEpislonStrategy(
                action_space=expl_env.action_space,
                max_sigma=exploration_kwargs['noise'],
                min_sigma=exploration_kwargs['noise'],  # constant sigma
                epsilon=0,
            )
            expl_policy = PolicyWrappedWithExplorationStrategy(
                exploration_strategy=es,
                policy=expl_policy,
            )
        else:
            error

    if variant.get('replay_buffer_class', EnvReplayBuffer) == AWREnvReplayBuffer:
        main_replay_buffer_kwargs = variant['replay_buffer_kwargs']
        main_replay_buffer_kwargs['env'] = expl_env
        main_replay_buffer_kwargs['qf1'] = qf1
        main_replay_buffer_kwargs['qf2'] = qf2
        main_replay_buffer_kwargs['policy'] = policy
    else:
        main_replay_buffer_kwargs=dict(
            max_replay_buffer_size=variant['replay_buffer_size'],
            env=expl_env,
        )
    replay_buffer_kwargs = dict(
        max_replay_buffer_size=variant['replay_buffer_size'],
        env=expl_env,
    )

    replay_buffer = variant.get('replay_buffer_class', EnvReplayBuffer)(
        **main_replay_buffer_kwargs,
    )
    trainer = AWRSACTrainer(
        env=eval_env,
        policy=policy,
        qf1=qf1,
        qf2=qf2,
        target_qf1=target_qf1,
        target_qf2=target_qf2,
        buffer_policy=buffer_policy,
        **variant['trainer_kwargs']
    )
    if variant['collection_mode'] == 'online':
        expl_path_collector = MdpStepCollector(
            expl_env,
            policy,
        )
        algorithm = TorchOnlineRLAlgorithm(
            trainer=trainer,
            exploration_env=expl_env,
            evaluation_env=eval_env,
            exploration_data_collector=expl_path_collector,
            evaluation_data_collector=eval_path_collector,
            replay_buffer=replay_buffer,
            max_path_length=variant['max_path_length'],
            batch_size=variant['batch_size'],
            num_epochs=variant['num_epochs'],
            num_eval_steps_per_epoch=variant['num_eval_steps_per_epoch'],
            num_expl_steps_per_train_loop=variant['num_expl_steps_per_train_loop'],
            num_trains_per_train_loop=variant['num_trains_per_train_loop'],
            min_num_steps_before_training=variant['min_num_steps_before_training'],
        )
    else:
        expl_path_collector = MdpPathCollector(
            expl_env,
            expl_policy,
        )
        algorithm = TorchBatchRLAlgorithm(
            trainer=trainer,
            exploration_env=expl_env,
            evaluation_env=eval_env,
            exploration_data_collector=expl_path_collector,
            evaluation_data_collector=eval_path_collector,
            replay_buffer=replay_buffer,
            max_path_length=variant['max_path_length'],
            batch_size=variant['batch_size'],
            num_epochs=variant['num_epochs'],
            num_eval_steps_per_epoch=variant['num_eval_steps_per_epoch'],
            num_expl_steps_per_train_loop=variant['num_expl_steps_per_train_loop'],
            num_trains_per_train_loop=variant['num_trains_per_train_loop'],
            min_num_steps_before_training=variant['min_num_steps_before_training'],
        )
    algorithm.to(ptu.device)

    demo_train_buffer = EnvReplayBuffer(
        **replay_buffer_kwargs,
    )
    demo_test_buffer = EnvReplayBuffer(
        **replay_buffer_kwargs,
    )

    if variant.get('save_paths', False):
        algorithm.post_train_funcs.append(save_paths)
    if variant.get('load_demos', False):
        path_loader_class = variant.get('path_loader_class', MDPPathLoader)
        path_loader = path_loader_class(trainer,
            replay_buffer=replay_buffer,
            demo_train_buffer=demo_train_buffer,
            demo_test_buffer=demo_test_buffer,
            **path_loader_kwargs
        )
        path_loader.load_demos()
    if variant.get('pretrain_policy', False):
        trainer.pretrain_policy_with_bc()
    if variant.get('pretrain_rl', False):
        trainer.pretrain_q_with_bc_data()
    if variant.get('save_pretrained_algorithm', False):
        p_path = osp.join(logger.get_snapshot_dir(), 'pretrain_algorithm.p')
        pt_path = osp.join(logger.get_snapshot_dir(), 'pretrain_algorithm.pt')
        data = algorithm._get_snapshot()
        data['algorithm'] = algorithm
        torch.save(data, open(pt_path, "wb"))
        torch.save(data, open(p_path, "wb"))
    if variant.get('train_rl', True):
        algorithm.train()


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# https://qualitysphere.github.io/images/point.png
# https://qualitysphere.gitee.io/images/point.png

from pony.orm import db_session, get, select
from models.models import *
import numpy

def __reload_issue_policy_data(policy_body: dict):
    """
    重新加载缺陷规则的数据
    :param policy_body:
    :return:
    """
    __policy_info = dict()
    for level in ['project', 'sprint']:
        __policy_info[level] = dict()
        for item in policy_body[level].items():
            __policy_info[level][item['key']] = {
                'weight': item['weight'],
                's': item['s'],
                'a': item['a'],
                'b': item['b'],
                'c': item['c'],
                'd': item['d'],
            }
    return __policy_info

@db_session
def add_grade_policy(body: dict):
    item = GradePolicy(
        name=body['name'],
        policy_conf={
            'project': body['project'],
            'sprint': body['sprint'],
        }
    )
    return {
        'id': item.uuid,
        'name': item.name,
    }

def update_grade_policy(policy_id: str):
    item = get(p for p in GradePolicy if str(p.uuid) == policy_id)
    item.name = ''
    item.policy_conf = ''
    item.update_time = ''
    return True

def get_project_grade_report():
    pass

@db_session
def get_sprint_grade_report(sprint_id: str):
    pass

@db_session
def generate_grade_report(policy_id: str, sprint_id: str):
    """初始化 report"""
    report = {
        'project': {
            'bug_count_risk': 0,
            'prod_bug_risk': 0,
            'bug_found_rate': 0,
            'bug_density_risk': 0,
        },
        'sprint': {
            'bug_severity_rate': 0,
            'reopen_bug_rate': 0,
            'regression_bug_rate': 0,
            'previous_bug_rate': 0,
            'feature_found_bug_rate': 0,
            'improve_found_bug_rate': 0,
            'qamissed_bug_rate': 0,
            'rc_bug_risk': 0,
            'bug_growth_risk': 0,
            'bug_fix_rate': 0,
            'bug_fix_risk': 0,
            'bug_verify_rate': 0,
            'bug_verify_risk': 0,
            'feature_bug_risk': 0,
            # 'case_execution_rate': 0,
            # 'case_pass_rate': 0,
        },
    }
    """获取 issue 数据并进行指标计算 https://qualitysphere.gitee.io/images/point.png"""
    sprint = get(s for s in Sprint if str(s.uuid) == sprint_id)
    project = get(p for p in Project if p.uuid == sprint.project.uuid)
    # bugs in the Sprint ➗ average(bugs in all Sprints)
    bugs_in_the_sprint = get(i for i in IssueCaptureStaticProject if i.sprint.uuid == sprint.uuid).in_release['total']
    average_bugs_in_all_sprints = list()
    for item in select(i for i in IssueCaptureStaticProject if i.sprint.project.uuid == project.uuid):
        average_bugs_in_all_sprints.append(item.in_release['total'])
    report['project']['bug_count_risk'] = numpy.divide(
        bugs_in_the_sprint,
        numpy.mean(average_bugs_in_all_sprints)
    )
    # customer found bugs in the Sprint ➗ customer found bugs in the previous Sprint
    # team found bugs ➗ ( team found bugs + customer found bugs ) * 100%
    # ( the Sprint total bugs ➗ the Sprint new code lines ) ➗ ( previous Sprint total bugs ➗ previous Sprint new code lines)

    """创建 Sprint 评级报告"""
    item = GradeReportSprint(
        policy=get(p for p in GradePolicy if str(p.uuid) == policy_id),
        sprint=get(s for s in Sprint if str(s.uuid) == sprint_id),
        report=report
    )

if __name__ == '__main__':
    print(u'This is SERVICE for Grade')


import src.config as cfg
import src.feature_subsets as subset
from sklearn.ensemble import GradientBoostingClassifier
from sklearn.model_selection import cross_val_score
from sklearn.metrics import classification_report, confusion_matrix, roc_auc_score, f1_score, precision_score, \
    recall_score, accuracy_score
from sklearn.model_selection import train_test_split
from datetime import datetime
from sklearn.ensemble import RandomForestClassifier
from src.multiscorer import MultiScorer
from numpy import average
from tqdm import tqdm
import os
import csv
from train_model import random_forest_classifyer_default,random_forest_classifyer_tuned,gradient_boost_classifier_default,gradient_boost_classifier_tuned

def main():
    connection_rw_size = cfg.CONNECTION_RW_SIZE
    timw_rw_size_min = cfg.TIME_RW_SIZE

    features_csv = f'../../data/processed/full_ft_netflow_crw_{5000}_trw_{5}.csv'

    label = subset.get_target(features_csv)

    baseline_ft, basene_ft_name = subset.get_baseline_ft(features_csv)
    time_ft, time_ft_name = subset.get_time_ft(features_csv)
    conn_ft, conn_ft_name = subset.get_conn_ft(features_csv)
    comb_ft, comb_ft_name = subset.get_comb_ft(features_csv)

    time_fwd_ft, time_fwd_ft_name = subset.get_time_fwd_ft(features_csv)
    time_bwd_ft, time_bwd_ft_name = subset.get_time_bwd_ft(features_csv)

    conn_fwd_ft, conn_fwd_ft_name = subset.get_conn_fwd_ft(features_csv)
    conn_bwd_ft, conn_bwd_ft_name = subset.get_conn_bwd_ft(features_csv)

    comb_fwd_ft, comb_fwd_ft_name = subset.get_comb_fwd_ft(features_csv)
    comb_bwd_ft, comb_bwd_ft_name = subset.get_comb_bwd_ft(features_csv)

    feature_subsets= [ [baseline_ft, 'baseline'],
                       [time_ft,'time'],
                       [conn_ft,'conn'],
                       [comb_ft,'comb'],
                       [time_fwd_ft,'time_fwd'],
                       [time_bwd_ft,'time_bwd'],
                       [conn_fwd_ft,'conn_fwd'],
                       [conn_bwd_ft,'conn_bwd'],
                       [comb_fwd_ft,'comb_fwd'],
                       [comb_bwd_ft,'comb_bwd']]

    rw_size_set =[[1000,1],
             [2000,2],
             [5000,5],
             [10000,10],
             [15000,15],
             [20000,20]]

    metric_list_list=[]
    # for ft_set in feature_subsets:
    #     print(f' ////////////// {ft_set[1]}  ////////////// ')
    #     print('Starting default model')
    #
    #     model_info_default, time_info_default = random_forest_classifyer_default(ft_set[0], label,
    #                                                                              random_state=42)
    #     print('Starting tuned model')
    #     model_info_tuned, time_info_tuned = random_forest_classifyer_tuned(ft_set[0], label, random_state=42)
    #     print('Calculating default metrics')
    #     metric_list_list.append(get_model_general_results(model_info_default[0], model_info_default[1], model_info_default[3], f'RF - DF - {ft_set[1]} - RW {connection_rw_size}/{timw_rw_size_min}'))
    #     print('Calculating tuned metrics')
    #     metric_list_list.append(
    #         get_model_general_results(model_info_tuned[0], model_info_tuned[1], model_info_tuned[3],
    #                                   f'RF - Tuned - {ft_set[1]} - RW {connection_rw_size}/{timw_rw_size_min}'))
    #     file = open(os.getcwd() + 'metric_list.csv', 'w', newline='', encoding='utf-8')
    #     print('\n')
    #     print(metric_list_list)
    #     print('\n')
    # for ft_set in feature_subsets:
    #     print(f' ////////////// {ft_set[1]}  ////////////// ')
    #     print('Starting default model')
    #
    #     model_info_default, time_info_default = gradient_boost_classifier_default(ft_set[0], label,
    #                                                                              random_state=42)
    #     print('Starting tuned model')
    #     model_info_tuned, time_info_tuned = gradient_boost_classifier_tuned(ft_set[0], label, random_state=42)
    #     print('Calculating default metrics')
    #     metric_list_list.append(get_model_general_results(model_info_default[0], model_info_default[1], model_info_default[3], f'GB - DF - {ft_set[1]} - RW {connection_rw_size}/{timw_rw_size_min}'))
    #     print('Calculating tuned metrics')
    #     metric_list_list.append(
    #         get_model_general_results(model_info_tuned[0], model_info_tuned[1], model_info_tuned[3],
    #                                   f'GB - Tuned - {ft_set[1]} - RW {connection_rw_size}/{timw_rw_size_min}'))
    #     file = open(os.getcwd() + 'metric_list.csv', 'w', newline='', encoding='utf-8')
    #     print('\n')
    #     print(metric_list_list)
    #     print('\n')

    # for rw_size in rw_size_set:
    #     print(f' ////////////// {rw_size}  ////////////// ')
    #
    #     features_csv = f'../../data/processed/full_ft_netflow_crw_{rw_size[0]}_trw_{rw_size[1]}.csv'
    #     label = subset.get_target(features_csv)
    #     conn_ft, conn_ft_name = subset.get_conn_ft(features_csv)
    #     print('Trainning model')
    #     model_info_tuned, time_info_tuned = random_forest_classifyer_tuned(conn_ft, label, random_state=42)
    #     print('Generating metrics')
    #     metric_list_list.append(
    #                 get_model_general_results(model_info_tuned[0], model_info_tuned[1], model_info_tuned[3],
    #                                           f'RF - Tuned - CONN - RW {rw_size[0]}/{rw_size[1]}'))
    # for rw_size in rw_size_set:
    #     print(f' ////////////// {rw_size}  ////////////// ')
    #
    #     features_csv = f'../../data/processed/full_ft_netflow_crw_{rw_size[0]}_trw_{rw_size[1]}.csv'
    #     label = subset.get_target(features_csv)
    #     conn_ft, conn_ft_name = subset.get_conn_ft(features_csv)
    #     print('Trainning model')
    #     model_info_tuned, time_info_tuned = gradient_boost_classifier_tuned(conn_ft, label, random_state=42)
    #     print('Generating metrics')
    #     metric_list_list.append(
    #                 get_model_general_results(model_info_tuned[0], model_info_tuned[1], model_info_tuned[3],
    #                                           f'GB - Tuned - CONN - RW {rw_size[0]}/{rw_size[1]}'))
    #
    #
    # file = open(f'{os.getcwd()}/metric_list.csv', 'a', newline='', encoding='utf-8')
    # with file:
    #     writer=csv.writer(file)
    #     for rf_entry in metric_list_list:
    #         writer.writerow(rf_entry)
    print(f'GB Tuned || {datetime.now()}')
    model_info, time_info_tuned = gradient_boost_classifier_tuned(conn_fwd_ft, label, random_state=42)
    gb_tuned = model_info[2]
    print(f'GB Def || {datetime.now()}')
    model_info, time_info_tuned = gradient_boost_classifier_default(conn_fwd_ft, label, random_state=42)
    gb_def = model_info[2]
    print(f'RF Tuned || {datetime.now()}')
    model_info, time_info_tuned = random_forest_classifyer_tuned(conn_fwd_ft, label, random_state=42)
    rf_tuned = model_info[2]
    print(f'RF Def || {datetime.now()}')
    model_info, time_info_tuned = random_forest_classifyer_default(conn_fwd_ft, label, random_state=42)
    rf_def = model_info[2]

    time_gb_tuned=[]
    time_gb_def = []
    time_rf_tuned = []
    time_rf_def = []
    for x in range(0, 10):
        print(f'Iteration: {x} || {datetime.now()}')
        time_start=datetime.now()
        gb_tuned.predict(conn_fwd_ft)
        end_time=datetime.now()-time_start
        time_gb_tuned.append(end_time.total_seconds())

        time_start = datetime.now()
        gb_def.predict(conn_fwd_ft)
        end_time = datetime.now() - time_start
        time_gb_def.append(end_time.total_seconds())

        time_start = datetime.now()
        rf_tuned.predict(conn_fwd_ft)
        end_time = datetime.now() - time_start
        time_rf_tuned.append(end_time.total_seconds())

        time_start = datetime.now()
        rf_def.predict(conn_fwd_ft)
        end_time = datetime.now() - time_start
        time_rf_def.append(end_time.total_seconds())

    print(f' GB tuned: {sum(time_gb_tuned)/10} || {time_gb_tuned}')
    print(f' GB def: {sum(time_gb_def)/10} || {time_gb_def}')
    print(f' RF tuned: {sum(time_rf_tuned)/10} || {time_rf_tuned}')
    print(f' RF def: {sum(time_rf_def)/10} || {time_rf_def}')

    gb_tune_pred= gb_tuned.predict(conn_fwd_ft)
    gb_def_pred = gb_def.predict(conn_fwd_ft)
    rf_tune_pred = rf_tuned.predict(conn_fwd_ft)
    rf_def_pred = rf_def.predict(conn_fwd_ft)

    print('GB Tuned Results')
    model_general_results(gb_tune_pred, label)
    print('GB Tuned Results CV')
    get_model_cross_val_results(gb_tuned, conn_fwd_ft, label)
    print('GB Def Results')
    model_general_results(gb_def_pred, label)
    print('RF Tuned Results')
    model_general_results(rf_tune_pred, label)
    print('RF Tuned Results CV')
    get_model_cross_val_results(rf_tuned, conn_fwd_ft, label)
    print('RF Def Results')
    model_general_results(rf_def_pred, label)
    print('RF Def Results CV')
    get_model_cross_val_results(rf_def, conn_fwd_ft, label)

    import src.visualization as vis

    # vis.plot_cf_matrix(gb_tuned, conn_ft, label, 'GB Tuned')
    # vis.plot_cf_matrix(gb_def, conn_ft, label, 'GB def')
    # vis.plot_cf_matrix(rf_tuned, conn_ft, label, 'RF Tuned')
    # vis.plot_cf_matrix(rf_def, conn_ft, label, 'RF Def')

def model_general_results(predictions, y_test):
    print("=== Confusion Matrix ===")
    cf_matrix = confusion_matrix(y_test, predictions)
    print(cf_matrix)
    print('\n')
    print("=== Classification Report ===")
    class_report = classification_report(y_test, predictions)
    print(class_report)
    print('\n')
    print("=== ROC AUC Scores ===")
    print(roc_auc_score(y_test, predictions))
    print('\n')

def model_cross_val_results(model, ft_data, label):
    """
    Code adapted from https://github.com/StKyr/multiscorer/
    """
    scorer = MultiScorer({  # Create a MultiScorer instance
        'Accuracy': (accuracy_score, {}),
        'Precision': (precision_score, {'average': None}),
        'Recall': (recall_score, {'average': None}),
        'F-1': (f1_score, {'average': None}),
        'Area Under ROC Curve': (roc_auc_score, {})
    })

    cross_val_score(model, ft_data, label, scoring=scorer, cv=10)
    results = scorer.get_results()
    for metric in results.keys():  # Iterate and use the results
        if metric in ["Precision", 'Recall', 'F-1']:
            metric_vpn = []
            metric_non_vpn = []
            for sub_array in results[metric]:
                metric_non_vpn.append(sub_array[0])
                metric_vpn.append(sub_array[1])
            print(f'{metric}-Non VPN: {average(metric_non_vpn)}')
            print(f'{metric}-VPN: {average(metric_vpn)}')

        else:
            print(f'{metric}: {average(results[metric])}')

def get_model_general_results(predictions, y_test, prediction_probabilities, ft_set):
    # f1, f1_vpn = f1_score(y_test, predictions, average=None)
    prec, prec_vpn = precision_score(y_test, predictions, average=None)
    recc, recc_vpn = recall_score(y_test, predictions, average=None)
    auc = roc_auc_score(y_test, prediction_probabilities[:, 1])
    tn, fp, fn, tp = confusion_matrix(y_test, predictions).ravel()

    acc = accuracy_score(y_test, predictions)
    model_metrics_list = [ft_set, acc, prec, prec_vpn, recc, recc_vpn, auc, tn, fp, fn, tp]

    return model_metrics_list

def get_model_cross_val_results(model, ft_data, label):
    """
        Code adapted from https://github.com/StKyr/multiscorer/
        """
    scorer = MultiScorer({  # Create a MultiScorer instance
        'precision': (precision_score, {'average': None}),
        'recall': (recall_score, {'average': None}),
        'f1': (f1_score, {'average': None}),
        'accuracy': (accuracy_score, {}),
        'auc': (roc_auc_score, {})
    })

    model_cross_val_metrics_list = []
    cross_val_score(model, ft_data, label, scoring=scorer, cv=10)
    results = scorer.get_results()
    for metric in results.keys():  # Iterate and use the results
        if metric in ["precision", 'recall', 'f1']:
            metric_vpn = []
            metric_non_vpn = []
            for sub_array in results[metric]:
                metric_non_vpn.append(sub_array[0])
                metric_vpn.append(sub_array[1])

            model_cross_val_metrics_list.append(average(metric_non_vpn))
            model_cross_val_metrics_list.append(average(metric_vpn))

        else:

            model_cross_val_metrics_list.append(average(results[metric]))

    model_cross_val_metrics_string = ""
    for ele in model_cross_val_metrics_list:
        model_cross_val_metrics_string += f', {str(ele)}'
    model_cross_val_metrics_string = model_cross_val_metrics_string[2:]

    return model_cross_val_metrics_string, model_cross_val_metrics_list

if __name__ == '__main__':
    main()


import numpy as np
import pybullet as p
from .env import AssistiveEnv

import math
import os

from .agents.human import Human
from .agents.human import right_arm_joints, left_arm_joints, torso_joints, head_joints
controllable_joints = right_arm_joints + left_arm_joints + torso_joints + head_joints

def configure_human(human):
    human.impairment = None
    # human.set_all_joints_stiffness(0.02)
    human.set_whole_body_frictions(lateral_friction=50., spinning_friction=10., rolling_friction=10.)

    joint_pos = default_sitting_pose(human)
    human.setup_joints(joint_pos, use_static_joints=False, reactive_force=None)

    start_pos = [0, 0.05, 0.875]
    start_orient = [0, 0, 0, 1]
    human.set_base_pos_orient(start_pos, start_orient)
    # human.set_on_ground()

    joint_i = [pose[0] for pose in joint_pos]
    joint_th = [pose[1] for pose in joint_pos]
    joint_gains = [10.] * len(joint_i)
    # forces = [50.] * len(joint_i)
    forces = [1.] * len(joint_i)

    # tweak joint control
    for i in range(len(joint_gains)):
        if i not in controllable_joints:
            joint_gains[i] = 0.
            forces[i] = 0.

    human.control(joint_i, joint_th, joint_gains, forces)

# def set_joint_stiffnesses(human):
#     human.set_joint_stiffness(human.j_)

def default_sitting_pose(human):
    # Arms
    joint_pos = [(human.j_right_shoulder_x, 30.),
                 (human.j_left_shoulder_x, -30.),
                 (human.j_right_shoulder_y, 0.),
                 (human.j_left_shoulder_y, 0.),
                 (human.j_right_elbow, -90.),
                 (human.j_left_elbow, -90.)]

    # Legs
    joint_pos += [(human.j_right_knee, 90.),
                  (human.j_left_knee, 90.),
                  (human.j_right_hip_x, -90.),
                  (human.j_left_hip_x, -90.)]

    # Torso
    joint_pos += [(human.j_waist_x, 0.)]
    return joint_pos

class BasePoseEnv(AssistiveEnv):
    def __init__(self, human):
        super(BasePoseEnv, self).__init__(robot=None, human = human, task='pose_analysis', )

    def step(self, action):
        n_dof = len(action)
        action = np.zeros(n_dof).ravel()
        self.take_step(action)
        observation = self._get_obs()
        self.steps += 1
        reward = 0
        done = False if self.steps < self.max_steps else True
        info = {"n/a": 'n/a'}  # must be a dict
        return observation, reward, done, info

    def _get_obs(self, agent=None):
        # observation = self.human.get_joint_angles()
        observation = self.human.get_pos_orient(self.human.head)
        return observation

    def reset(self):
        super(BasePoseEnv, self).reset()
        self.build_assistive_env(fixed_human_base=False)
        plane_path = os.path.join(self.directory, "primitives", "plane_chair.urdf")
        plane_chair = p.loadURDF(plane_path, flags=p.URDF_USE_SELF_COLLISION, physicsClientId=self.id)

        # Human init
        configure_human(self.human)

        p.resetDebugVisualizerCamera(cameraDistance=1.10, cameraYaw=40, cameraPitch=-45, cameraTargetPosition=[-0.2, 0, 0.75], physicsClientId=self.id)
        p.configureDebugVisualizer(p.COV_ENABLE_RENDERING, 1, physicsClientId=self.id)
        return self._get_obs()

class StableRestingPoseEnv(BasePoseEnv):
    def __init__(self):
        human = Human(controllable_joint_indices=controllable_joints, controllable=True)
        super(BasePoseEnv, self).__init__(human=human)
        self.steps = 0
        self.max_steps = 100


import os
import os.path as osp
import bisect
import argparse
import multiprocessing as mp
import numpy as np
from tqdm import tqdm
import megengine as mge
from megengine import distributed as dist
from megengine import optimizer as optim
import megengine.autodiff as autodiff
from megengine import jit
import dataset, network
from config import config as cfg
from misc_utils import ensure_dir
import socket
import pdb

ensure_dir(cfg.output_dir)
logger = mge.get_logger(__name__)
log_path = osp.join(cfg.output_dir, 'logger.log')
mge.set_log_file(log_path, mode='a')

def find_free_port():

    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.bind(("", 0))
    port = sock.getsockname()[1]
    sock.close()
    return port

def allreduce_cb(param, grad, group=dist.WORLD):
    return dist.functional.all_reduce_sum(grad, group) / group.size

def train_one_epoch(model, gm, data_iter, opt, max_steps, rank, epoch_id, gpu_num):

    # @jit.trace(symbolic=False,)
    def propagate():
        with gm:
            loss_dict = model(model.inputs)
            total_loss = sum([loss_dict[key].mean() for key in loss_dict.keys()])
            gm.backward(total_loss)
            opt.step().clear_grad()
            loss_dict['total_loss'] = total_loss
        return loss_dict

    for step in range(max_steps):
        # learing rate
        if epoch_id == 0 and step < cfg.warm_iters:
            base_lr = (
                cfg.basic_lr * gpu_num * cfg.batch_per_gpu
                * (cfg.lr_decay_rate ** bisect.bisect_right(cfg.lr_decay_sates, epoch_id)
                )
            )
            lr_factor = (step + 1.0) / cfg.warm_iters
            for param_group in opt.param_groups:
                param_group['lr'] = (0.33 + 0.67 * lr_factor) * base_lr

        mini_batch = next(data_iter)
        im_info = mini_batch["im_info"]
        image = mini_batch["data"][:, :, :int(im_info[0, 0]), :int(im_info[0, 1])]
        model.inputs["image"].set_value(image)
        model.inputs["gt_boxes"].set_value(mini_batch["boxes"])
        model.inputs["im_info"].set_value(mini_batch["im_info"])
        del mini_batch
        losses = propagate()

        print_str = ' '
        for loss_name, loss_value in losses.items():
            print_str += ', {}: {:.4f}'.format(loss_name, loss_value.numpy())

        if rank == 0:
            if step % cfg.log_dump_interval == 0:

                logger.info(
                    "epoch-{}, {}/{}, lr: {:.4f}{}.\n{}".format(
                    epoch_id,
                    step,
                    max_steps,
                    opt.param_groups[0]["lr"],
                    print_str,
                    cfg.workspace,
                    )
                )

def worker(rank, gpu_num, args):
    # using sublinear
    os.environ["MGB_COMP_GRAPH_OPT"] = "enable_sublinear_memory_opt=1;seq_opt.enable_seq_comp_node_opt=0"
    os.environ["MGB_SUBLINEAR_MEMORY_GENETIC_NR_ITER"] = '10'
    os.environ['MGB_CUDA_RESERVE_MEMORY'] = '1'
    # establish the server if is the master

    dist_port = args.port
    if rank == 0:
        dist.Server(port=dist_port)
    if gpu_num> 1:

        dist.init_process_group(
            master_ip="localhost",
            port=dist_port,
            world_size=gpu_num,
            rank=rank,
            device=rank,
        )
        logger.info("Init process group for gpu%d done", rank)

    model = network.Network()
    params = model.parameters(requires_grad=True)
    model.train()

    # Autodiff gradient manager
    gm = autodiff.GradManager().attach(
        model.parameters(),
        callbacks=allreduce_cb,
    )

    opt = optim.SGD(
        params,
        lr=cfg.basic_lr * gpu_num * cfg.batch_per_gpu,
        momentum=cfg.momentum,
        weight_decay=cfg.weight_decay,
    )

    if cfg.pretrain_weight is not None:
        weights = mge.load(cfg.pretrain_weight)
        del weights['fc.weight']
        del weights['fc.bias']
        model.resnet50.load_state_dict(weights)

    start_epoch = 0
    if args.resume_weights is not None:
        assert osp.exists(args.resume_weights)
        model_file = args.resume_weights
        model_dict = mge.load(model_file)
        start_epoch, weights = model_dict['epoch'] + 1, model_dict['state_dict']
        model.load_state_dict(weights, strict=False)

    logger.info("Prepare dataset")
    train_loader = dataset.train_dataset(rank)
    logger.info("Training...")
    for epoch_id in range(start_epoch, cfg.max_epoch):
        for param_group in opt.param_groups:
            param_group["lr"] = (
                cfg.basic_lr * gpu_num * cfg.batch_per_gpu
                * (cfg.lr_decay_rate ** bisect.bisect_right(cfg.lr_decay_sates, epoch_id))
            )

        max_steps = cfg.nr_images_epoch // (cfg.batch_per_gpu * gpu_num)
        train_one_epoch(model, gm, train_loader, opt, max_steps, rank, epoch_id, gpu_num)
        if rank == 0:
            save_path = osp.join(cfg.model_dir, 'epoch-{}.pkl'.format(epoch_id + 1))
            state_dict = model.state_dict()
            names = [k for k, _ in state_dict.items()]
            for name in names:
                if name.startswith('inputs.'):
                    del state_dict[name]

            mge.save(
                {"epoch": epoch_id, "state_dict": state_dict}, save_path,
            )
            logger.info("dump weights to %s", save_path)

def train(args):
    # ------------------------ begin training -------------------------- #
    valid_nr_dev = mge.get_device_count("gpu")
    gpu_num = min(valid_nr_dev, args.num_gpus)
    assert gpu_num > 0
    logger.info('Device Count: {}'.format(gpu_num))

    ensure_dir(cfg.model_dir)

    if not osp.exists('output'):
        os.symlink(cfg.output_dir,'output')

    if gpu_num > 1:
        args.port =find_free_port()
        mp.set_start_method("spawn")
        processes = list()
        for i in range(gpu_num):
            process = mp.Process(target=worker, args=(i, gpu_num, args))
            process.start()
            processes.append(process)

        for p in processes:
            p.join()
    else:
        worker(0, 1, args)

def run_train():

    parser = argparse.ArgumentParser()

    parser.add_argument("--num_gpus", "-d", default=-1, type=int, help="total number of gpus for training")
    parser.add_argument('--resume_weights', '-r', default=None, type=str)
    parser.add_argument('--progressbar', '-p', action='store_true', default=False)
    parser.add_argument('--port', '-pt', type=int, default=11123)
    args = parser.parse_args()
    train(args)

if __name__ == '__main__':
    run_train()


import keras
import pandas as pd
#import datetime
#import serial.tools.list_ports
import os
import numpy as np
from scipy import signal
import matplotlib.pyplot as plt
import tensorflow as tf
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder
from sklearn.preprocessing import OneHotEncoder
X = []
y = []
maindirs = 'Feature_extract'
dirs = os.listdir(maindirs)
emosi = ['kaget','marah','santai','senang']
df = pd.read_csv(maindirs+"/"+"tes_extracted.csv")
d_t = df.drop('EMOSI',axis=1)
label = pd.get_dummies(df['EMOSI'])
print(label)
data_len = int(len(d_t))
for i in range (0,data_len):
    temp = d_t.loc[i]
    temp_list = temp.values.tolist()
    print(temp_list)
    X.append(temp_list)
for j in range(0,data_len):
    temp1 = label.iloc[j]
    # print(temp1)
    temp1_list = temp1.values.tolist()
    y.append(temp1_list)
print(X)
print(y)