import warnings

import astropy.units as u
import numpy as np
from astropy.coordinates import SkyCoord, SkyOffsetFrame, ICRS
from astropy.time import Time

from sora.config.decorators import deprecated_alias
from .meta import BaseEphem

__all__ = ['EphemHorizons', 'EphemJPL', 'EphemKernel', 'EphemPlanete']

class EphemPlanete(BaseEphem):
    """Class used to simulate former Fortran programs `ephem_planete` and
    `fit_d2_ksi_eta`.

    Attributes
    ----------
    ephem : `file`, required
        Input file with JD (UTC), geocentric RA (deg), DEC (deg), and
        distance (AU).

    name : `str`, optional, default=None
        Name of the object to search in the JPL database.

    radius : `int`, `float`, optional, default: online database
        Object radius, in km.

    error_ra : `int`, `float`, optional, default: online database
        Ephemeris RA*cosDEC error, in arcsec.

    error_dec : `int`, `float`, optional, default: online database
        Ephemeris DEC error, in arcsec.

    mass : `int`, `float`, optional. default=0
        Object mass, in kg.

    H : `int`, `float`, optional, default=NaN
        Object absolute magnitude.

    G : `int`, `float`, optional, default=NaN
        Object phase slope.

    """

    def __init__(self, ephem, name=None, spkid=None, **kwargs):

        super().__init__(name=name, spkid=spkid, **kwargs)
        data = np.loadtxt(ephem, unpack=True)
        self.time = Time(data[0], format='jd')
        self.ephem = SkyCoord(data[1] * u.deg, data[2] * u.deg, data[3] * u.AU)
        self.__reftime = self.time[0]
        self.min_time = Time(data[0].min(), format='jd')
        self.max_time = Time(data[0].max(), format='jd')

    def get_position(self, time, observer='geocenter'):
        """Returns the object's geocentric position.

        Parameters
        ----------
        time : `str`, `astropy.time.Time`
            Reference time to calculate the object position. It can be a string
            in the ISO format (yyyy-mm-dd hh:mm:ss.s) or an astropy Time object.

        observer : `any`
            This parameter is present in EphemPlanete for compatibility with the
            remaining ephem classes. The returned positions are based on the given
            ephemeris despite the observer.

        Returns
        -------
        coord : `astropy.coordinates.SkyCoord`
            Astropy SkyCoord object with the object coordinates at the given time.
        """
        ksi, eta = self.get_ksi_eta(time=time) * u.km
        distance = self.ephem.distance.mean()
        off_ra = np.arctan2(ksi, distance)
        off_dec = np.arctan2(eta, distance)
        coord_frame = SkyOffsetFrame(origin=self.star)
        pos = SkyCoord(lon=off_ra, lat=off_dec, distance=distance, frame=coord_frame)
        return pos.icrs

    @deprecated_alias(log='verbose')  # remove this line in v1.0
    def fit_d2_ksi_eta(self, star, verbose=True):
        """Fits the projected position (orthographic projection) of the object in
        the tangent sky plane relative to a star.

        Parameters
        ----------
        star : `str`, `astropy.coordinates.SkyCoord`
            The coordinate of the star in the same reference frame as the ephemeris.

        verbose : `bool`, optional, default=True
            Enable log printing.
        """
        if type(star) == str:
            star = SkyCoord(star, unit=(u.hourangle, u.deg))
        if hasattr(self, 'star') and self.star.to_string('hmsdms', precision=5) == star.to_string('hmsdms',
                                                                                precision=5):
            return
        self.star = star
        target = self.ephem.transform_to(SkyOffsetFrame(origin=star))
        da = target.cartesian.y
        dd = target.cartesian.z
        dt = (self.time - self.__reftime) / (self.max_time - self.min_time)

        self.ksi = np.polyfit(dt, da.to(u.km).value, 2)
        self.eta = np.polyfit(dt, dd.to(u.km).value, 2)
        ksi = np.poly1d(self.ksi)
        eta = np.poly1d(self.eta)
        dksi = da.to(u.km).value - ksi(dt)
        deta = dd.to(u.km).value - eta(dt)
        rmsk = np.sqrt(np.mean(np.square(dksi)))
        rmse = np.sqrt(np.mean(np.square(deta)))
        if verbose:
            output = ('Fitting ephemeris position relative to star coordinate {}\n'
                      'ksi = aksi*t\u00b2 + bksi*t + cksi\n'
                      'eta = aeta*t\u00b2 + beta*t + ceta\n'
                      't=(jd-{})/({}-{})\n'
                      '        aksi={}\n'
                      '        bksi={}\n'
                      '        cksi={}\n'
                      '        aeta={}\n'
                      '        beta={}\n'
                      '        ceta={}\n'
                      'Residual RMS: ksi={:.3f} km, eta={:.3f} km'.format(
                          self.star.to_string('hmsdms'), self.__reftime.jd, self.max_time.jd,
                          self.min_time.jd, *self.ksi, *self.eta, rmsk, rmse)
                      )
            print(output)

    def get_ksi_eta(self, time, star=None):
        """Returns the projected position (orthographic projection) of the object
        in the tangent sky plane relative to a star.

        Parameters
        ----------
        time : `str`, `astropy.time.Time`, required
            Reference time to calculate the position. It can be a string
            in the ISO format (yyyy-mm-dd hh:mm:ss.s) or an astropy Time object.

        star : `str`, `astropy.coordinates.SkyCoord`, optional, default=None
            The coordinate of the star in the same reference frame as the ephemeris.

        Returns
        -------
        ksi, eta : `float` array
            Projected position (orthographic projection) of the object in the
            tangent sky plane relative to a star.
            ``ksi`` is in the North-South direction (North positive).
            ``eta`` is in the East-West direction (East positive).
        """
        if star:
            self.fit_d2_ksi_eta(star)
        time = Time(time)
        if not time.isscalar:
            if any(time < self.min_time) or any(time > self.max_time):
                raise ValueError('time must be in the interval [{},{}]'.format(
                    self.min_time, self.max_time))
        elif time.isscalar:
            if time < self.min_time or time > self.max_time:
                raise ValueError('time must be in the interval [{},{}]'.format(
                    self.min_time, self.max_time))
        if hasattr(self, 'ksi') and hasattr(self, 'eta'):
            ksi = np.poly1d(self.ksi)
            eta = np.poly1d(self.eta)
            k = ksi((time - self.__reftime) / (self.max_time - self.min_time))
            e = eta((time - self.__reftime) / (self.max_time - self.min_time))
            if hasattr(self, 'offset'):
                dist = self.ephem[0].distance.to(u.km).value
                dksi = np.sin(self.offset.d_lon_coslat) * dist
                deta = np.sin(self.offset.d_lat) * dist
                k = k + dksi
                e = e + deta
            return k.value, e.value
        else:
            raise ValueError('A "star" parameter is missing. Please run fit_d2_ksi_eta first.')

    def __str__(self):
        """ String representation of the EphemPlanete Class.
        """
        validity = 'Valid from {} until {}'.format(self.min_time.iso, self.max_time.iso)
        out = super().__str__().format(ephem_info=validity)
        if hasattr(self, 'star'):
            out += ("\nFitted ephemeris position relative to star coordinate {}\n"
                    "    ksi = aksi*t\u00b2 + bksi*t + cksi\n"
                    "    eta = aeta*t\u00b2 + beta*t + ceta\n"
                    "    t=(jd-{})/({}-{})\n"
                    "        aksi={}\n"
                    "        bksi={}\n"
                    "        aeta={}\n"
                    "        beta={}\n"
                    "        ceta={}\n".format(
                        self.star.to_string('hmsdms'), self.__reftime.jd, self.max_time.jd,
                        self.min_time.jd, *self.ksi, *self.eta)
                    )
            target = self.ephem.transform_to(SkyOffsetFrame(origin=self.star))
            da = -target.cartesian.y
            dd = -target.cartesian.z
            dt = (self.time - self.__reftime) / (self.max_time - self.min_time)
            ksi = np.poly1d(self.ksi)
            eta = np.poly1d(self.eta)
            dksi = da.to(u.km).value - ksi(dt)
            deta = dd.to(u.km).value - eta(dt)
            rmsk = np.sqrt(np.mean(np.square(dksi)))
            rmse = np.sqrt(np.mean(np.square(deta)))
            out += 'Residual RMS: ksi={:.3f} km, eta={:.3f} km\n'.format(rmsk, rmse)
        return out

class EphemHorizons(BaseEphem):
    """Obtains the ephemeris from Horizons/JPL service.

    Note
    ----
    Web tool URL: https://ssd.jpl.nasa.gov/horizons.cgi

    Attributes
    ----------
    name : `str`, required
        Name of the object to search in the JPL database.

    id_type: `str`, default='smallbody'
        Type of object options: ``smallbody``, ``majorbody`` (planets but
        also anything that is not a small body), ``designation``, ``name``,
        ``asteroid_name``, ``comet_name``, ``id`` (Horizons id number), or
        ``smallbody`` (find the closest match under any id_type).

    radius : `int`, `float`, default: online database
        Object radius, in km.

    error_ra : `int`, `float`, default: online database
        Ephemeris RA*cosDEC error, in arcsec.

    error_dec : `int`, `float`, default: online database
        Ephemeris DEC error, in arcsec.

    mass : `int`, `float`, default=0
        Object mass, in kg.

    H : `int`, `float`, default=NaN
        Object absolute magnitude.

    G : `int`, `float`, default=NaN
        Object phase slope.

    """

    def __init__(self, name, id_type='smallbody', spkid=None, **kwargs):

        super().__init__(name=name, spkid=spkid, **kwargs)
        self.id_type = id_type
        _ = self.get_position(Time.now())  # test if Horizons can proceed for this object
        self.meta = {'kernels':'horizons'}

    def get_position(self, time, observer='geocenter'):
        """Returns the ICRS position of the object for observer.

        Parameters
        ----------
        time : `str`, `astropy.time.Time`
            Reference time to calculate the object position. It can be a string
            in the ISO format (yyyy-mm-dd hh:mm:ss.s) or an astropy Time object.

        observer : `str`, `sora.Observer`, `sora.Spacecraft`
            IAU code of the observer (must be present in given list of kernels),
            a SORA observer object or a string: ['geocenter', 'barycenter']

        Returns
        -------
        coord : `astropy.coordinates.SkyCoord`
            Astropy SkyCoord object with the object coordinates at the given time.
        """
        from .utils import ephem_horizons

        coord = ephem_horizons(time=time, target=self.name, observer=observer, id_type=self.id_type, output='ephemeris')
        if hasattr(self, 'offset'):
            pos_frame = SkyOffsetFrame(origin=coord)
            new_pos = SkyCoord(lon=self.offset.d_lon_coslat, lat=self.offset.d_lat,
                               distance=coord.distance, frame=pos_frame)
            coord = new_pos.transform_to(ICRS)
        if not coord.isscalar and len(coord) == 1:
            coord = coord[0]
        return coord

    def __str__(self):
        """ String representation of the EphemHorizons Class.
        """
        out = super().__str__().format(ephem_info='Ephemeris are downloaded from Horizons website')
        return out

# remove this block for v1.0
class EphemJPL(EphemHorizons):
    def __init__(self, name, id_type='smallbody', spkid=None, **kwargs):
        warnings.warn('EphemJPL is deprecated and will be removed in v1.0. Please use EphemHorizons.')
        super().__init__(name=name, id_type=id_type, spkid=spkid, **kwargs)

    __init__.__doc__ = EphemHorizons.__init__.__doc__
# end of block removal for v1.0

class EphemKernel(BaseEphem):
    """Gets the ephemeris from BSP kernels.

    Parameters
    ----------
    name : `str`,  optional, default=None
        Name of the object to search in the JPL database.

    spkid : `str`, required
        `spkid` of the targeting object. Former 'code' (v0.1).

    kernels : `list`, required
        List of paths for kernels files.

    radius : `int`, `float`, optional, default: online database
        Object radius, in km.

    error_ra : `int`, `float`, optional, default: online database
        Ephemeris RA*cosDEC error, in arcsec .

    error_dec : `int`, `float`, optional, default: online database
        Ephemeris DEC error, in arcsec.

    mass : `int`, `float`, optional, default=0
        Object Mass, in kg.

    H : `int`, `float`, optional, default=NaN
        Object Absolute Magnitude.

    G : `int`, `float`, optional, default=NaN
        Object Phase slope.

    """

    @deprecated_alias(code='spkid')  # remove this line for v1.0
    def __init__(self, kernels, spkid, name=None, **kwargs):
        import spiceypy as spice

        super().__init__(name=name, spkid=spkid, **kwargs)
        self.meta = {}
        kerns = []
        for arg in kernels:
            spice.furnsh(arg)
            kerns.append(arg.split('/')[-1].split('.')[0].upper())
        spice.kclear()
        self.meta['kernels'] = '/'.join(kerns)
        self.__kernels = kernels

    def get_position(self, time, observer='geocenter'):
        """Returns the object geocentric position.

        Parameters
        ----------
        time : `str`, `astropy.time.Time`
            Reference time to calculate the object position. It can be a string
            in the ISO format (yyyy-mm-dd hh:mm:ss.s) or an astropy Time object.

        observer : `str`, `sora.Observer`, `sora.Spacecraft`
            IAU code of the observer (must be present in given list of kernels),
            a SORA observer object or a string: ['geocenter', 'barycenter']

        Returns
        -------
        coord : `astropy.coordinates.SkyCoord`
            Astropy SkyCoord object with the object coordinates at the given time.
        """
        from .utils import ephem_kernel
        pos = ephem_kernel(time=time, target=self.spkid, observer=observer, kernels=self.__kernels)
        if hasattr(self, 'offset'):
            pos_frame = SkyOffsetFrame(origin=pos)
            new_pos = SkyCoord(lon=self.offset.d_lon_coslat, lat=self.offset.d_lat,
                               distance=pos.distance, frame=pos_frame)
            pos = new_pos.transform_to(ICRS)
        pos = SkyCoord(ra=pos.ra, dec=pos.dec, distance=pos.distance)
        if not pos.isscalar and len(pos) == 1:
            pos = pos[0]
        return pos

    def __str__(self):
        """ String representation of the EphemKernel Class.
        """
        out = super().__str__().format(ephem_info=self.meta['kernels'])
        return out

import numpy,import astropy,from astropy


# -*- coding: utf-8 -*-

from __future__ import absolute_import

import contextlib
import re
import copy
import operator
from types import MethodType
import collections

from . import mixture
from . import element
from . import compoundfromformula
from . import xrayspectrum
from . import compoundfromlist
from . import types
from ..utils import units
from ..utils import instance
from ..patch import xraylib
from ..io.mca import save as savemca
from ..utils.copyable import Copyable

import numpy as np
import scipy.interpolate
import scipy.optimize
import matplotlib.pyplot as plt
import fisx
from PyMca5.PyMcaPhysics.xrf import ClassMcaTheory
from PyMca5.PyMcaPhysics.xrf import ConcentrationsTool
from PyMca5.PyMcaIO import ConfigDict

try:
    from silx.gui import qt
    from PyMca5.PyMcaGui.physics.xrf import McaAdvancedFit
except ImportError:
    qt = None
    McaAdvancedFit = None

class PymcaBaseHandle(Copyable):
    def __init__(self):
        self.mcafit = ClassMcaTheory.McaTheory()
        self.ctool = ConcentrationsTool.ConcentrationsTool()
        self.app = None

    def loadfrompymca(self, filename=None, config=None):
        if filename is not None:
            fconfig = ConfigDict.ConfigDict()
            fconfig.read(filename)
            self.mcafit.configure(fconfig)
        if config is not None:
            self.mcafit.configure(config)

    def savepymca(self, filename):
        self.mcafit.config.write(filename)

    def configfromfitresult(self, digestedresult):
        config = copy.deepcopy(digestedresult["config"])
        # Global parameters
        parammap = {
            "Zero": ("detector", "zero"),
            "Gain": ("detector", "gain"),
            "Noise": ("detector", "noise"),
            "Fano": ("detector", "fano"),
            "Sum": ("detector", "sum"),
            "ST AreaR": ("peakshape", "st_arearatio"),
            "ST SlopeR": ("peakshape", "st_sloperatio"),
            "LT AreaR": ("peakshape", "lt_arearatio"),
            "LT SlopeR": ("peakshape", "lt_sloperatio"),
            "STEP HeightR": ("peakshape", "step_heightratio"),
            "Eta Factor": ("peakshape", "eta_factor"),
        }
        nparams = len(digestedresult["parameters"])
        ngroups = len(digestedresult["groups"])
        nglobal = nparams - ngroups
        for i in range(nglobal):
            param = digestedresult["parameters"][i]
            if param in parammap:
                key = parammap[param]
                config[key[0]][key[1]] = digestedresult["fittedpar"][i]
        return config

    @staticmethod
    def _parse_fitresult_massfractions(wfrac1, wfrac0, exclude):
        # Group per element
        wfrac2 = {}
        for zgroup, w in wfrac1.items():
            element = zgroup.element
            if element in exclude:
                continue
            group = zgroup.element.shells[0]
            if element not in wfrac2:
                wfrac2[element] = {}
            wfrac2[element][group] = w
        # Select mass fraction from most reliable group
        wfrac3 = {}
        for element, groupdict in wfrac2.items():
            w = sorted(groupdict.items(), key=operator.itemgetter(0))[0][1]
            wfrac3[element] = w
        # Modify original mass fractions
        wfracr = dict(wfrac0)
        snew = 1 - sum(w for element, w in wfrac0.items() if element not in wfrac3)
        sold = sum(wfrac3.values())
        norm = snew / sold
        for k, v in wfrac3.items():
            if v:
                wfracr[k] = v * norm
        return wfracr

    def _pymcaconfig_massfractions(self, config):
        """
        Extract elemental mass fractions from the pymca config sample
        """
        materials = collections.OrderedDict()
        enabled, name, density = config["attenuators"]["Matrix"][:3]
        if name == "MULTILAYER":
            for i in range(len(config["multilayer"])):
                enabled, name, density = config["multilayer"]["Layer{}".format(i)][:3]
                if enabled:
                    mat = self.loadfrompymca_material(config, name, density)
                    materials[name] = mat.elemental_massfractions()
        elif enabled:
            mat = self.loadfrompymca_material(config, name, density)
            materials[name] = mat.elemental_massfractions()
        return materials

    def configadjustconcentrations(self, config, lmassfractions, exclude=None):
        """
        Modify pymca config sample materials based on list of elemental concentrations
        """
        if not exclude:
            exclude = []
        materials = self._pymcaconfig_massfractions(config)
        if not materials:
            return
        if len(materials) != len(lmassfractions):
            raise RuntimeError("Number of layers in config and fitresult do not match")
        for (name, wfrac0), wfrac1 in zip(materials.items(), lmassfractions):
            wfrac2 = self._parse_fitresult_massfractions(wfrac1, wfrac0, exclude)
            CompoundList = [str(k) + "1" for k in wfrac2.keys()]
            CompoundFraction = list(wfrac2.values())
            config["materials"][name]["CompoundList"] = CompoundList
            config["materials"][name]["CompoundFraction"] = CompoundFraction

    def setdata(self, y):
        x = np.arange(len(y), dtype=np.float32)
        self.mcafit.setData(x, y)

    def fit(
        self,
        loadfromfit=False,
        concentrations=True,
        spectra=True,
        environ_elements=None,
    ):
        # Fit
        self.mcafit.estimate()
        fitresult, digestedresult = self.mcafit.startfit(digest=1)
        # Background
        yback = digestedresult.pop("continuum")  # polynomial + snip
        if self.mcafit.STRIP:
            ysnip = np.ravel(self.mcafit.zz)
        else:
            ysnip = 0
        digestedresult["yback"] = yback
        digestedresult["ysnip"] = ysnip
        # Parse result
        result = self.miscfromfitresult(digestedresult)
        if concentrations:
            self.concentrationsfromfitresult(digestedresult, out=result)
        if spectra:
            self.spectrafromfitresult(digestedresult, out=result)
        # Load parameters from fit (if you did "load from fit")
        if loadfromfit:
            config = self.configfromfitresult(digestedresult)
            if concentrations:
                self.configadjustconcentrations(
                    config, result["lmassfractions"], exclude=environ_elements
                )
            self.loadfrompymca(config=config)
        return result

    def fitgui(
        self,
        ylog=False,
        legend="data",
        loadfromfit=False,
        concentrations=True,
        spectra=True,
        environ_elements=None,
    ):
        if self.app is None:
            self.app = qt.QApplication([])
        w = McaAdvancedFit.McaAdvancedFit()

        # Copy mcafit
        x = self.mcafit.xdata0
        y = self.mcafit.ydata0
        w.setData(x, y, legend=legend, xmin=0, xmax=len(y) - 1)
        w.mcafit.configure(self.mcafit.getConfiguration())

        # GUI for fitting
        if ylog:
            w.graphWindow.yLogButton.click()
        w.graphWindow.energyButton.click()
        # w.graphWindow.setGraphYLimits(min(y[y>0]),None)
        w.refreshWidgets()
        w.show()
        result = self.app.exec_()

        # Parse result
        result = {}
        try:
            digestedresult = w.mcafit.digestresult()
        except AttributeError:
            digestedresult = None

        if digestedresult:
            # Load parameters from fit (if you did "load from fit")
            self.miscfromfitresult(digestedresult, out=result)
            if concentrations:
                self.concentrationsfromfitresult(digestedresult, out=result)
            if spectra:
                self.spectrafromfitresult(digestedresult, out=result)
            if loadfromfit:
                config = w.mcafit.getConfiguration()
                if concentrations:
                    self.configadjustconcentrations(
                        config, result["lmassfractions"], exclude=environ_elements
                    )
                self.loadfrompymca(config=config)
        else:
            if concentrations or spectra:
                result = self.fit(
                    loadfromfit=loadfromfit,
                    concentrations=concentrations,
                    spectra=spectra,
                )

        return result

    def processfitresult(self, digestedresult, originalconcentrations=False):
        ctoolcfg = self.ctool.configure()
        ctoolcfg.update(digestedresult["config"]["concentrations"])
        return self.ctool.processFitResult(
            config=ctoolcfg,
            fitresult={"result": digestedresult},
            elementsfrommatrix=originalconcentrations,
            fluorates=self.mcafit._fluoRates,
            addinfo=True,
        )

    def concentrationsfromfitresult(self, digestedresult, out=None):
        # Mass fractions:
        #
        # For a group (e.g Fe-K) and i loops over the layers:
        #  grouparea = flux.time.solidangle/(4.pi).sum_i[massfrac_i.grouprate_i]
        #
        # When Fe is present in only one layer j:
        #  massfrac_j = grouparea/(flux.time.solidangle/(4.pi).grouprate_i)
        #
        # When Fe present in multiple layers:
        #  grouparea = flux.time.solidangle/(4.pi).massfrac_avg.sum_i[b_i.grouprate_i]
        # where b_i is 1 or 0 depending on whether it is present in the matrix definition
        # When the element is not in the matrix definition b_i=1 for all layers.
        #
        #  massfrac_avg = grouparea/(flux.time.solidangle/(4.pi).sum_i[b_i.grouprate_i])
        if out is None:
            out = {}
        conresult, addinfo = self.processfitresult(
            digestedresult, originalconcentrations=False
        )

        out["massfractions"] = {
            element.Element.fluozgroup(k): v
            for k, v in conresult["mass fraction"].items()
        }

        if "layerlist" in conresult:
            nlayers = len(conresult["layerlist"])
            if nlayers > 0:
                out["lmassfractions"] = [
                    {
                        element.Element.fluozgroup(k): v
                        for k, v in conresult[k]["mass fraction"].items()
                    }
                    for k in conresult["layerlist"]
                ]
            else:
                nlayers = 1
                out["lmassfractions"] = []
        else:
            nlayers = 1
            out["lmassfractions"] = []

        if len(out["lmassfractions"]) == 0:
            out["lmassfractions"] = [out["massfractions"]]

        # Group rates:
        #   grouprate = solidangle/(4.pi).sum_i[massfrac_i.grouprate_i]   (i loops over the layers)
        out["rates"] = {}  # Ifluo / (Flux.time)
        safrac = addinfo["SolidAngle"]
        # assert(self.sample.geometry.solidangle/(4*np.pi)==addinfo["SolidAngle"])
        # assert(self.flux*self.time==addinfo["I0"])
        for zgroup in out["massfractions"]:
            ele = str(zgroup.element)
            group = str(zgroup.group)
            grouprate = 0.0
            for layer in range(1, nlayers + 1):
                if ele in self.mcafit._fluoRates[layer]:
                    massfrac_l = self.mcafit._fluoRates[layer][ele]["mass fraction"]
                    grouprate_l = self.mcafit._fluoRates[layer][ele]["rates"][
                        "{} xrays".format(group)
                    ]
                    grouprate += massfrac_l * grouprate_l
            grouprate *= safrac
            out["rates"][zgroup] = grouprate

        # Fitted areas (elements)
        out["fitareas"] = {
            element.Element.fluozgroup(k): v for k, v in conresult["fitarea"].items()
        }

        # Fitted areas (scatter)
        scatterpeaks = {
            "Compton": {"energy": [], "fitarea": []},
            "Rayleigh": {"energy": [], "fitarea": []},
        }
        nen = len(instance.asarray(digestedresult["energy"]))
        for k, v in digestedresult.items():
            if k.startswith("Scatter"):
                if "Compton" in k:
                    peak = "Compton"
                else:
                    peak = "Rayleigh"
                energy, fitarea = zip(
                    *[(v[k2]["energy"], v[k2]["fitarea"]) for k2 in v["peaks"]]
                )
                scatterpeaks[peak]["energy"] += energy
                scatterpeaks[peak]["fitarea"] += fitarea
        for peak in scatterpeaks:
            energy = np.array(scatterpeaks[peak]["energy"])
            fitarea = np.array(scatterpeaks[peak]["fitarea"])
            if peak == "Compton":
                k = xrayspectrum.ComptonLine(energy)
            else:
                k = xrayspectrum.RayleighLine(energy)
            out["fitareas"][k] = fitarea

        # Fitted rates
        out["fitrates"] = {k: v / addinfo["I0"] for k, v in out["fitareas"].items()}

        # self._print_pymcainternals_rates()
        return out

    def _pymcainternals_solidanglefrac(self):
        radius2 = self.mcafit.config["concentrations"]["area"] / np.pi
        distance = self.mcafit.config["concentrations"]["distance"]
        return 0.5 * (1.0 - (distance / np.sqrt(distance ** 2 + radius2)))

    def _print_pymcainternals_rates(self):
        safrac = self._pymcainternals_solidanglefrac()

        rowfmt = "{:>6}{:>8}{:>20}{:>10}{:>10}{:>20}"
        print(rowfmt.format("Layer", "Element", "MassFrac", "Line", "Energy", "Rate"))
        for i, layer in enumerate(self.mcafit._fluoRates):
            if i == 0:
                continue
            for ele in sorted(layer):
                lines = []
                for k, v in layer[ele].items():
                    if k.endswith("xrays"):
                        lines.extend(v)
                for line in lines:
                    w = layer[ele]["mass fraction"]
                    rate = layer[ele][line]["rate"] * w * safrac
                    energy = layer[ele][line]["energy"]
                    if rate > 0:
                        print(rowfmt.format(i - 1, ele, w, line, energy, rate))

    def miscfromfitresult(self, digestedresult, out=None):
        if out is None:
            out = {}
        out["chisq"] = digestedresult["chisq"]
        return out

    def spectrafromfitresult(self, digestedresult, out=None):
        conresult, addinfo = self.processfitresult(
            digestedresult, originalconcentrations=True
        )

        # Group rates:
        if False:
            # TODO: set rates of elements not in matrix to zero
            out["rates"] = {}
            for group in out["area"]:
                out["rates"][element.fluozgroup(group)] = (
                    conresult["area"][group] / addinfo["I0"]
                )

        # Matrix spectrum
        nparams = len(digestedresult["parameters"])
        ngroups = len(digestedresult["groups"])
        nglobal = nparams - ngroups
        parameters = list(digestedresult["fittedpar"])
        for i, group in enumerate(digestedresult["groups"]):
            if group in conresult["area"]:
                # Replace fitted with theoretical peak area
                parameters[nglobal + i] = conresult["area"][group]
            else:
                # Scattering peaks don't have a theoretical peak area
                parameters[nglobal + i] = 0.0
        ymatrix = self.mcafit.mcatheory(parameters, digestedresult["xdata"])

        yback = digestedresult["yback"]  # polynomial + snip
        ysnip = digestedresult["ysnip"]

        def interpol(x, spectrum):
            return scipy.interpolate.interp1d(
                x,
                spectrum,
                kind="nearest",
                bounds_error=False,
                fill_value=(spectrum[0], spectrum[-1]),
            )

        def interpol_energy(prof):
            return interpol(digestedresult["energy"], prof)

        def interpol_channel(prof):
            return interpol(digestedresult["xdata"], prof)

        if out is None:
            out = {}
        out["energy"] = digestedresult["energy"]
        out["channels"] = digestedresult["xdata"]
        out["y"] = digestedresult["ydata"]
        out["yfit"] = digestedresult["yfit"]
        out["yback"] = yback
        out["interpol_energy"] = interpol_energy
        out["interpol_channel"] = interpol_channel
        out["ypileup"] = digestedresult["pileup"]
        # polynomial is already in ymatrix, snip not
        out["ymatrix"] = ymatrix + ysnip

        def _plot(ylog=False, label="data"):
            plt.plot(out["energy"], out["y"], "+", label=label)
            plt.plot(out["energy"], out["yfit"], label=label + " fit")
            plt.plot(out["energy"], out["ymatrix"], label=label + " matrix")
            plt.plot(out["energy"], out["yback"], label=label + " background")
            if ylog:
                plt.gca().set_yscale("log", basey=10)
            plt.legend()
            plt.xlabel("MCA (keV)")
            plt.ylabel("Counts")

        out["plot"] = _plot

        return out

class PymcaHandle(PymcaBaseHandle):
    """Class that converts spectrocrunch to pymca parameters
    and vice versa. Also used to call the pymca fitting
    API for single spectra.
    """

    def __init__(
        self,
        sample=None,
        emin=None,
        emax=None,
        energy=None,
        weights=None,
        scatter=None,
        flux=1e9,
        time=0.1,
        escape=True,
        pileup=False,
        ninteractions=1,
        linear=False,
        snip=True,
        continuum=0,
        SingleLayerStrategy=None,
        noisepropagation=True,
    ):
        self.sample = sample
        self.set_source(energy=energy, weights=weights, scatter=scatter)
        self.emin = emin
        self.emax = emax
        self.linear = linear
        self.escape = escape
        self.pileup = pileup
        self.snip = snip
        self.continuum = continuum
        self.noisepropagation = noisepropagation
        self.ninteractions = ninteractions
        self.SingleLayerStrategy = SingleLayerStrategy
        self.flux = flux
        self.time = time
        super(PymcaHandle, self).__init__()

    def __str__(self):
        s = zip(
            instance.asarray(self.energy),
            instance.asarray(self.weights),
            instance.asarray(self.scatter),
        )
        s = "\n ".join(
            "{} keV: {} % (Scatter: {})".format(k, v * 100, sc) for k, v, sc in s
        )
        if self.sample:
            sample = str(self.sample)
            geometry = str(self.sample.geometry)
        else:
            sample = geometry = None
        return "Flux = {:~e}\nTime = {:~}\nSource lines:\n {}\n{}\n{}".format(
            self.flux, self.time, s, sample, geometry
        )

    def set_source(self, energy=None, weights=None, scatter=None):
        if energy is None:
            return
        self.energy = units.umagnitude(energy, "keV")
        if weights is None:
            self.weights = np.ones_like(self.energy)
        else:
            self.weights = weights
        if scatter is None:
            self.scatter = np.ones_like(self.energy)
        elif isinstance(scatter, bool):
            if scatter:
                self.scatter = np.ones_like(self.energy)
            else:
                self.scatter = np.zeros_like(self.energy)
        else:
            self.scatter = scatter

    @property
    def flux(self):
        return self._flux

    @flux.setter
    def flux(self, value):
        self._flux = units.Quantity(value, "Hz").to("Hz")

    @property
    def time(self):
        return self._time

    @time.setter
    def time(self, value):
        self._time = units.Quantity(value, "s").to("s")

    @property
    def I0(self):
        return (self.flux * self.time).to("dimensionless").magnitude

    @property
    def emax_strict(self):
        return np.max(self.energy)

    @property
    def emax(self):
        if self._emax is None:
            # include elastic scattering peaks
            p = np.max(self.energy)
            if self.sample:
                s = np.sqrt(self.sample.geometry.detector.gaussianVAR(p))
            else:
                s = 0.1
            return p + 3 * s
        else:
            return self._emax

    @emax.setter
    def emax(self, value):
        self._emax = value

    @property
    def emin(self):
        if self._emin is None:
            # include Al-K lines
            p = xraylib.LineEnergy(13, xraylib.KL3_LINE)
            if self.sample:
                s = np.sqrt(self.sample.geometry.detector.gaussianVAR(p))
            else:
                s = 0.1
            return p - 3 * s
        else:
            return self._emin

    @emin.setter
    def emin(self, value):
        self._emin = value

    def addtopymca_lstsq(self, cfg):
        cfg["fit"]["maxiter"] = 500

    def addtopymca_beam(self, cfg):
        # TODO: move to source?
        cfg["fit"]["energy"] = instance.asarray(self.energy).tolist()
        cfg["fit"]["energyweight"] = instance.asarray(self.weights).tolist()
        cfg["fit"]["energyscatter"] = (
            instance.asarray(self.scatter).astype(int).tolist()
        )
        cfg["fit"]["energyflag"] = np.ones_like(
            cfg["fit"]["energy"], dtype=int
        ).tolist()
        cfg["fit"]["scatterflag"] = int(any(cfg["fit"]["energyscatter"]))

        # Not just the source:
        cfg["concentrations"]["flux"] = self.flux.to("Hz").magnitude
        cfg["concentrations"]["time"] = self.time.to("s").magnitude

    def loadfrompymca_beam(self, cfg):
        ind = np.asarray(cfg["fit"]["energyflag"], dtype=bool).tolist()

        self.energy = instance.asarray(cfg["fit"]["energy"])
        self.energy = self.energy[ind].astype(float)
        self.weights = np.asarray(cfg["fit"]["energyweight"], dtype=float)
        self.weights = self.weights[ind]
        self.scatter = np.asarray(cfg["fit"]["energyscatter"], dtype=int)
        self.scatter = self.scatter[ind]
        self.flux = cfg["concentrations"]["flux"]
        self.time = cfg["concentrations"]["time"]

    def addtopymca_background(self, cfg):
        # bmodelbkg = self.sample.geometry.detector.bstail or \
        #            self.sample.geometry.detector.bltail or \
        #            self.sample.geometry.detector.bstep
        cfg["fit"]["stripflag"] = int(self.snip)
        cfg["fit"]["stripalgorithm"] = 1
        cfg["fit"]["snipwidth"] = 100
        cfg["fit"]["continuum"] = int(self.continuum)

    def loadfrompymca_background(self, cfg):
        self.snip = bool(cfg["fit"]["stripflag"])
        self.continuum = int(cfg["fit"]["continuum"])

    def addtopymca_other(self, cfg):
        cfg["fit"]["escapeflag"] = int(self.escape)
        cfg["fit"]["sumflag"] = int(self.pileup)
        cfg["fit"]["linearfitflag"] = int(self.linear)
        cfg["concentrations"]["usemultilayersecondary"] = self.ninteractions - 1
        cfg["fit"]["fitweight"] = int(self.noisepropagation)

    def addtopymca_custom(self, cfg):
        pass

    def addcustomnpymcamethod(self, method):
        self.addtopymca_custom = MethodType(method, self)

    def addtopymca_strategy(self, cfg):
        strategy = "SingleLayerStrategy"
        dic = getattr(self, strategy)
        if dic:
            dic = dict(dic)
            cfg["fit"]["strategy"] = strategy
            cfg["fit"]["strategyflag"] = 1
            if strategy not in cfg:
                cfg[strategy] = {}
                cfg[strategy]["layer"] = "Auto"
                cfg[strategy]["iterations"] = 3

            def func(mat):
                return mat if instance.isstring(mat) else mat.pymcaname

            cfg[strategy]["completer"] = func(dic.pop("completer", "-"))
            dic2 = {}
            for el, mat in dic.items():
                if el in cfg["peaks"]:
                    k = "{} {}".format(el, sorted(cfg["peaks"][el])[0])
                    dic2[k] = func(mat)
            cfg[strategy]["peaks"] = dic2.keys()
            cfg[strategy]["materials"] = dic2.values()
            cfg[strategy]["flags"] = [1] * len(dic2)
        else:
            cfg["fit"]["strategyflag"] = 0

    def loadfrompymca_other(self, cfg):
        self.escape = bool(cfg["fit"]["escapeflag"])
        self.pileup = bool(cfg["fit"]["sumflag"])
        self.linear = bool(cfg["fit"]["linearfitflag"])
        self.ninteractions = cfg["concentrations"]["usemultilayersecondary"] + 1
        self.noisepropagation = bool(cfg["fit"]["fitweight"])

    def addtopymca_material(self, cfg, material, defaultthickness=1e-4):
        return material.topymca(cfg, defaultthickness=defaultthickness)

    def loadfrompymca_material(self, cfg, matname, density):
        formula = re.compile("^(([A-Z][a-z]?)([0-9]*))+$")
        purelement = re.compile("^(?P<element>[A-Z][a-z]?)1?$")

        # Material name is a formula:
        name = mixture.Mixture.namefrompymca(matname)
        mf = formula.match(name)
        if mf:
            me = purelement.match(name)
            if me:
                return element.Element(me.groupdict()["element"])
            else:
                return compoundfromformula.CompoundFromFormula(name, density)

        # Parse material info:
        dic = cfg["materials"][matname]
        if not instance.isarray(dic["CompoundList"]):
            dic["CompoundList"] = [dic["CompoundList"]]
        if not instance.isarray(dic["CompoundFraction"]):
            dic["CompoundFraction"] = [dic["CompoundFraction"]]
        massfractions = np.array(dic["CompoundFraction"]) / sum(dic["CompoundFraction"])
        density = dic["Density"]

        # List of materials:
        lst = [
            self.loadfrompymca_material(cfg, mat, density)
            for mat in dic["CompoundList"]
        ]

        # All elements:
        if all([isinstance(e, element.Element) for e in lst]):
            return compoundfromlist.CompoundFromList(
                lst, massfractions, types.fraction.mass, density, name=name
            )
        else:
            return mixture.Mixture(lst, massfractions, types.fraction.mass, name=name)

    def loadfrompymca(self, filename=None, config=None):
        """Update pymca parameters and convert to spectrocrunch"""
        super(PymcaHandle, self).loadfrompymca(filename=filename, config=config)
        config = self.mcafit.getConfiguration()
        self.loadfrompymca_beam(config)
        self.loadfrompymca_background(config)
        self.loadfrompymca_other(config)
        self.sample.loadfrompymca(self, config)

    def addtopymca(self, fresh=False, onlygeometry=False):
        """Convert spectrocrunch to pymca parameters"""
        if fresh:
            config = self.mcafit.getStartingConfiguration()
            config["attenuators"] = {}
            try:
                self.mcafit.useFisxEscape(True)
            except:
                pass
            self.mcafit.disableOptimizedLinearFit()
        else:
            config = self.mcafit.getConfiguration()

        if onlygeometry:
            self.sample.geometry.addtopymca(self, config)
        else:
            self.addtopymca_beam(config)
            self.addtopymca_background(config)
            self.addtopymca_lstsq(config)
            self.addtopymca_other(config)
            self.sample.addtopymca(self, config)
            self.addtopymca_strategy(config)
        self.addtopymca_custom(config)
        self.mcafit.configure(config)

    def configurepymca(self, **kwargs):
        """Convert spectrocrunch to pymca parameters"""
        self.addtopymca(**kwargs)

    def savepymca(self, filename):
        """Convert spectrocrunch to pymca parameters and then save"""
        self.configurepymca()
        super(PymcaHandle, self).savepymca(filename)

    @contextlib.contextmanager
    def fitenergy_context(self):
        # Keep original config
        self.configurepymca()
        config = self.mcafit.getConfiguration()
        keep = self.emin, self.emax, self.energy

        # Change config (extend emin and emax)
        emin = np.min(self.energy)
        line = xrayspectrum.ComptonLine(emin)
        emin = line.energy(polar=np.pi)
        s = np.sqrt(self.sample.geometry.detector.gaussianVAR(emin))
        emin = min(keep[0], emin - 5 * s)
        emin = max(emin, 0)

        emax = np.max(self.energy)
        s = np.sqrt(self.sample.geometry.detector.gaussianVAR(emax))
        emax = max(keep[1], emax + 3 * s)

        self.emin, self.emax = emin, emax
        self.configurepymca()

        yield

        # Reset config
        self.emin, self.emax, self.energy = keep
        self.mcafit.configure(config)
        self.configurepymca()

    def fitenergy(self, plot=False):
        with self.fitenergy_context():
            # Fit reduced range
            result = self.fit()
            if plot:
                plt.plot(result["energy"], result["y"], label="data")
                plt.plot(result["energy"], result["yfit"], label="pymca fit")
                plt.legend()
                plt.show()

            config = self.mcafit.getConfiguration()

            # Repeat fit
            def func(params):
                params = instance.asarray(params)
                self.energy = params[:-1]
                self.sample.geometry.angleout = params[-1]
                self.configurepymca()

                self.mcafit.configure(config)

                fitresult, digestedresult = self.mcafit.startfit(digest=1)
                return digestedresult["yfit"] - digestedresult["ydata"]

            guess = instance.asarray(self.energy).tolist()
            guess.append(self.sample.geometry.angleout)

            params, success = scipy.optimize.leastsq(func, guess)

        success = success > 0 and success < 5
        if success:
            params = instance.asarray(params)
            self.energy = params[:-1]
            self.sample.geometry.angleout = params[-1]
            self.configurepymca()

    def xrayspectrum(self, **kwargs):
        return self.sample.xrayspectrum(
            self.energy,
            emin=self.emin,
            emax=self.emax,
            weights=self.weights,
            ninteractions=self.ninteractions,
            **kwargs
        )

    def xraygrouprates(self, **kwargs):
        spectrum = self.xrayspectrum(**kwargs)
        groups = spectrum.linegroups(convert=True)
        pattern = "^(?P<element>[A-Z][a-z]?)-(?P<group>[A-Z]).?$"
        groups2 = {}
        for k in groups:
            m = re.match(pattern, k)
            if m:
                m = m.groupdict()
                g = "{}-{}".format(m["element"], m["group"])
                A = sum(groups[k].values())
                if g in groups2:
                    groups2[g] += A
                else:
                    groups2[g] = A
            else:
                groups2[k] = sum(groups[k].values())

        return groups2

    def xraygroupareas(self, **kwargs):
        return {k: v * self.I0 for k, v in self.xraygrouprates(**kwargs).items()}

    def mca(self, histogram=True, decomposed=0, energy=False, full=False, **kwargs):
        spectrum = self.xrayspectrum(**kwargs)
        if decomposed == 0:
            x, y, ylabel = spectrum.sumspectrum(fluxtime=self.I0, histogram=histogram)
        elif decomposed == 1:
            x, y, ylabel = spectrum.peakspectrum(
                fluxtime=self.I0, histogram=histogram, group=True
            )
        else:
            x, y, ylabel = spectrum.peakspectrum(
                fluxtime=self.I0, histogram=histogram, group=False
            )
        if not energy:
            a, b = spectrum.channellimits
            nchan = int(2 ** np.ceil(np.log2(b + 1)))
            mca = np.zeros(nchan, dtype=y.dtype)
            mca[a : b + 1] = y
            y = mca
            x = np.arange(nchan)
        if full:
            return x, y, ylabel
        else:
            return y

    def savemca(self, filename, mode="w", func=None, **kwargs):
        y = self.mca(**kwargs)
        if instance.iscallable(func):
            y = func(y)
        savemca(
            y,
            filename,
            mode=mode,
            zero=self.sample.geometry.detector.mcazero,
            gain=self.sample.geometry.detector.mcagain,
        )

class FisxConfig:

    FISXMATERIALS = None

    def init(self):
        if self.FISXMATERIALS is None:
            self.FISXMATERIALS = fisx.Elements()
            self.FISXMATERIALS.initializeAsPyMca()

    @contextlib.contextmanager
    def init_ctx(self):
        self.init()
        yield

    def setDetector(self, setup, detector):
        with self.init_ctx():
            detector.addtofisx(setup, self.FISXMATERIALS)

    def addtofisx_material(self, material, defaultthickness=1e-4):
        with self.init_ctx():
            return material.tofisx(
                self.FISXMATERIALS, defaultthickness=defaultthickness
            )


"""This module implements an ORIGEN v2.2 transmutation solver.
"""
from __future__ import print_function, division

import os
import subprocess
import tempfile
from collections import Mapping
from warnings import warn
from pyne.utils import QAWarning

import numpy as np

from pyne import data
from pyne import rxname
from pyne import nucname
from pyne import nuc_data
from pyne import origen22
from pyne import utils
from pyne.material import Material, from_atom_frac
from pyne.xs.data_source import NullDataSource, SimpleDataSource, EAFDataSource
from pyne.xs.cache import XSCache

warn(__name__ + " is not yet QA compliant.", QAWarning)

class Transmuter(object):
    """A class for transmuting materials using ORIGEN v2.2."""

    def __init__(self, t=0.0, phi=0.0, temp=300.0, tol=1e-7, cwd='',
                 base_tape9=origen22.BASE_TAPE9, xscache=None,
                 o2exe='o2_therm_linux.exe', *args, **kwargs):
        """Parameters
        ----------
        t : float
            Transmutations time [sec].
        phi : float or array of floats
            Neutron flux vector [n/cm^2/sec].  Currently this must either be
            a scalar or match the group structure of EAF.
        temp : float, optional
            Temperature [K] of material, defaults to 300.0.
        tol : float
            Tolerance level for chain truncation.
        cwd : str, optional
            Current working directory for origen runs. Defaults to this dir.
        base_tape9 : str or dict, optional
            A base TAPE9.INP file.  If this is a str it is interpreted as a path
            to a file, which is then read in and parsed.  If this is a dict, it is
            assumed to be in the format described in the main origen22 module.
        xscache : XSCache, optional
            A cross section cache to generate cross sections with.
        o2exe : str, optional
            Name or path to ORIGEN 2.2 executable.
        args : tuple, optional
            Other arguments ignored for compatibility with other Transmuters.
        kwargs : dict, optional
            Other keyword arguments ignored for compatibility with other Transmuters.
        """
        if not isinstance(base_tape9, Mapping):
            base_tape9 = origen22.parse_tape9(tape9=base_tape9)
        self.base_tape9 = base_tape9

        if xscache is None:
            eafds = EAFDataSource()
            eafds.load(temp=temp)
            gs = np.array([eafds.src_group_struct[0], eafds.src_group_struct[-1]])
            eafds.dst_group_struct = gs
            xscache = XSCache(group_struct=gs, data_source_classes=[SimpleDataSource,
                                                                    NullDataSource])
            xscache.load(temp=temp)
            xscache.data_sources.insert(0, eafds)
        self.xscache = xscache

        self.t = t
        self._phi = None
        self.phi = phi
        self.temp = temp
        self.tol = tol
        self.cwd = os.path.abspath(cwd)
        self.o2exe = o2exe

    @property
    def phi(self):
        return self._phi

    @phi.setter
    def phi(self, flux):
        """Ensures that the flux is correctly formatted."""
        flux = np.asarray(flux)
        if flux.ndim == 0:
            _ = np.empty(175, float)
            _.fill(flux / 175.0)
            flux = _
        elif flux.ndim == 1 and flux.shape[0] != 175:
            raise ValueError("Group structure must match EAF.")
        elif flux.ndim > 1:
            raise ValueError("The flux vector must be 0- or 1-dimensional.")
        if not np.all(flux >= 0.0):
            raise ValueError("Flux entries must be non-negative.")
        for ds in self.xscache.data_sources:
            ds.src_phi_g = flux
        self.xscache['phi_g'] = np.array([flux.sum()])
        self._phi = flux

    def transmute(self, x, t=None, phi=None, tol=None, cwd=None, xscache=None,
                  o2exe=None, *args, **kwargs):
        """Transmutes a material into its daughters.

        Parameters
        ----------
        x : Material or similar
            Input material for transmutation.
        t : float
            Transmutations time [sec].
        phi : float or array of floats
            Neutron flux vector [n/cm^2/sec].  Currently this must either be
            a scalar or match the group structure of EAF.
        tol : float
            Tolerance level for chain truncation.
        cwd : str, optional
            Current working directory for origen runs. Defaults to this dir.
        xscache : XSCache, optional
            A cross section cache to generate cross sections with.
        o2exe : str, optional
            Name or path to ORIGEN 2.2 executable.

        Returns
        -------
        y : Material
            The output material post-transmutation.

        """
        if not isinstance(x, Material):
            x = Material(x)
        if t is not None:
            self.t = t
        if phi is not None:
            self.phi = phi
        if tol is not None:
            self.tol = tol
        if cwd is not None:
            self.cwd = os.path.abspath(cwd)
        if xscache is not None:
            self.xscache = xscache
        if o2exe is not None:
            self.o2exe = o2exe

        # prepare new tape9
        nucs = set(x.comp.keys())
        base_tape9 = self.base_tape9
        decay_nlb, xsfpy_nlb = origen22.nlbs(base_tape9)
        new_tape9 = origen22.xslibs(nucs=nucs, xscache=self.xscache, nlb=xsfpy_nlb)
        t9 = origen22.merge_tape9([new_tape9, base_tape9])

        # write out files
        origen22.write_tape4(x, outfile=os.path.join(self.cwd, 'TAPE4.INP'))
        origen22.write_tape5_irradiation('IRF', self.t/86400.0, self.xscache['phi_g'][0],
            outfile=os.path.join(self.cwd, 'TAPE5.INP'), decay_nlb=decay_nlb,
            xsfpy_nlb=xsfpy_nlb, cut_off=self.tol)
        origen22.write_tape9(t9, outfile=os.path.join(self.cwd, 'TAPE9.INP'))

        # run origen & get results
        f = tempfile.NamedTemporaryFile()
        try:
            subprocess.check_call([self.o2exe], cwd=self.cwd, stdout=f, stderr=f)
        except subprocess.CalledProcessError:
            f.seek(0)
            print("ORIGEN output:\n\n{0}".format(f.read()))
            raise
        finally:
            f.close()
        t6 = origen22.parse_tape6(tape6=os.path.join(self.cwd, 'TAPE6.OUT'))
        y = t6['materials'][-1]
        return y


import numpy as np
import pytest

from jesse.config import config, reset_config
from jesse.store import store

def set_up():
    """

    """
    reset_config()
    config['app']['considering_candles'] = [('Sandbox', 'BTCUSD')]
    store.reset()
    store.orderbooks.init_storage()

def test_add_orderbook_and_orderbook_getters():
    set_up()

    asks = [[9188.76, 52.66], [9189.68, 27.58], [9190.6, 29.36], [9191.52, 47.02], [9192.44, 31.44], [9193.36, 37.89],
            [9194.28, 44.37], [9195.2, 33.84], [9196.12, 26.23], [9197.04, 27.38], [9197.96, 51.1], [9198.88, 56.8],
            [9199.8, 52.01], [9200.72, 46.82], [9201.64, 58.11], [9220.0, 7.579], [9221.45, 1.0], [9232.71, 0.573],
            [9234.15, 0.595], [9236.9, 0.55], [9237.22, 0.573], [9237.92, 0.555], [9238.13, 0.562], [9238.52, 0.546],
            [9238.7, 0.557], [9239.0, 50.0], [9239.42, 0.546], [9240.59, 0.571], [9241.0, 1.0], [9243.0, 24.165],
            [9255.0, 4.0], [9256.66, 0.005], [9274.0, 95.479], [9279.52, 0.129], [9280.0, 0.05], [9300.0, 222.311],
            [9312.0, 4.336], [9330.0, 45.0], [9340.0, 190.0], [9345.0, 50.0], [9349.83, 0.03], [9406.19, 20.0],
            [9431.0, 5.347], [9450.0, 1.0], [9458.0, 2.0], [9460.0, 0.02], [9489.0, 215.11], [9550.0, 134.357],
            [9562.0, 25.0], [9582.0, 17.304], [9600.0, 21.784], [9640.0, 50.0], [9650.0, 1.0], [9692.0, 1000.0],
            [9715.0, 1.0], [9738.0, 0.5], [9792.0, 1.0], [9800.0, 1.0], [9813.0, 0.5], [9827.0, 50.0], [9990.0, 80.0],
            [9999.0, 1000.0], [10000.0, 2000.0], [10001.0, 50.0], [10119.0, 1.5], [10132.56, 1.957], [10186.56, 0.001],
            [10270.0, 2.0], [10447.0, 25.0], [10500.0, 200.0], [10600.0, 41.954], [71080.0, 1.0]]
    bids = [[9188.75, 53.43], [9187.83, 21.68], [9186.91, 40.64], [9185.99, 24.93], [9185.07, 51.37], [9184.15, 20.09],
            [9183.23, 49.1], [9182.31, 24.39], [9181.39, 20.79], [9181.12, 1.089], [9180.47, 35.67], [9179.55, 21.63],
            [9178.63, 35.59], [9177.71, 52.51], [9177.0, 11.866], [9176.79, 56.2], [9175.87, 37.16], [9168.0, 0.01],
            [9153.25, 0.01], [9150.0, 10.5], [9148.0, 0.01], [9130.0, 1.0], [9128.0, 0.01], [9108.0, 0.01],
            [9106.78, 0.01], [9100.0, 0.5], [9088.0, 0.01], [9080.0, 180.0], [9078.0, 35.0], [9068.0, 0.01],
            [9060.32, 0.011], [9050.0, 0.5], [9048.0, 0.01], [9028.0, 80.0], [9009.0, 7.0], [9002.0, 100.0],
            [9000.0, 40.325], [8999.0, 13.231], [8989.0, 0.1], [8965.0, 5.0], [8961.0, 300.0], [8960.06, 1.0],
            [8955.0, 300.0], [8911.0, 16.254], [8902.0, 340.0], [8900.0, 5.0], [8871.0, 25.0], [8866.42, 0.001],
            [8800.0, 1.0], [8777.0, 38.114], [8705.0, 10.0], [8700.0, 34.173], [8647.17, 10.0], [8600.0, 13.0],
            [8595.0, 20.0], [8550.0, 10.0], [8525.86, 5.823], [8500.0, 3.25], [8488.8, 0.007], [8465.11, 62.411],
            [8450.0, 80.0], [8443.8, 2.047], [8429.4, 0.001], [8400.0, 2.25], [8390.0, 160.0], [8388.0, 2.0],
            [8300.0, 7.25], [8200.0, 2.25], [8100.0, 54.0], [8085.0, 4.52], [8008.0, 3.747], [8000.0, 15.5],
            [7900.0, 124.505], [7855.0, 2.5], [7851.0, 100.0], [7821.0, 0.49], [7803.0, 0.001], [7800.0, 1.0],
            [7760.0, 5.92], [7702.0, 50.0], [7700.0, 1.0], [7675.0, 0.23], [7600.0, 1.0], [7568.3, 0.001],
            [7501.0, 50.0], [7500.0, 1.001], [7421.0, 50.0], [7405.0, 100.0], [7401.0, 100.0], [7011.0, 21.339],
            [7000.0, 206.0], [6986.0, 50.0], [6985.0, 0.076], [6950.0, 0.076], [6942.0, 0.02], [6915.0, 0.076],
            [6879.0, 0.03], [6798.0, 0.02], [6721.0, 100.0], [6450.0, 1.5], [6415.0, 172.154], [6404.95, 0.01],
            [6400.0, 1.0], [6372.92, 0.01], [6358.0, 0.182], [6341.05, 0.01], [6309.34, 0.01], [6301.0, 501.5],
            [6280.56, 0.085], [6277.79, 0.01], [6246.4, 0.01], [6220.0, 200.0], [6215.16, 0.01], [6184.08, 0.01],
            [6153.15, 0.01], [6122.38, 0.01], [6091.76, 0.01], [6061.3, 0.01], [6030.99, 0.01], [6021.0, 30.0],
            [6000.83, 0.01], [6000.0, 16.701]]
    store.orderbooks.add_orderbook('Sandbox', 'BTCUSD', asks, bids)

    # test get_best_ask
    np.testing.assert_equal(
        store.orderbooks.get_best_ask('Sandbox', 'BTCUSD'),
        np.array(
            [9189.0, 52.66]
        )
    )
    # test get_current_asks
    np.testing.assert_equal(
        store.orderbooks.get_current_asks('Sandbox', 'BTCUSD')[0],
        np.array(
            [9189.0, 52.66]
        )
    )
    np.testing.assert_equal(
        store.orderbooks.get_current_asks('Sandbox', 'BTCUSD')[1],
        np.array(
            [9190.0, 27.58]
        )
    )

    # test get_best_bid
    np.testing.assert_equal(
        store.orderbooks.get_best_bid('Sandbox', 'BTCUSD'),
        np.array(
            [9188.0, 53.43]
        )
    )
    # test get_current_bids
    np.testing.assert_equal(
        store.orderbooks.get_current_bids('Sandbox', 'BTCUSD')[0],
        np.array(
            [9188.0, 53.43]
        )
    )
    np.testing.assert_equal(
        store.orderbooks.get_current_bids('Sandbox', 'BTCUSD')[1],
        np.array(
            [9187.0, 21.68]
        )
    )

    # test get_current_orderbook
    np.testing.assert_equal(
        store.orderbooks.get_current_orderbook('Sandbox', 'BTCUSD')[0][0],
        np.array(
            [9189.0, 52.66]
        )
    )

    # test get_orderbooks
    np.testing.assert_equal(
        store.orderbooks.get_orderbooks('Sandbox', 'BTCUSD')[-1][0][0],
        np.array(
            [9189.0, 52.66]
        )
    )

def test_fix_array_len():
    from jesse.store.state_orderbook import _fix_array_len

    a = np.array([
        1, 2, 3, 4, 5
    ], dtype=float)

    a = _fix_array_len(a, 7)
    b = np.array([
        1, 2, 3, 4, 5
    ], dtype=float)

    assert np.array_equiv(a[:5], b)
    assert np.isnan(a[5])
    assert np.isnan(a[6])

    c = np.array([
        1, 2, 3, 4, 5
    ], dtype=float)

    # assert that len has to be >= len(a)
    with pytest.raises(ValueError):
        _fix_array_len(c, 3)
def test_trim_orderbook_list():
    from jesse.store.state_orderbook import _trim_orderbook_list

    trimmed_asks = _trim_orderbook_list(
        [[9414.61, 44.84], [9415.55, 50.31], [9416.49, 49.13], [9417.43, 49.57], [9418.37, 58.71], [9419.31, 40.67],
         [9420.25, 32.2], [9421.19, 32.06], [9422.13, 40.38], [9423.07, 56.83], [9424.01, 53.35], [9424.45, 1.0],
         [9424.95, 35.19], [9425.89, 24.71], [9427.77, 23.11], [9450.0, 10.0], [9454.2, 0.005], [9468.25, 279.229],
         [9470.0, 6.104], [9478.0, 194.239], [9490.0, 0.11], [9500.0, 5.1], [9500.87, 157.017], [9511.0, 20.0],
         [9592.0, 232.788], [9599.0, 0.13], [9600.0, 11.796], [9639.0, 10.0], [9650.0, 1.0], [9692.0, 1000.0],
         [9715.0, 1.0], [9792.0, 1.0], [9800.0, 23.0], [9999.0, 1000.0], [10000.0, 2000.0], [10001.0, 50.0],
         [10119.0, 1.5], [10186.56, 0.001], [10270.0, 2.0], [10500.0, 200.0], [10600.0, 41.954], [71080.0, 1.0]],
        ascending=True
    )
    np.testing.assert_equal(trimmed_asks,
                            [[9415, 44.84], [9416, 50.31], [9417, 49.13], [9418, 49.57], [9419, 58.71], [9420, 40.67],
                             [9421, 32.2],
                             [9422, 32.06], [9423, 40.38], [9424, 56.83], [9425, 89.53999999999999], [9426, 24.71],
                             [9428, 23.11],
                             [9450, 10.0], [9455, 0.005], [9469, 279.229], [9470, 6.104], [9478, 194.239], [9490, 0.11],
                             [9500, 5.1],
                             [9501, 157.017], [9511, 20.0], [9592, 232.788], [9599, 0.13], [9600, 11.796], [9639, 10.0],
                             [9650, 1.0],
                             [9692, 1000.0], [9715, 1.0], [9792, 1.0], [9800, 23.0], [9999, 1000.0], [10000, 2000.0],
                             [10001, 50.0],
                             [10119, 1.5], [10187, 0.001], [10270, 2.0], [10500, 200.0], [10600, 41.954]]
                            )
    trimmed_bids = _trim_orderbook_list(
        [[9414.6, 97.2], [9410.84, 54.62], [9408.96, 41.98], [9408.02, 34.65], [9401.82, 56.84], [9400.88, 23.81],
         [9386.87, 0.005], [9348.25, 7.156], [9346.0, 1.0], [9330.0, 30.0], [9327.13, 100.0], [9303.0, 0.02],
         [9300.99, 1.0], [9300.0, 3.0], [9296.0, 31.0], [9221.0, 7.0], [9200.0, 21.991], [9120.0, 10.0],
         [9080.0, 180.0], [9060.0, 180.0], [9056.0, 30.0], [9028.0, 80.0], [9002.0, 100.0], [9000.0, 41.825],
         [8999.0, 13.231], [8989.0, 0.1], [8970.0, 107.0], [8965.0, 5.0], [8961.0, 300.0], [8960.06, 1.0],
         [8955.0, 300.0], [8940.71, 0.001], [8920.0, 20.0], [8911.0, 16.254], [8909.32, 532.426], [8902.0, 340.0],
         [8900.0, 5.5], [8888.0, 20.0], [8871.92, 0.001], [8800.0, 4.0], [8777.0, 38.114], [8705.0, 10.0],
         [8700.0, 33.173], [8647.17, 10.0], [8600.0, 20.0], [8595.0, 20.0], [8550.0, 35.0], [8525.86, 5.823],
         [8500.0, 1.0], [8488.8, 0.007], [8465.11, 62.411], [8450.0, 80.0], [8429.4, 0.001], [8400.0, 4.5],
         [8390.0, 160.0], [8388.0, 2.0], [8300.0, 6.0], [8200.0, 1.0], [8100.0, 51.5], [8085.0, 4.52], [8008.0, 3.747],
         [8000.0, 14.0], [7900.0, 124.505], [7855.0, 2.5], [7851.0, 100.0], [7821.0, 0.49], [7803.0, 0.001],
         [7800.0, 1.0], [7760.0, 5.92], [7702.0, 50.0], [7700.0, 1.0], [7675.0, 0.23], [7600.0, 1.0], [7568.3, 0.001],
         [7501.0, 50.0], [7500.0, 1.001], [7421.0, 50.0], [7405.0, 100.0], [7401.0, 100.0], [7011.0, 21.339],
         [7000.0, 206.0], [6986.0, 50.0], [6985.0, 0.076], [6950.0, 0.076], [6942.0, 0.02], [6915.0, 0.076],
         [6879.0, 0.03], [6798.0, 0.02], [6721.0, 100.0], [6450.0, 1.5], [6415.0, 172.154], [6404.95, 0.01],
         [6400.0, 1.0], [6372.92, 0.01], [6358.0, 0.182], [6341.05, 0.01], [6309.34, 0.01], [6301.0, 501.5],
         [6280.56, 0.085], [6277.79, 0.01], [6246.4, 0.01], [6220.0, 200.0], [6215.16, 0.01], [6184.08, 0.01],
         [6153.15, 0.01], [6122.38, 0.01], [6091.76, 0.01], [6061.3, 0.01], [6030.99, 0.01], [6021.0, 30.0],
         [6000.83, 0.01], [6000.0, 16.701]],
        ascending=False
    )
    print(trimmed_bids)
    np.testing.assert_equal(
        trimmed_bids,
        [[9414.0, 97.2], [9410.0, 54.62], [9408.0, 76.63], [9401.0, 56.84], [9400.0, 23.81], [9386.0, 0.005],
         [9348.0, 7.156], [9346.0, 1.0], [9330.0, 30.0], [9327.0, 100.0], [9303.0, 0.02], [9300.0, 4.0], [9296.0, 31.0],
         [9221.0, 7.0], [9200.0, 21.991], [9120.0, 10.0], [9080.0, 180.0], [9060.0, 180.0], [9056.0, 30.0],
         [9028.0, 80.0], [9002.0, 100.0], [9000.0, 41.825], [8999.0, 13.231], [8989.0, 0.1], [8970.0, 107.0],
         [8965.0, 5.0], [8961.0, 300.0], [8960.0, 1.0], [8955.0, 300.0], [8940.0, 0.001], [8920.0, 20.0],
         [8911.0, 16.254], [8909.0, 532.426], [8902.0, 340.0], [8900.0, 5.5], [8888.0, 20.0], [8871.0, 0.001],
         [8800.0, 4.0], [8777.0, 38.114], [8705.0, 10.0], [8700.0, 33.173], [8647.0, 10.0], [8600.0, 20.0],
         [8595.0, 20.0], [8550.0, 35.0], [8525.0, 5.823], [8500.0, 1.0], [8488.0, 0.007], [8465.0, 62.411],
         [8450.0, 80.0]]
    )
    assert len(trimmed_bids) == 50


import numpy as np
from polytrack.general import cal_dist, check_sight_coordinates
# import polytrack.bg_subtraction as polytrack_bgs
import itertools as it
from scipy.optimize import linear_sum_assignment
from polytrack.deep_learning import detect_deep_learning
from polytrack.config import pt_cfg

max_dist_bs = pt_cfg.POLYTRACK.MAX_DIST_BS #Maximum detection threshold for background subtraction-based detection
max_dist_dl = pt_cfg.POLYTRACK.MAX_DIST_DL
max_bg_changes = pt_cfg.POLYTRACK.MAX_BG_CHANGES #Threshold value of background changes (if background changes exceeds this value, deep learning-based detection will be used)

def evaluate_noisy(noisy):
    if noisy:
        pt_cfg.POLYTRACK.NOISY = True
        pt_cfg.POLYTRACK.CLEAR_FRAMES = 0
    else:
        if not pt_cfg.POLYTRACK.NOISY:
            pass
        else:
            if pt_cfg.POLYTRACK.CLEAR_FRAMES < pt_cfg.POLYTRACK.CLEAR_FRAMES_THRESH:
                pt_cfg.POLYTRACK.CLEAR_FRAMES += 1
            else:
                pt_cfg.POLYTRACK.NOISY = False
    return None

def track(frame, _predicted_position, _insectsBS):
    _predictions = _predicted_position
    _detections_BS = _insectsBS

    if(len(_insectsBS)<=max_bg_changes):
        _associated_det_BS, _missing_BS, _not_associated_BS = associate_detections_BS(_insectsBS, _predicted_position)
        evaluate_noisy(False)

    else:
        _missing_BS = [i[0] for i in _predicted_position]
        _associated_det_BS = []
        _not_associated_BS  =[]
        evaluate_noisy(True)

    #print(pt_cfg.POLYTRACK.NOISY)

    if (run_DL(_missing_BS, _not_associated_BS, _insectsBS)):
        #print('here')

        _predictions_DL = np.zeros(shape=(0,3))

        for pred in np.arange(len(_missing_BS)):
            _predictions_DL = np.vstack([_predictions_DL,([row for row in _predictions if _missing_BS[pred] == row[0]])])

        _detections_DL = detect_deep_learning(frame)

        _associated_det_DL, _missing, _pot_new_insect = associate_detections_DL(_detections_DL, _predictions_DL, max_dist_dl)
        _new_insect = verify_new_insect(_pot_new_insect,_associated_det_BS)
    else:
        _associated_det_DL, _missing, _new_insect =[],[], []

    return _associated_det_BS, _associated_det_DL, _missing, _new_insect

def run_DL(_missing_BS, _not_associated_BS, _insectsBS):
    n_missing = len(_missing_BS)
    n_not_assciated_BS = len(_not_associated_BS)
    noisy_background = len(_insectsBS)>max_bg_changes

    if((n_missing>0) or (n_not_assciated_BS>0) or noisy_background):
        _run_DL = True
    else:
        _run_DL = False

    return _run_DL

def new_insect_mode(pot_new_insect):
    if (len(pot_new_insect))>0:
        pt_cfg.POLYTRACK.NEW_INSECT_MODE = True
    else:
       pt_cfg.POLYTRACK.NEW_INSECT_MODE = False

    return None

def dark_spot_check(_pot_new_insect):
    _new_insect = []
    for _insect in _pot_new_insect:
        if check_sight_coordinates(_insect):
            _new_insect.append(_insect)

        else:
            pass

    return _new_insect

def verify_new_insect(_pot_new_insect,_associated_det_BS):
    # print(_pot_new_insect, _associated_det_BS)
    recorded_BS = []
    for bsdet in _associated_det_BS:
        for pni in np.arange(len(_pot_new_insect)):
            dist = cal_dist(bsdet[1],bsdet[2],int(float(_pot_new_insect[pni][0])),int(float(_pot_new_insect[pni][1])))
            if dist <= max_dist_bs:
                recorded_BS.append(pni)

    _pot_new_insect = np.delete(_pot_new_insect,recorded_BS,0)

    _new_insects = dark_spot_check(_pot_new_insect)

    new_insect_mode(_new_insects)

    return _new_insects

def associate_detections_DL(_detections, _predictions, _max_dist_dl):
    _missing = []
    _assignments = Hungarian_method(_detections, _predictions)
    _insects = [i[0] for i in _predictions]

    _not_associated = np.zeros(shape=(0,5))
    for _nass in (_assignments[len(_insects):]):
        _not_associated = np.vstack([_not_associated,(_detections[_nass])])

    _associations_DL = np.zeros(shape=(0,6))
    for ass in np.arange(len(_insects)):
        _record = _assignments[ass]

        if (_record <= len(_detections)-1):
            _xc, _yc, _area, _lable, _conf = _detections[_assignments[ass]][0],_detections[_assignments[ass]][1],_detections[_assignments[ass]][2],_detections[_assignments[ass]][3],_detections[_assignments[ass]][4]
            _dist = cal_dist(_xc,_yc,_predictions[ass][1],_predictions[ass][2])
            if(_dist>_max_dist_dl) and not low_confident_ass(_detections, _predictions, max_dist_dl,_dist, False):
                _missing.append(_predictions[ass][0])
            else:
                _associations_DL = np.vstack([_associations_DL,(_predictions[ass][0],_detections[_assignments[ass]][0],_detections[_assignments[ass]][1],_detections[_assignments[ass]][2],_detections[_assignments[ass]][3],_detections[_assignments[ass]][4])])

        else:
            _missing.append(_predictions[ass][0])

    return _associations_DL, _missing, _not_associated

def cal_threshold_dist(_max_dist_dl,bs_mode):
    if bs_mode:
        threshold_dist = _max_dist_dl*1.5
    else:
        threshold_dist = _max_dist_dl*2

    return threshold_dist

def low_confident_ass(_detections, _predictions,_max_dist_dl,_dist,bs_mode):

    threshold_dist = cal_threshold_dist(_max_dist_dl,bs_mode)

    if (len(_detections) == len(_predictions)) and (_dist <= threshold_dist) and pt_cfg.POLYTRACK.NEW_INSECT_MODE:
        pt_cfg.POLYTRACK.NEW_INSECT_MODE = False
        can_associate = True
    else:
        can_associate = False

    return can_associate

#not associated = when there are more detections than predictions
#missing = when there are less detections than predictions. 1) less number of insects detected, 2) detected insects cannot be associated with a detection
def associate_detections_BS(_detections, _predictions):
    _assignments = Hungarian_method(_detections, _predictions)
    _missing = []
    _insects = [i[0] for i in _predictions]
    _not_associated = np.zeros(shape=(0,3))
    for _nass in (_assignments[len(_insects):]):
        _not_associated = np.vstack([_not_associated,(_detections[_nass])])

    _associations_BS = np.zeros(shape=(0,4))
    for ass in np.arange(len(_insects)):
        _record = _assignments[ass]
        if (_record <= len(_detections)-1):
            _xc, _yc, _area = _detections[_assignments[ass]][0],_detections[_assignments[ass]][1],_detections[_assignments[ass]][2]
            _dist = cal_dist(_xc,_yc,_predictions[ass][1],_predictions[ass][2])
            #print(len(_detections), len(_predictions))
            if(_dist>max_dist_bs) and not low_confident_ass(_detections, _predictions, max_dist_dl,_dist, True):
                _missing.append(_predictions[ass][0])
            else:
                _associations_BS = np.vstack([_associations_BS,(int(float(_predictions[ass][0])),int(_xc), int(_yc), int(_area))])
        else:
            _missing.append(_predictions[ass][0])

    return _associations_BS, _missing,_not_associated

def Hungarian_method(_detections, _predictions):
    num_detections, num_predictions = len(_detections), len(_predictions)
    mat_shape = max(num_detections, num_predictions)
    hun_matrix = np.full((mat_shape, mat_shape),0)
    for p in np.arange(num_predictions):
        for d in np.arange(num_detections):
            hun_matrix[p][d] = cal_dist(_predictions[p][1],_predictions[p][2],_detections[d][0],_detections[d][1])

    row_ind, col_ind = linear_sum_assignment(hun_matrix)

    return col_ind

def prepare_to_track(nframe, vid, idle, new_insect, video_start_frame):
    #print((new_insect), idle)

    if idle and (len(new_insect)>0):
        nframe = max((nframe - pt_cfg.POLYTRACK.BACKTRACK_FRAMES), video_start_frame)
        reset_frame = nframe - video_start_frame
        vid.set(1, reset_frame)
        idle = False
        new_insect = []
        pt_cfg.POLYTRACK.IDLE_OUT = True

    else:
        pass

    return nframe, idle, new_insect


import copy
import glob
import math
import os
import time
from collections import deque
from distutils.dir_util import copy_tree

import numpy as np
import torch

from a2c_ppo_acktr import algo
from a2c_ppo_acktr.arguments import get_args
from envs.envs import make_vec_envs, get_vec_normalize
from a2c_ppo_acktr.model import Policy
from a2c_ppo_acktr.storage import RolloutStorage
from a2c_ppo_acktr.utils import update_linear_schedule
from a2c_ppo_acktr.visualize import visdom_plot
from envs.pipelines import pipelines

args = get_args()
assert args.algo in ['a2c', 'ppo', 'acktr']
if args.recurrent_policy:
    assert args.algo in ['a2c', 'ppo'], \
        'Recurrent policy is not implemented for ACKTR'

use_metric = args.trg_succ_rate is not None
torch.manual_seed(args.seed)
torch.cuda.manual_seed_all(args.seed)

if args.cuda and torch.cuda.is_available() and args.cuda_deterministic:
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True

def main(env, scene_path):
    try:
        os.makedirs(args.log_dir)
    except OSError:
        files = glob.glob(os.path.join(args.log_dir, '*.monitor.csv'))
        for f in files:
            os.remove(f)
    save_path = os.path.join(args.save_dir, args.algo)

    eval_x = []
    eval_y = []

    torch.set_num_threads(1)
    device = torch.device("cuda:0" if args.cuda else "cpu")

    initial_policies = torch.load(os.path.join(args.load_dir, args.algo,
                                               args.initial_policy + ".pt")) \
        if args.initial_policy else None

    if args.reuse_residual:
        residual, ob_rms, initial_policies = initial_policies
    else:
        residual = None
        ob_rms = None

    pose_estimator = torch.load(os.path.join(args.load_dir, "pe",
                                             args.pose_estimator + ".pt")) \
        if args.pose_estimator else None

    envs = make_vec_envs(env, scene_path, args.seed, args.num_processes, args.gamma, args.log_dir,
                         device, False, initial_policies, pose_estimator=pose_estimator,
                         init_control=not args.dense_ip)
    if args.reuse_residual:
        vec_norm = get_vec_normalize(envs)
        if vec_norm is not None:
            vec_norm.eval()
            vec_norm.ob_rms = ob_rms

    base_kwargs = {'recurrent': args.recurrent_policy}
    base = residual.base if args.reuse_residual else None
    dist = residual.dist if args.reuse_residual else None
    actor_critic = Policy(envs.observation_space.shape, envs.action_space, base_kwargs=base_kwargs,
                          zero_last_layer=True, base=base, dist=dist)
    actor_critic.to(device)

    if args.algo == 'a2c':
        agent = algo.A2C_ACKTR(actor_critic, args.value_loss_coef, args.entropy_coef, lr=args.lr,
                               eps=args.eps, alpha=args.alpha, max_grad_norm=args.max_grad_norm)
    elif args.algo == 'ppo':
        agent = algo.PPO(actor_critic, args.clip_param, args.ppo_epoch, args.num_mini_batch,
                         args.value_loss_coef, args.entropy_coef, lr=args.lr, eps=args.eps,
                         max_grad_norm=args.max_grad_norm,
                         burn_in=initial_policies is not None and not args.reuse_residual)
    elif args.algo == 'acktr':
        agent = algo.A2C_ACKTR(actor_critic, args.value_loss_coef, args.entropy_coef, acktr=True)

    rollouts = RolloutStorage(args.num_steps, args.num_processes,
                              envs.observation_space.shape, envs.action_space,
                              actor_critic.recurrent_hidden_state_size)

    obs = envs.reset()
    rollouts.obs[0].copy_(obs)
    rollouts.to(device)

    episode_rewards = deque(maxlen=64)

    num_updates = int(args.num_env_steps) // args.num_steps // args.num_processes
    total_num_steps = 0
    j = 0
    max_succ = -1
    max_mean_rew = - math.inf
    mean_ep_rew = - math.inf
    evals_without_improv = 0

    start = time.time()
    start_update = start
    while (not use_metric and j < num_updates) or (use_metric and max_succ < args.trg_succ_rate):
        if args.eval_interval is not None and j % args.eval_interval == 0:
            print("Evaluating current policy...")
            i = 0
            total_successes = 0
            max_trials = 50
            eval_recurrent_hidden_states = torch.zeros(
                args.num_processes, actor_critic.recurrent_hidden_state_size, device=device)
            eval_masks = torch.zeros(args.num_processes, 1, device=device)
            while i + args.num_processes <= max_trials:

                with torch.no_grad():
                    _, action, _, eval_recurrent_hidden_states = actor_critic.act(
                        obs, eval_recurrent_hidden_states, eval_masks, deterministic=True)

                obs, _, dones, infos = envs.step(action)

                if np.all(dones):  # Rigid - assumes episodes are fixed length
                    rews = []
                    for info in infos:
                        rews.append(info['rew_success'])
                    i += args.num_processes
                    rew = sum([int(rew > 0) for rew in rews])
                    total_successes += rew

            p_succ = (100 * total_successes / i)
            eval_x += [total_num_steps]
            eval_y += [p_succ]

            end = time.time()
            print(f"Evaluation: {total_successes} successful out of {i} episodes - "
                  f"{p_succ:.2f}% successful. Eval length: {end - start_update}")
            torch.save([eval_x, eval_y], os.path.join(args.save_as + "_eval.pt"))
            start_update = end

            if p_succ > max_succ:
                max_succ = p_succ
                max_mean_rew = mean_ep_rew
                evals_without_improv = 0
            elif mean_ep_rew > max_mean_rew:
                print("Unimproved success rate, higher reward")
                max_mean_rew = mean_ep_rew
                evals_without_improv = 0
            else:
                evals_without_improv += 1

            if evals_without_improv == 10 or max_succ >= args.trg_succ_rate:
                save_model = actor_critic
                if args.cuda:
                    save_model = copy.deepcopy(actor_critic).cpu()

                save_model = [save_model, getattr(get_vec_normalize(envs), 'ob_rms', None),
                              initial_policies]
                extra = "_final" if evals_without_improv == 5 else ""
                torch.save(save_model, os.path.join(save_path, args.save_as + f"{extra}.pt"))
                break

        # save for every interval-th episode or for the last epoch
        if ((not use_metric and (j % args.save_interval == 0 or j == num_updates - 1))
                or (use_metric and evals_without_improv == 0)) and args.save_dir != "":
            os.makedirs(save_path, exist_ok=True)

            save_model = actor_critic
            if args.cuda:
                save_model = copy.deepcopy(actor_critic).cpu()

            if pose_estimator is not None:
                save_model = [save_model, pose_estimator, initial_policies]
            else:
                save_model = [save_model, getattr(get_vec_normalize(envs), 'ob_rms', None),
                              initial_policies]

            torch.save(save_model, os.path.join(save_path, args.save_as + ".pt"))
            # torch.save(save_model, os.path.join(save_path, args.save_as + f"{j * args.num_processes * args.num_steps}.pt"))

        if args.use_linear_lr_decay:
            # decrease learning rate linearly
            if args.algo == "acktr":
                # use optimizer's learning rate since it's hard-coded in kfac.py
                update_linear_schedule(agent.optimizer, j, num_updates, agent.optimizer.lr)
            else:
                update_linear_schedule(agent.optimizer, j, num_updates, args.lr)

        if args.algo == 'ppo' and args.use_linear_clip_decay:
            agent.clip_param = args.clip_param  * (1 - j / float(num_updates))

        for step in range(args.num_steps):
            # Sample actions
            with torch.no_grad():
                value, action, action_log_prob, recurrent_hidden_states = actor_critic.act(
                        rollouts.obs[step],
                        rollouts.recurrent_hidden_states[step],
                        rollouts.masks[step])

            # Obser reward and next obs
            obs, reward, done, infos = envs.step(action)

            for info in infos:
                if 'episode' in info.keys():
                    episode_rewards.append(info['episode']['r'])

            # If done then clean the history of observations.
            masks = torch.FloatTensor([[0.0] if done_ else [1.0]
                                       for done_ in done])
            rollouts.insert(obs, recurrent_hidden_states, action, action_log_prob, value, reward, masks)

        with torch.no_grad():
            next_value = actor_critic.get_value(rollouts.obs[-1],
                                                rollouts.recurrent_hidden_states[-1],
                                                rollouts.masks[-1]).detach()

        rollouts.compute_returns(next_value, args.use_gae, args.gamma, args.tau)

        value_loss, action_loss, dist_entropy = agent.update(rollouts)

        rollouts.after_update()

        total_num_steps = (j + 1) * args.num_processes * args.num_steps

        if j % args.log_interval == 0 and len(episode_rewards) > 1:
            mean_ep_rew = np.mean(episode_rewards)
            if mean_ep_rew > max_mean_rew:
                print("Improved max mean reward")
                max_mean_rew = mean_ep_rew
                evals_without_improv = 0
            end = time.time()
            print("Updates {}, num timesteps {}, FPS {} \n Last {} training episodes: mean/median reward {:.1f}/{:.1f}, min/max reward {:.1f}/{:.1f}".
                format(j, total_num_steps,
                       int(total_num_steps / (end - start)),
                       len(episode_rewards),
                       mean_ep_rew,
                       np.median(episode_rewards),
                       np.min(episode_rewards),
                       np.max(episode_rewards), dist_entropy,
                       value_loss, action_loss))
            print("Update length: ", end - start_update)
            start_update = end

        if args.vis and (j % args.vis_interval == 0 or (not use_metric and j == num_updates - 1)):
            try:
                # Sometimes monitor doesn't properly flush the outputs
                visdom_plot(args.log_dir, args.save_as, args.algo, total_num_steps)
            except IOError:
                pass

        j += 1

    if use_metric:
        if max_succ >= args.trg_succ_rate:
            print(f"Achieved greater than {args.trg_succ_rate}% success, advancing curriculum.")
        else:
            print(f"Policy converged with max success rate < {args.trg_succ_rate}%")
    # Copy logs to permanent location so new graphs can be drawn.
    copy_tree(args.log_dir, os.path.join('logs', args.save_as))
    envs.close()
    return total_num_steps

def train_with_metric(pipeline, train, save_base):
    if args.use_linear_clip_decay:
        raise ValueError("Cannot use clip decay with unbounded metric-based training length.")
    if args.eval_interval is None:
        raise ValueError("Need to set eval_interval to evaluate success rate")
    args.use_linear_lr_decay = False

    training_lengths = []
    save_path = os.path.join(args.save_dir, args.algo)
    target = args.trg_succ_rate
    for i in range(0, args.num_seeds * 16, 16):
        args.seed = i
        args.reuse_residual = False
        args.trg_succ_rate = target
        args.initial_policy = None
        print(f"Using seed {args.seed}")
        training_lengths += [train(pipeline, f"{save_base}_{i}")]
        torch.save(training_lengths,
                   os.path.join(save_path, f"{save_base}_{args.pipeline}_train_lengths.pt"))

    print(training_lengths)
    total_train_times = [sum(lengths) for lengths in training_lengths]
    print(total_train_times)
    torch.save([total_train_times, training_lengths],
               os.path.join(save_path, f"{save_base}_{args.pipeline}_train_lengths.pt"))
    return total_train_times, training_lengths

def execute_curriculum(pipeline, save_base):
    training_lengths = []
    criteria_string = f"until {args.trg_succ_rate}% successful" if use_metric \
        else f"for {args.num_env_steps} timesteps"
    for scene in pipeline['curriculum']:
        print(f"Training {scene} {criteria_string}")
        args.save_as = f'{save_base}_{scene}'
        training_lengths += [main(pipeline['sparse'], scene)]
        args.reuse_residual = True
        args.initial_policy = args.save_as
    scene = pipeline['task']
    print(f"Training on {scene} full task")
    args.save_as = f'{save_base}_{scene}'
    args.trg_succ_rate = 101  # Does not affect fixed length curriculum
    training_lengths += [main(pipeline['sparse'], scene)]
    return training_lengths

def train_baseline(pipeline, save_base):
    training_lengths = []
    scene = pipeline['task']
    print(f"Training {scene} until {args.trg_succ_rate}% successful with dense rewards")
    args.save_as = f'{save_base}_dense_{scene}'
    training_lengths += [main(pipeline['dense'], scene)]
    args.initial_policy = args.save_as
    print(f"Training on {scene} until convergence")
    args.save_as = f'{save_base}_sparse_{scene}'
    args.trg_succ_rate = 101
    training_lengths += [main(pipeline['sparse'], scene)]
    return training_lengths

if __name__ == "__main__":
    if args.scene_name is not None:
        main(pipelines[args.pipeline]['sparse'], args.scene_name)  # TODO
    elif args.dense_ip:
        train_with_metric(pipelines[args.pipeline], train_baseline, args.save_as)
    elif use_metric:
        train_with_metric(pipelines[args.pipeline], execute_curriculum, args.save_as)
    else:
        execute_curriculum(pipelines[args.pipeline], args.save_as)