"""Objects representing pure chemical components and pseudo-components.

Attributes
----------
COMP_FAM : list
    Pre-defined component family pick list based on definitions in _[1] and _[2].

Notes
-----

References
----------
[1] Perry's Chemical Engineers' Handbook; Perry, R. H., Southard, M. Z., Eds.; McGraw-Hill Education: New York, 2019.
[2] Tihic, A.; Kontogeorgis, G. M.; von Solms, N.;Michelsen, M. L. Applications of the simplified perturbed-chain SAFT
equation of state using an extended parameter table. Fluid Phase Equilib. 2006, 248, 29-43.
"""

COMP_FAM = ['Alkanes', 'Alkenes', 'Alkynes', 'Cycloalkanes' 'Aromatics', 'Polynuclear Aromatics', 'Aldehydes',
            'Ketones', 'Heterocyclics', 'Elements', 'Alcohols', 'Phenols', 'Ethers', 'Acids', 'Esters', 'Amines',
            'Amides', 'Nitriles', 'Nitro Compounds', 'Isocyanates', 'Mercaptans', 'Sulfides',
            'Halogenated Hydrocarbons', 'Silanes', 'Inorganics', 'Multifunctional']

import numpy as np
from utilities import *

class Comp(object):
    """A pure chemical component."""
    def __init__(self, name=None):
        """
        Parameters
        ----------
        name : str
        """
        # General Attributes
        if name is None:
            raise ValueError("Comp object must be initialized with a name and molecular weight.")
        else:
            # Metadata and constants.
            self.name = name
            self.cas_no = None
            self.formula = None
            self.family = None
            self.mw = None
            self.vdwv = None
            self.vdwa = None
            self.rgyr = None
            self.dipole = None
            self.quadrupole = None
            self.acentric = None
            self.tc = None
            self.pc = None
            self.vc = None
            self.rhoc = None
            self.tt = None
            self.pt = None
            self.bp = None
            self.mp = None
            self.hfus = None
            self.hsub = None
            self.ig_hform = None
            self.ig_gform = None
            self.ig_entr = None
            self.hcomb = None

            # Temperature dependent properties.
            self.pvap_l = None
            self.hvap_l = None
            self.den_s = None
            self.den_l = None
            self.beta_s = None
            self.beta_l = None
            self.cp_s = None
            self.cp_l = None
            self.cp_ig = None
            self.visc_l = None
            self.visc_ig = None
            self.tcond_s = None
            self.tcond_l = None
            self.tcond_ig = None
            self.sigma = None

            # TODO: Ensure CEOS physical terms are built with getter/setter checks for CEOS objects.
            # Cubic EOS physical parameter dictionaries.
            self.srk_phys = {'a': None, 'b': None, 'm': None}
            self.cpa_phys = {'a0': None, 'b': None, 'c1': None}
            self.pr_phys = {'a': None, 'b': None, 'm': None}
            self.gpr_phys = {'a': None, 'b': None, 'a': None, 'b': None, 'c': None}
            self.tpr_phys = {'a': None, 'b': None, 'l': None, 'm': None, 'n': None}

            # SAFT EOS physical parameter dictionaries.
            self.spc_saft_phys = {'m': None, 'sig': None, 'eps': None}
            self.pc_saft_phys = {'m': None, 'sig': None, 'eps': None}

            # SAFT EOS association parameter objects.
            self.assoc_sites = None
            self.cpa_assoc = None
            self.spc_saft_assoc = None
            self.pc_saft_assoc = None

    @property
    def name(self):
        """str : Name of chemical compound.

        Value can only be set whe creating a new Comp instance.
        """
        return self._name

    @name.setter
    def name(self, value):
        try:
            self._name
        except AttributeError:
            if isinstance(value, str):
                self._name = value
            else:
                raise TypeError("name must be a string.")

    @property
    def cas_no(self):
        """str: Chemical Abstracts Service Registry Number."""
        return self._cas_no

    @cas_no.setter
    def cas_no(self, value):
        if isinstance(value, str) or value is None:
            self._cas_no = value
        else:
            raise TypeError("cas_no must be a string.")

    @property
    def formula(self):
        """str : Chemical formula."""
        return self._formula

    @formula.setter
    def formula(self, value):
        if isinstance(value, str) or value is None:
            self._formula = value
        else:
            raise TypeError("formula must be a string.")

    @property
    def family(self):
        """str : Chemical family."""
        return self._family

    @family.setter
    def family(self, value):
        if isinstance(value, str) or value is None:
            if value in COMP_FAM or value is None:
                self._family = value
            else:
                raise ValueError("family must be one of the pre-defined values.")
        else:
            raise TypeError("family must be a string.")

    @property
    def mw(self):
        """float : Molecular weight, g/mol."""
        return self._mw

    @mw.setter
    def mw(self, value):
        if (isinstance(value, float) and value >= 0.0) or value is None:
            self._mw = value
        else:
            raise TypeError("mw must be a positive float.")

    @property
    def vdwv(self):
        """float : Van der Waal's volume, unit TBD."""
        return self._vdwv

    @vdwv.setter
    def vdwv(self, value):
        if (isinstance(value, float) and value >= 0.0) or value is None:
            self._vdwv = value
        else:
            raise TypeError("vdwv must be a positive float.")

    @property
    def vdwa(self):
        """float : Van der Waal's surface area, unit TBD."""
        return self._vdwa

    @vdwa.setter
    def vdwa(self, value):
        if (isinstance(value, float) and value >= 0.0) or value is None:
            self._vdwa = value
        else:
            raise TypeError("vdwa must be a positive float.")

    @property
    def rgyr(self):
        """float : Radius of gyration, unit TBD."""
        return self._rgyr

    @rgyr.setter
    def rgyr(self, value):
        if (isinstance(value, float) and value >= 0.0) or value is None:
            self._rgyr = value
        else:
            raise TypeError("rgyr must be a positive float.")

    @property
    def dipole(self):
        """float : Gas phase dipole moment, unit TBD."""
        return self._dipole

    @dipole.setter
    def dipole(self, value):
        if (isinstance(value, float) and value >= 0.0) or value is None:
            self._dipole = value
        else:
            raise TypeError("dipole must be a positive float.")

    @property
    def quadrupole(self):
        """float : Gas phase quadrupole moment, unit TBD."""
        return self._quadrupole

    @quadrupole.setter
    def quadrupole(self, value):
        if (isinstance(value, float) and value >= 0.0) or value is None:
            self._quadrupole = value
        else:
            raise TypeError("quadrupole must be a positive float.")

    @property
    def acentric(self):
        """float : Acentric factor, dimensionless."""
        return self._acentric

    @acentric.setter
    def acentric(self, value):
        if isinstance(value, float) or value is None:
            self._acentric = value
        else:
            raise TypeError("acentric must be a float.")

    @property
    def tc(self):
        """float : Critical temperature, K."""
        return self._tc

    @tc.setter
    def tc(self, value):
        if (isinstance(value, float) and value >= 0.0) or value is None:
            self._tc = value
        else:
            raise TypeError("tc must be a positive float.")

    @property
    def pc(self):
        """float : Critical pressure, Pa."""
        return self._pc

    @pc.setter
    def pc(self, value):
        if (isinstance(value, float) and value >= 0.0) or value is None:
            self._pc = value
        else:
            raise TypeError("pc must be a positive float.")

    @property
    def vc(self):
        """float : Critical volume, m**3/mol."""
        return self._vc

    @vc.setter
    def vc(self, value):
        if (isinstance(value, float) and value >= 0.0) or value is None:
            self._vc = value
        else:
            raise TypeError("vc must be a positive float.")

    @property
    def rhoc(self):
        """float : Critical density, mol/m**3"""
        return self._rhoc

    @rhoc.setter
    def rhoc(self, value):
        if (isinstance(value, float) and value >= 0.0) or value is None:
            self._rhoc = value
        else:
            raise TypeError("rhoc must be a positive float.")

    @property
    def tt(self):
        """float : Triple point temperature, K."""
        return self._tt

    @tt.setter
    def tt(self, value):
        if (isinstance(value, float) and value >= 0.0) or value is None:
            self._tt = value
        else:
            raise TypeError("tt must be a positive float.")

    @property
    def pt(self):
        """float : Triple point pressure, Pa."""
        return self._pt

    @pt.setter
    def pt(self, value):
        if (isinstance(value, float) and value >= 0.0) or value is None:
            self._pt = value
        else:
            raise TypeError("pt must be a positive float.")

    @property
    def bp(self):
        """float : Boiling point, K."""
        return self._bp

    @bp.setter
    def bp(self, value):
        if (isinstance(value, float) and value >= 0.0) or value is None:
            self._bp = value
        else:
            raise TypeError("bp must be a positive float.")

    @property
    def mp(self):
        """float : Melting point, K."""
        return self._mp

    @mp.setter
    def mp(self, value):
        if (isinstance(value, float) and value >= 0.0) or value is None:
            self._mp = value
        else:
            raise TypeError("mp must be a positive float.")

    @property
    def hfus(self):
        """float : Enthalpy of solid-liquid fusion, J/mol."""
        return self._hfus

    @hfus.setter
    def hfus(self, value):
        if isinstance(value, float) or value is None:
            self._hfus = value
        else:
            raise TypeError("hfus must be a float.")

    @property
    def hsub(self):
        """float : Enthalpy of sublimation, J/mol."""
        return self._hsub

    @hsub.setter
    def hsub(self, value):
        if isinstance(value, float) or value is None:
            self._hsub = value
        else:
            raise TypeError("hsub must be a float.")

    @property
    def ig_hform(self):
        """float : Ideal gas enthalpy of formation at 298.15, unit TBD.

        The compounds are considered to be formed from the elements in their standard states at 298.15K and 1 bar.
        These include  C (graphite) and S (rhombic).
        """
        return self._ig_hform

    @ig_hform.setter
    def ig_hform(self, value):
        if isinstance(value, float) or value is None:
            self._ig_hform = value
        else:
            raise TypeError("ig_hform must be a float.")

    @property
    def ig_gform(self):
        """float : Ideal gas gibbs energy of formation at 298.15, unit TBD."""
        return self._ig_gform

    @ig_gform.setter
    def ig_gform(self, value):
        if isinstance(value, float) or value is None:
            self._ig_gform = value
        else:
            raise TypeError("ig_gform must be a float.")

    @property
    def ig_entr(self):
        """float : Ideal gas gibbs entropy, unit TBD."""
        return self._ig_entr

    @ig_entr.setter
    def ig_entr(self, value):
        if isinstance(value, float) or value is None:
            self._ig_entr = value
        else:
            raise TypeError("ig_entr must be a float.")

    @property
    def hcomb(self):
        """float : Standard net enthalpy of combustion, unit TBD.

        Enthalpy of combustion is the net value for the compound in its standard state at 298.15K and 1 bar.  Products
        of combustion are taken to be CO2 (gas), H2O (gas), F2 (gas), Cl2 (gas), Br2 (gas), I2 (gas), SO2 (gas), N2
        (gas), P4O10 (crystalline), SiO2 (crystobalite), and Al2O3 (crystal, alpha).
        """
        return self._hcomb

    @hcomb.setter
    def hcomb(self, value):
        if isinstance(value, float) or value is None:
            self._hcomb = value
        else:
            raise TypeError("hcomb must be a float.")

    @property
    def pvap_l(self):
        """float : Liquid vapor pressure, Pa."""
        return self._pvap_l

    @pvap_l.setter
    def pvap_l(self, value):
        if isinstance(value, Corel) or value is None:
            self._pvap_l = value
        else:
            raise TypeError("pvap_l must be an instance of Corel.")

    @property
    def hvap_l(self):
        """float : Enthalpy of saturated liquid vaporization, J/mol."""
        return self._hvap_l

    @hvap_l.setter
    def hvap_l(self, value):
        if isinstance(value, Corel) or value is None:
            self._hvap_l = value
        else:
            raise TypeError("hvap_l must be an instance of Corel.")

    @property
    def den_s(self):
        """float : Solid density, kg/m3 or mol/m3."""
        return self._den_s

    @den_s.setter
    def den_s(self, value):
        if isinstance(value, Corel) or value is None:
            self._den_s = value
        else:
            raise TypeError("den_S must be an instance of Corel.")

    @property
    def den_l(self):
        """float : Liquid density, kg/m3 or mol/m3."""
        return self._den_l

    @den_l.setter
    def den_l(self, value):
        if isinstance(value, Corel) or value is None:
            self._den_l = value
        else:
            raise TypeError("den_l must be an instance of Corel.")

    @property
    def beta_s(self):
        """float : Isothermal solid compressibility, Unit TBD.

        beta = -(1/V)*(dV/dP) = (1/rho)*(drho/dP), sign change due to conversion from V to rho with chain rule."""
        return self._beta_s

    @beta_s.setter
    def beta_s(self, value):
        if isinstance(value, Corel) or value is None:
            self._beta_s = value
        else:
            raise TypeError("beta_s must be an instance of Corel.")

    @property
    def beta_l(self):
        """float : Isothermal liquid compressibility, Unit TBD.

        beta = -(1/V)*(dV/dP) = (1/rho)*(drho/dP), sign change due to conversion from V to rho with chain rule."""
        return self._beta_l

    @beta_l.setter
    def beta_l(self, value):
        if isinstance(value, Corel) or value is None:
            self._beta_l = value
        else:
            raise TypeError("beta_l must be an instance of Corel.")

    @property
    def cp_s(self):
        """float : Solid heat capacity, J/mol.K."""
        return self._cp_s

    @cp_s.setter
    def cp_s(self, value):
        if isinstance(value, Corel) or value is None:
            self._cp_s = value
        else:
            raise TypeError("cp_s must be an instance of Corel.")

    @property
    def cp_l(self):
        """float : Saturated liquid heat capacity, J/mol.K."""
        return self._cp_l

    @cp_l.setter
    def cp_l(self, value):
        if isinstance(value, Corel) or value is None:
            self._cp_l = value
        else:
            raise TypeError("cp_l must be an instance of Corel.")

    @property
    def cp_ig(self):
        """float : Ideal gas heat capacity, J/mol.K."""
        return self._cp_ig

    @cp_ig.setter
    def cp_ig(self, value):
        if isinstance(value, Corel) or value is None:
            self._cp_ig = value
        else:
            raise TypeError("cp_ig must be an instance of Corel.")

    @property
    def visc_l(self):
        """float : Saturated liquid viscosity, unit TBD."""
        return self._visc_l

    @visc_l.setter
    def visc_l(self, value):
        if isinstance(value, Corel) or value is None:
            self._visc_l = value
        else:
            raise TypeError("visc_l must be an instance of Corel.")

    @property
    def visc_ig(self):
        """float : Ideal gas viscosity, unit TBD."""
        return self._visc_ig

    @visc_ig.setter
    def visc_ig(self, value):
        if isinstance(value, Corel) or value is None:
            self._visc_ig = value
        else:
            raise TypeError("visc_ig must be an instance of Corel.")

    @property
    def tcond_s(self):
        """float : Solid thermal conductivity, unit TBD."""

    @tcond_s.setter
    def tcond_s(self, value):
        if isinstance(value, Corel) or value is None:
            self._tcond_s = value
        else:
            raise TypeError("tcond_s must be an instance of Corel.")

    @property
    def tcond_l(self):
        """float : Saturated liquid thermal conductivity, unit TBD."""
        return self._tcond_l

    @tcond_l.setter
    def tcond_l(self, value):
        if isinstance(value, Corel) or value is None:
            self._tcond_l = value
        else:
            raise TypeError("tcond_l must be an instance of Corel.")

    @property
    def tcond_ig(self):
        """float : Ideal gas thermal conductivity, unit TBD."""
        return self._tcond_ig

    @tcond_ig.setter
    def tcond_ig(self, value):
        if isinstance(value, Corel) or value is None:
            self._tcond_ig = value
        else:
            raise TypeError("tcond_ig must be an instance of Corel.")

    @property
    def sigma(self):
        """float : Surface tension, unit TBD."""
        return self._sigma

    @sigma.setter
    def sigma(self, value):
        if isinstance(value, Corel) or value is None:
            self._sigma = value
        else:
            raise TypeError("sigma must be an instance of Corel.")

    def k_wilson(self, p=None, t=None):
        """Wilson's equilibrium ratio (Ki = yi / xi).

        Parameters
        ----------
        p : float
            Pressure, Pa.
        t : float
            Temperature, K.

        Returns
        -------
        float
            Wilson's K-factor.
        """
        if isinstance(p, float) and isinstance(t, float) and p > 0.0 and t > 0.0:
            return (self._pc / p) * np.exp(5.37 * (1.0 + self._acentric) * (1.0 - self._tc / t))
        else:
            raise RuntimeError("p and t must both be positive floats.")

    def melting_curve(self, t=None):
        """Melting curve for a pure solid.

        Melting pressure can be estimated with the Clausius-Clapeyron equation. A very simple approach assumes that
        changes in the enthalpy of fusion, liquid molar volume, and solid molar volume are negligible over the
        temperature range of interest.  With these assumptions, the equation of the melting curve is as follows:

        dp/dt = delta_h / (t * delta_v)

        p - p0 = (delta_h / delta_v) * ln(t / t0)

        Variation of the enthalpy of fusion with temperature is related to differences in heat capacity between phases.

        ddelta_h/dt = dh_l/dt - dh_s/dt = cp_l - cp_s = detla_cp

        delta_h = integral_from_t0_to_t(delta_cp * dt) + delta_h0

        Temperature dependence of the heat capacity difference is usually weak and can be neglected.

        delta_h = delta_cp * (t - t0) + delta_h0

        Liquid and solid heat capacities can be evaluated at the melting temperature. These values are be derived from
        cp_l(t) and cp_s(t) correlations OR it can be entered as a constant (a useful adjustable parameter for
        correlating melting curves, sublimation curves, and solubility. The equation of the melting curve is as follows:

        p - p0 = (delta_cp / delta_v) * (t - t0) + ((delta_h - t0 * delta_cp) / delta_v) * ln(t / t0)

        Parameters
        ----------
        t : float
            Temperature specification, K.

        Returns
        -------
        float
            Melting pressure, Pa.

        References
        ----------
        [1] Tosun, I. The thermodynamics of phase and reaction equilibria, 1st ed., Elsevier: Amsterdam, 2013.
        [2] Poling, B. E.; Praunitz, J. M.; O'Connell, J. P. The properties of gases and liquids, 5th ed.,
        McGraw-Hill, 2000.
        [3] Pappa, G. D.; Voutsas, E. C.; Magoulas, K.; Tassios, D. P. Estimation of the differential molar heat
        capacities of organic compounds at their melting point. Ind. Eng. Chem. Res. 2005, 44, 3799–3806.
        """
        if isinstance(t, float) and t > 0.0:
            if self._hfus and self._mp and self._den_l and self._den_s:
                p0 = 101325.0
                t0 = self._mp
                tavg = (t + t0) / 2.0
                v_l = 1.0 / self._den_l(tavg)
                v_s = 1.0 / self._den_s(tavg)
                delta_v = v_l - v_s
                if self._cp_l and self._cp_s:
                    delta_cp = self._cp_l(tavg) - self._cp_s(tavg)
                    return p0 + (delta_cp/delta_v) * (t - t0) + ((self._hfus - t0 * delta_cp)/delta_v) * np.log(t/t0)
                else:
                    return p0 + (self._hfus/delta_v) * np.log(t/t0)
            else:
                raise RuntimeError("Enthalpy of fusion at melting point, liquid density, and solid density required.")
        else:
            raise RuntimeError("t must be a positive float.")

    def sublimation_curve(self):
        """Sublimation curve for a pure solid.

        Sublimation pressure can be estimated with the Clausius-Clapeyron equation. A very simple approach assumes that
        changes in the enthalpy of sublimation and solid molar volume are negligible over the temperature range of
        interest.  With these assumptions, the equation of the melting curve is as follows:

        dp/dt = delta_h / (t * delta_v)

        delta_h_sub = delta_h_vap + delta_h_fusion

        delta_v = v_v - v_s ~ v_v = R*t/p_sub

        p - p0 = (delta_h / delta_v) * ln(t / t0)

        Parameters
        ----------
        t : float
            Temperature specification, K.

        Returns
        -------
        float
            Sublimation pressure, Pa.

        References
        ----------
        [1] Tosun, I. The thermodynamics of phase and reaction equilibria, 1st ed., Elsevier: Amsterdam, 2013.
        [2] Poling, B. E.; Praunitz, J. M.; O'Connell, J. P. The properties of gases and liquids, 5th ed.,
        McGraw-Hill, 2000.
        [3] Pappa, G. D.; Voutsas, E. C.; Magoulas, K.; Tassios, D. P. Estimation of the differential molar heat
        capacities of organic compounds at their melting point. Ind. Eng. Chem. Res. 2005, 44, 3799–3806.
        """
        return

    def density(self, t=None, p=None, phase='l', spec='molar'):
        """Pure component liquid and/or solid density estimation.

        This function evaluates liquid or solid densities at pressures deviating from the vapor pressure curve or
        melting curve. The first step is calculating an initial density along the vapor pressure curve or melting curve.
        The next step is to correct for pressure with the isothermal compressibility:

        rho = rho_0 * exp(beta * (p - p_0))

        In this expression, rho_0 and p_0 are the saturation or melting density.  If pressure is not specified, then the
        saturated liquid or melting point solid density is returned. IF

        Parameters
        ----------
        p : float
            Pressure, Pa.
        t : float
            Temperature, K.
        spec : str
            Phase specification (either 'mass' or 'molar').

        Returns
        -------
        float or None
            Density of saturated liquid or melting point solid. None returned if 't' is outside correlation temp limits.
        float or None
            Density of solid phase at temperature 't' (None if no solid density correlation available).
        str
            Unit ('kg/m3' or 'mol/m3').
        """
        if not isinstance(p, (float, None)):
            raise TypeError("p must be None or a positive float.")
        elif not isinstance(t, float) and p > 0.0:
            raise TypeError("t must be a positive float.")
        elif phase not in ['l', 's']:
            raise ValueError("phase must be either 'l' or 's'.")
        elif spec not in ['mass', 'molar']:
            raise ValueError("spec must be either 'molar' or 'mass'.")
        else:
            if self._beta_l and p:
                p_corr_l = np.exp(self._beta_l(t) * ())
            else:
                beta_l = 0.0

            if self._beta_s:
                beta_s = self._beta_s(t)
            else:
                beta_s = 0.0

        if self._den_l and not self._den_s:
            # Liquid density correlation exists, solid density correlation does not.
            if spec == 'molar':
                if self._den_l.unit == 'kg/m3':
                    return self._den_l(t) * 1000.0 / self._mw
                elif self._den_l.unit == 'mol/m3':
                    return self._den_l(t)
                else:
                    raise RuntimeError("liquid density correlation does not have a valid unit.")
            elif spec == 'mass':
                return
            else:
                raise ValueError("spec is not valid.")
        elif not self._den_l and self._den_s:
            # Liquid density correlation does not exist, solid density correlation exists.
            if spec == 'molar':
                return
            elif spec == 'mass':
                return
            else:
                raise ValueError("spec is not valid.")
        elif self._den_l and self._den_s:
            # Liquid density and solid density correlations exist.
            if spec == 'molar':
                return
            elif spec == 'mass':
                return
            else:
                raise ValueError("spec is not valid.")
        else:
            raise RuntimeError("liquid and solid density correlations not loaded for Comp.")

    # TODO: Improve interface by making these checks part of getter/setter methods?
    def _check_assoc_sites(self):
        # Check to ensure there are no duplicate association sites.
        if self.assoc_sites is not None:
            if not isinstance(self.assoc_sites, (list, tuple)):
                raise TypeError("Comp objects must be a list of unique assoc_site objects.")
            elif len(self.assoc_sites) != len(set(self.assoc_sites)):
                raise ValueError("Comp objects must be a list of unique assoc_site objects.")
            else:
                pass

    def _check_assoc_parameters(self):
        # Check to ensure association parameter are lists of AssocSiteInter objects.
        if self.assoc_sites is not None:
            if self.cpa_assoc is not None:
                for asi in self.cpa_assoc:
                    if not isinstance(asi, AssocSiteInter):
                        raise TypeError("Association parameter lists must contain AssocSiteInter objects.")

    def __eq__(self, other):
        if isinstance(other, Comp):
            name_eq = self.name == other.name
            return name_eq
        return False

    def __ne__(self, other):
        return not self.__eq__(other)

    def __hash__(self):
        return hash(self.name)

    def __str__(self):
        return "Name: {}, CAS No.: {}".format(self.name, self.cas_no)

class PseudoComp(object):
    """A pseudo-component (polymer, distillation cut, asphaltene, etc.).

    Notes
    -----
    # TODO: Improve agreement bewteeen Comp and PseudoComp classes. Currently a skeleton implementation.
    # TODO: Add Riazi correlations to estimate Tc, Pc, Vc, Rhoc from mw, sg, NBP.
    """

    def __init__(self, name):
        """
        Parameters
        ----------
        name : str
        """
        # General Attributes
        self.name = name
        self.mw = None
        self.sg = None
        self.nbp = None
        self.acentric = None
        self.tc = None
        self.pc = None
        self.vc = None
        self.rhoc = None

        # EOS Physical Attributes
        self.srk_phys = None
        self.pr_phys = None
        self.gpr_phys = None
        self.cpa_phys = None
        self.spc_saft_phys = None

        # EOS Association Attributes
        self.assoc_sites = None
        self.cpa_assoc = None
        self.spc_saft_assoc = None

    @property
    def name(self):
        """str : Name of pseudo-component.

        Value can only be set whe creating a new Comp instance.
        """
        return self._name

    @name.setter
    def name(self, value):
        try:
            self._name
        except AttributeError:
            self._name = value

    def __eq__(self, other):
        if isinstance(other, PseudoComp):
            name_eq = self.name == other.name
            return name_eq
        return False

    def __ne__(self, other):
        return not self.__eq__(other)

    def __hash__(self):
        return hash(self.name)

    def __str__(self):
        return "Name: {}, MW: {}, SG: {}, NBP: {}".format(self.name, self.mw, self.sg, self.nbp)

class CompSet(object):
    """Set of components or pseudo-components.

    Notes
    -----
    """
    def __init__(self, comps=None):
        """
        Parameters
        ----------
        comps : list or tuple of Comp objects
        """
        if comps is None:
            raise ValueError("comps must be provided to create an instance of CompSet.")
        else:
            self.comps = comps

    @property
    def comps(self):
        """list or tuple of Comp objects : A collection of Comp objects.

        Values can only be set whe creating a new CompSet instance.
        """
        return self._comps

    @comps.setter
    def comps(self, value):
        try:
            self._comps
        except AttributeError:
            if isinstance(value, (list, tuple)):
                if len(value) == 0:
                    raise ValueError("comps must contain at least one Comp object.")
                elif len(value) != len(set(value)):
                    raise ValueError("comps cannot contain duplicate Comp objects.")
                elif all(isinstance(item, (Comp, PseudoComp)) for item in value):
                    self._comps = value
                else:
                    raise TypeError("comps must contain only Comp or PseudoComp objects.")
            else:
                raise TypeError("comps must be a list or tuple.")

    @property
    def size(self):
        """int : The number of Comp or PseudoComp objects in 'comps'."""
        return len(self._comps)

    @property
    def can_associate(self):
        """list of bool : Boolean indicating if Comp or PseudoComp objects in 'comps' can associate."""
        result = []
        for comp in self._comps:
            if comp.assoc_sites is not None:
                result.append(True)
            else:
                result.append(False)
        return result

    @property
    def mw(self):
        """list of float or None : Molecular weight for each Comp or PseudoComp objects in 'comps'.

        Returns None if any Comp or PseudoComp object is missing molecular weight."""
        result = []
        for comp in self._comps:
            if comp.mw is None:
                return None
            else:
                result.append(comp.mw)
        return np.array(result)

    def __eq__(self, other):
        if isinstance(other, CompSet):
            return self.comps == other.comps
        return False

    def __ne__(self, other):
        return not self.__eq__(other)

    def __hash__(self):
        return hash(self._comps)

    def __str__(self):
        result = []
        for comp in self._comps:
            result.append(comp.name)
        return ",".join(tuple(result))


#!/usr/bin/env python
# coding=utf-8
'''

yxhu@NPU-ASLP in Sogou inc.

'''

import numpy as np
import scipy
import torch
import random
import torch.nn as nn
from torch.utils.data import Dataset
import torch.utils.data as tud
import os
import sys
sys.path.append(os.path.dirname(sys.path[0])+'/tools/speech_processing_toolbox/')
sys.path.append('/home/work_nfs3/yxhu/workspace/se-cldnn-torch/tools/speech_processing_toolbox/')
sys.path.append(
    os.path.dirname(__file__) + '/tools/speech_processing_toolbox/')
sys.path.append(
    os.path.dirname(__file__))
import soundfile as sf
import voicetool.base as voicebox
from misc import read_and_config_file
import multiprocessing as mp
from Augmentation import  MixSpeaker, AddNoise
from drc import drc

class DataReader(object):
    def __init__(self, file_name,adaptation_wavs_dir, win_len=400, win_inc=100,left_context=0,right_context=0, fft_len=512, window_type='hamming', target_mode='MSA',sample_rate=16000):
        self.left_context = left_context
        self.right_context = right_context
        self.win_len = win_len
        self.win_inc = win_inc
        self.fft_len = fft_len
        self.window = {
                        'hamming':np.hamming(self.win_len)/1.2607934,
                        'hanning':np.hanning(self.win_len),
                        'none':np.ones(self.win_len)
                      }[window_type]
        self.file_list = read_and_config_file(file_name, decode=True)
        self.label_type = target_mode
        self.adaptation_wavs_dir = adaptation_wavs_dir

    def extract_feature(self, path):

        # process input feature
        wave_path = path['inputs']
        utt_id = wave_path.split('/')[-1]
        data = voicebox.audioread(wave_path)

        # Get embedding
        emb_wave_path = path['emb']
        # just for Aishell-1
        if not os.path.isfile(emb_wave_path):
            spk_dir = emb_wave_path[6:11]
            emb_wave_path = os.path.join(self.adaptation_wavs_dir, spk_dir, emb_wave_path+'.wav')

        inputs = data[None,:].astype(np.float32)
        emb = voicebox.audioread(emb_wave_path).astype(np.float32)
        #emb = drc(emb, 16000, 15,0.03,0.98)
        emb = np.reshape(emb, [1, -1])
        nsamples = data.shape[0]
        return inputs, emb, np.angle(inputs), utt_id, nsamples

    def __len__(self):
        return len(self.file_list)

    def __getitem__(self, index):
        return self.extract_feature(self.file_list[index])

class FixDataset(Dataset):

    def __init__(
            self,
            scp_file_name,
            emb_wav_path,
            segement_length=8,
            sample_rate=16000,
            gender2spk=None
        ):
        '''
            scp_file_name: the list include:[input_wave_path, output_wave_path, duration]
            spk_emb_scp: a speaker embedding ark's scp
            segement_length: to clip data in a fix length segment, default: 4s
            sample_rate: the sample rate of wav, default: 16000
            processer: a processer class to handle wave data
            gender2spk: a list include gender2spk, default: None
        '''
        self.wav_list = read_and_config_file(scp_file_name)
        mgr = mp.Manager()
        self.index =mgr.list()#[d for b in buckets for d in b]
        self.segement_length = segement_length * sample_rate
        self.emb_wave_path = emb_wav_path
        _dochunk(self.wav_list, self.index, self.segement_length)

    def __len__(self):
        return len(self.index)

    def __getitem__(self, index):
        data_info, start_time = self.index[index]
        emb_wave_path = find_emb_path(data_info['emb'], self.emb_wave_path)
        inputs = clip_data(data_info['path']['inputs'], start_time, self.segement_length)
        labels = clip_data(data_info['path']['labels'], start_time, self.segement_length)
        emb = clip_data(emb_wave_path, -2, self.segement_length)
        return inputs, labels, emb

def clip_data(path, start, segement_length):
    data, fs = sf.read(path)
    if start == -2:
        if data.shape[0] > segement_length:
            start = 0
        else:
            start = -1
    if start == -1:
        data = np.concatenate([data, data[:segement_length-data.shape[0]]])
    else:
        data = data[start:start+segement_length]
    if data.shape[0] < segement_length:
        data = np.pad(data, [0,segement_length-data.shape[0]])
    return data

def find_emb_path(cond, emb_wave_path):
    idx = 0
    while True:
        # Attention and Warning !!!!!!!!
        # This is just for aishell-1
        # Please make sure the adaptation wav's path !!!
        # start
        emb_name = random.sample(cond,1)[0]
        if not os.path.isfile(emb_name):
            spk_dir = emb_name[6:11]
            if spk_dir[0] !='S':
                spk_dir = emb_name.split('_')[0]
            emb_path = emb_wave_path+'/'+spk_dir+'/'+emb_name+".wav"
        else:
            emb_path = emb_name
        idx+=1
        data, fs = sf.read(emb_path)
        #end
        if data.shape[0] > 3*fs:
        # chose not too short
            break
        if idx >=100:
            raise RuntimeError('can not find a useable embedding!!!')
    return emb_path

class AutoDataset(Dataset):

    def __init__(
            self,
            target_scp,
            infer_scp,
            noise_scp,
            segement_length=10,
            repeate=4,
            sample_rate=16000,
            gender2spk=None
        ):
        '''
            scp_file_name: the list include:[input_wave_path, output_wave_path, duration]
            spk_emb_scp: a speaker embedding ark's scp
            segement_length: to clip data in a fix length segment, default: 4s
            sample_rate: the sample rate of wav, default: 16000
            processer: a processer class to handle wave data
            gender2spk: a list include gender2spk, default: None
        '''
        self.loadscp(target_scp)
        mgr = mp.Manager()
        self.index =mgr.list()#[d for b in buckets for d in b]
        self.segement_length = segement_length * sample_rate
        _dochunk(self.wav_list, self.index, self.segement_length)
        self.index *= repeate
        self.augmentation = MixSpeaker(infer_scp, snr_range=[0,30], mode='first')
        self.addnosie = AddNoise(noise_scp, snr_range=[12,30])
        self.randstates = [ np.random.RandomState(idx) for idx in range(3000)]

    def loadscp(self, scp_list):
        self.wav_list = []
        with open(scp_list) as fid:
            for line in fid:
                tmp = line.strip().split()
                spkid = tmp[0]
                path = tmp[1]
                data, fs = sf.read(path)
                duration = data.shape[0]
                self.wav_list.append({'spkid':spkid, 'path':tmp[1], 'duration':duration})

    def __len__(self):
        return len(self.index)

    def __getitem__(self, index):
        data_info, start_time = self.index[index]
        spkid, path = data_info['spkid'], data_info['path']
        Y, X, emb, snrs = self.augmentation(
                                    (spkid, path),
                                    start=start_time,
                                    randstat=self.randstates[index%3000],
                                    segement_length=self.segement_length
                                )
        tmp = self.addnosie(Y,
                            randstat=self.randstates[(index+11)%3000])
        Y = tmp[0]
        X = tmp[3]*X[0]
        tmp = self.addnosie(emb, randstat=self.randstates[(index+17)%3000])
        emb = tmp[0]
        inputs = Y
        labels = X

        return inputs, labels, emb

def worker(target_list, result_list, start, end, segement_length):
    for item in target_list[start:end]:
        duration = item['duration']
        length = duration
        if length < segement_length:
            sample_index = -1
            if length * 2 < segement_length and length*4 > segement_length:
                sample_index = -2
            elif length * 2 > segement_length:
                sample_index = -1
            else:
                continue
            result_list.append([item, sample_index])
        else:
            sample_index = 0
            while sample_index + segement_length < length:
                result_list.append(
                        [item, sample_index])
                sample_index += segement_length

            if sample_index <= length:
                    result_list.append([
                        item,
                        int(length - segement_length),
                ])

def _dochunk(wav_list, index, segement_length=16000*4, num_threads=12):
        # mutliproccesing
        pc_list = []
        stride = len(wav_list) // num_threads
        if stride < 100:
            p = mp.Process(
                            target=worker,
                            args=(
                                    wav_list,
                                    index,
                                    0,
                                    len(wav_list),
                                    segement_length,
                                )
                        )
            p.start()
            pc_list.append(p)
        else:
            for idx in range(num_threads):
                if idx == num_threads-1:
                    end = len(wav_list)
                else:
                    end = (idx+1)*stride
                p = mp.Process(
                                target=worker,
                                args=(
                                    wav_list,
                                    index,
                                    idx*stride,
                                    end,
                                    segement_length,
                                )
                            )
                p.start()
                pc_list.append(p)
        for p in pc_list:
            p.join()

class Sampler(tud.sampler.Sampler):
    '''

    '''
    def __init__(self, data_source, batch_size):
        it_end = len(data_source) - batch_size + 1
        self.batches = [range(i, i+batch_size)
                        for i in range(0, it_end, batch_size)]
        self.data_source = data_source

    def __iter__(self):
        random.shuffle(self.batches)
        return (i for b in self.batches for i in b)

    def __len__(self):
        return len(self.data_source)

def zero_pad_concat(inputs):
    max_t = max(inp.shape[0] for inp in inputs)
    shape = np.array([len(inputs), max_t, inputs[0].shape[1]])
    inputs_mat = np.zeros(shape, np.float32)
    for idx, inp in enumerate(inputs):
        inputs_mat[idx, :inp.shape[0],:] = inp
    return inputs_mat

def collate_fn(data):
    #inputs, labels, embeddings, snrs = zip(*data)
    inputs, labels, embeddings = zip(*data)
    inputs = np.array(inputs, dtype=np.float32)
    labels = np.array(labels, dtype=np.float32)
    embeddings = np.array(embeddings, dtype=np.float32)
    return torch.from_numpy(inputs),\
            torch.from_numpy(labels),\
            torch.from_numpy(embeddings)

def make_fix_loader(scp_file_name, spk_emb_scp, batch_size, num_workers=12,):
    dataset = FixDataset(scp_file_name, spk_emb_scp,)
    sampler = Sampler(dataset, batch_size)
    loader = tud.DataLoader(dataset,
                            batch_size=batch_size,
                            num_workers=num_workers,
                            collate_fn=collate_fn,
                            sampler=sampler,
                            drop_last=False
                        )
    return loader, None

def make_auto_loader(target_scp, infer_scp, noise_scp, batch_size, num_workers=12, repeate=4):
    dataset = AutoDataset(target_scp, infer_scp, noise_scp, repeate=repeate)
    sampler = Sampler(dataset, batch_size)
    loader = tud.DataLoader(dataset,
                            batch_size=batch_size,
                            num_workers=num_workers,
                            collate_fn=collate_fn,
                            sampler=sampler,
                            drop_last=False
                        )
                            #shuffle=True,
    return loader, None #dataset

if __name__ == '__main__':
    torch.manual_seed(20)
    test_mode = 1
    if test_mode == 1:
        laoder, dataset = make_auto_loader(
                            '/search/odin/huyanxin/data/sogou_train/t1.scp',
                            '/search/odin/huyanxin/data/sogou_train/t1.scp',
                            #'/search/odin/huyanxin/data/sogou_train/spk_list/train_target_spk2wav.scp',
                            #'/search/odin/huyanxin/data/sogou_train/spk_list/train_spk2wav.scp',
                            '/search/odin/huyanxin/workspace/voicefilter-yxhu/data/musan_train.scp',
                            32, num_workers=16)
    else:
        laoder,_ = make_fix_loader(
                        '../data/dev.lst', '/search/odin/huyanxin/data/data_aishell/all_wavs/',
                            800, num_workers=16)
    import time
    stime = time.time()
    fid = open('log.csv', 'w')
    for epoch in range(3):
        if epoch == 0 or epoch == 4:
            torch.manual_seed(20)
        else:
            torch.manual_seed(20+epoch)
        '''
        idx=str(99)
        print(dataset)
        epoch=str(epoch)
        inputs = dataset[0][0]
        sf.write(epoch+'_'+idx+'_noisy.wav', inputs, 16000)
        '''
        print('num_batch', len(laoder))
        for idx, data in enumerate(laoder):
            inputs = data
            if (idx+1)%5 == 0:
                etime = time.time()
                print(epoch, idx, inputs[0].size(), (etime-stime)/100-1.)
                idx=str(idx)
                epoch=str(epoch)
                sf.write(epoch+'_'+idx+'_noisy.wav', inputs[0][0].numpy(), 16000)
                sf.write(epoch+'_'+idx+'_clean.wav', inputs[1][0].numpy(), 16000)
                sf.write(epoch+'_'+idx+'_emb.wav',  inputs[2][0].data.numpy(), 16000)
                break
                stime = etime


import numpy as np
import cv2
import os
import glob
import sys
from collections import defaultdict
from pathlib import Path

GT_PATH = '../../data/mot17/test/'
IMG_PATH = GT_PATH
SAVE_VIDEO = True
RESIZE = 2
IS_GT = False

def draw_bbox(img, bboxes, c=(255, 0, 255)):
  for bbox in bboxes:
    cv2.rectangle(img, (int(bbox[0]), int(bbox[1])),
      (int(bbox[0] + bbox[2]), int(bbox[1] + bbox[3])),
      c, 2, lineType=cv2.LINE_AA)
    ct = [bbox[0] + bbox[2] / 2, bbox[1] + bbox[3] / 2]
    txt = '{}'.format(bbox[4])
    cv2.putText(img, txt, (int(ct[0]), int(ct[1])),
                cv2.FONT_HERSHEY_SIMPLEX, 0.5,
                (255, 122, 255), thickness=1, lineType=cv2.LINE_AA)

if __name__ == '__main__':
  seqs = os.listdir(GT_PATH)
  if SAVE_VIDEO:
    save_path = sys.argv[1][:sys.argv[1].rfind('/res')] + '/video'
    if not os.path.exists(save_path):
      os.mkdir(save_path)
    print('save_video_path', save_path)
  for seq in sorted(seqs):
    print('seq', seq)
    # if len(sys.argv) > 2 and not sys.argv[2] in seq:
    #   continue
    if '.DS_Store' in seq:
      continue
    # if SAVE_VIDEO:
    #   fourcc = cv2.VideoWriter_fourcc(*'XVID')
    #   video = cv2.VideoWriter(
    #     '{}/{}.avi'.format(save_path, seq),fourcc, 10.0, (1024, 750))
    seq_path = '{}/{}/'.format(GT_PATH, seq)
    if IS_GT:
      ann_path = seq_path + 'gt/gt.txt'
    else:
      ann_path = seq_path + 'det/det.txt'
    anns = np.loadtxt(ann_path, dtype=np.float32, delimiter=',')
    print('anns shape', anns.shape)
    image_to_anns = defaultdict(list)
    for i in range(anns.shape[0]):
      if (not IS_GT) or (int(anns[i][6]) == 1 and float(anns[i][8]) >= 0.25):
        frame_id = int(anns[i][0])
        track_id = int(anns[i][1])
        bbox = (anns[i][2:6] / RESIZE).tolist()
        image_to_anns[frame_id].append(bbox + [track_id])

    image_to_preds = {}
    for K in range(1, len(sys.argv)):
      image_to_preds[K] = defaultdict(list)
      pred_path = sys.argv[K] + '/{}.txt'.format(seq)
      try:
        preds = np.loadtxt(pred_path, dtype=np.float32, delimiter=',')
      except:
        preds = np.loadtxt(pred_path, dtype=np.float32, delimiter=' ')
      for i in range(preds.shape[0]):
        frame_id = int(preds[i][0])
        track_id = int(preds[i][1])
        bbox = (preds[i][2:6] / RESIZE).tolist()
        image_to_preds[K][frame_id].append(bbox + [track_id])

    img_path = seq_path + 'img1/'
    images = os.listdir(img_path)
    num_images = len([image for image in images if 'jpg' in image])

    for i in range(num_images):
      frame_id = i + 1
      file_name = '{}/img1/{:06d}.jpg'.format(seq, i + 1)
      file_path = IMG_PATH + file_name
      img = cv2.imread(file_path)
      if RESIZE != 1:
        img = cv2.resize(img, (img.shape[1] // RESIZE, img.shape[0] // RESIZE))
      for K in range(1, len(sys.argv)):
        img_pred = img.copy()
        draw_bbox(img_pred, image_to_preds[K][frame_id])
        cv2.imshow('pred{}'.format(K), img_pred)
      draw_bbox(img, image_to_anns[frame_id])
      cv2.imshow('gt', img)
      cv2.waitKey()
      # if SAVE_VIDEO:
      #   video.write(img_pred)
    # if SAVE_VIDEO:
    #   video.release()


from os.path import join

from numpy import asarray
from pandas import isna, read_csv

from .DATA_DIRECTORY_PATH import DATA_DIRECTORY_PATH

def select_gene_symbol(
    gene_family_name_to_remove=(
        "18S ribosomal RNAs",
        "28S ribosomal RNAs",
        "45S pre-ribosomal RNAs",
        "5.8S ribosomal RNAs",
        "5S ribosomal RNAs",
        "Cytoplasmic transfer RNAs",
        "Long non-coding RNAs (published)",
        "MicroRNAs",
        "Mitochondrially encoded ribosomal RNAs",
        "Mitochondrially encoded tRNAs",
        "Nuclear-encoded mitochondrial transfer RNAs",
        "Piwi-interacting RNA clusters",
        "Ribosomal 45S RNA clusters",
        "Ribosomal 45S rRNA genes outside of clusters",
        "RNAs, 7SL, cytoplasmic",
        "RNAs, Ro-associated Y",
        "Small Cajal body-specific RNAs",
        "Small NF90 (ILF3) associated RNAs",
        "Small nuclear RNAs",
        "Small nucleolar RNAs, C/D box",
        "Small nucleolar RNAs, H/ACA box",
        "Vault RNAs",
        "L ribosomal proteins",
        "Mitochondrial ribosomal proteins",
        "S ribosomal proteins",
        "Mitochondrial complex II: succinate dehydrogenase subunits",
        "Mitochondrial complex III: ubiquinol-cytochrome c reductase complex subunits",
        "Mitochondrial complex IV: cytochrome c oxidase subunits",
        "Mitochondrial complex V: ATP synthase subunits",
        "NADH:ubiquinone oxidoreductase core subunits",
        "NADH:ubiquinone oxidoreductase supernumerary subunits",
    ),
    locus_type_to_keep=("gene with protein product",),
):

    hgnc = read_csv(join(DATA_DIRECTORY_PATH, "hgnc.tsv"), sep="\t", index_col=0)

    remove_by_gene_family_name = asarray(
        tuple(
            not isna(gene_family_name)
            and any(str_ in gene_family_name for str_ in gene_family_name_to_remove)
            for gene_family_name in hgnc["Gene Family Name"]
        )
    )

    print(
        "Removing {}/{} by gene family name ...".format(
            remove_by_gene_family_name.sum(), remove_by_gene_family_name.size
        )
    )

    keep_by_locus_type = asarray(
        tuple(
            not isna(locus_type)
            and any(str_ in locus_type for str_ in locus_type_to_keep)
            for locus_type in hgnc["Locus Type"]
        )
    )

    print(
        "Keeping {}/{} by locus type ...".format(
            keep_by_locus_type.sum(), keep_by_locus_type.size
        )
    )

    remove = remove_by_gene_family_name | ~keep_by_locus_type

    print("Removing {} ...".format(remove.sum()))

    for column_name in ("Gene Family Name", "Locus Type"):

        dataframe = hgnc.loc[remove, column_name].value_counts().to_frame()

        dataframe.index.name = column_name

        dataframe.columns = ("N Removed",)

        print(dataframe)

    gene_symbols = (
        hgnc.loc[~remove, ["Approved Symbol", "Previous Symbols"]]
        .unstack()
        .dropna()
        .unique()
    )

    print("Selected {} gene symbols.".format(gene_symbols.size))

    return gene_symbols.tolist()


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Jan 17 17:33:59 2020

@author: bene
"""
import numpy as np
import imagej

import imagej
ij = imagej.init('/Applications/Fiji.app')
#ij = imagej.init('sc.fiji:fiji:2.0.0-pre-10')
ij.getVersion()

#%%
macro = """
open("/Users/bene/Dropbox/Dokumente/Promotion/PROJECTS/STORMoChip/2020_01_13-EColi_TetraSpecs100_P20_Video_RAW_Fluct/E.Coli_2_my_raw_result-1.tif - drift corrected.tif");
run("Slice Keeper", "first=1 last=100 increment=1");

run("Estimate Drift", "time=1 max=50 reference=[first frame (default, better for fixed)] show_drift_plot choose=/Users/bene/Downloads/test.njt");
close();
close();
close();
run("Correct Drift", "choose=/Users/bene/Downloads/testDriftTable.njt");

run("SRRF Analysis", "ring=1 radiality_magnification=5 axes=6 frames_per_time-point=0 start=0 end=0 max=100 preferred=0");
saveAs("Tiff", "/Users/bene/Downloads/result.tif");
"""

result = ij.py.run_macro(macro)


#!/usr/bin/env python3

""" Plotting a specified tracking efficiency.

Tracking efficiencies are specified via ``AliAnalysisTaskEmcalJetHUtils``, so there
is a hard AliPhysics dependency.
"""

import coloredlogs
import logging
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
from typing import Any, cast, Callable, Dict, List, Tuple

# NOTE: This is out of the expected order, but it must be here to prevent ROOT from stealing the command
#       line options
from jet_hadron.base.typing_helpers import Hist

from pachyderm import histogram
from pachyderm.utils import epsilon

from jet_hadron.base import analysis_objects
from jet_hadron.base import labels
from jet_hadron.base import params
from jet_hadron.plot import base as plot_base

logger = logging.getLogger(__name__)

# Types
T_PublicUtils = Any

def generate_parameters(system: params.CollisionSystem) -> Tuple[np.ndarray, np.ndarray, int, Dict[int, params.SelectedRange]]:
    """ Generate the analysis parameters.

    This can be called multiple times if necessary to retrieve the parameters easily in any function.

    Args:
        system: Collision system.
    Returns:
        (pt_values, eta_values, n_cent_bins, centrality_ranges): Pt values where the efficiency should be evaluated,
            eta values where the efficiency should be evaluated, number of centrality bins, map from centrality bin
            number to centrality bin ranges.
    """
    pt_values = np.linspace(0.15, 9.95, 100 - 1)
    eta_values = np.linspace(-0.85, 0.85, 35)
    n_cent_bins = 4 if system != params.CollisionSystem.pp else 1
    centrality_ranges = {
        0: params.SelectedRange(0, 10),
        1: params.SelectedRange(10, 30),
        2: params.SelectedRange(30, 50),
        3: params.SelectedRange(50, 90),
    }

    return pt_values, eta_values, n_cent_bins, centrality_ranges

def setup_AliPhysics(period: str) -> Tuple[Callable[..., float], Any, T_PublicUtils]:
    """ Retrieve the efficiency function, period and internal efficiency functions and params from AliPhysics.

    Re-factored to a separate function so we can isolate the dependency.

    Args:
        period: Name of the data period.
    Returns:
        (efficiency_function, efficiency_period, T_PublicUtils): The general efficiency function, and
        the period enum value used in calling that function, and a derived class to make the internal
        efficiency functions and parameters public for further testing.
    """
    import ROOT
    # The tasks are in the PWGJE::EMCALJetTasks namespace, so we first retrieve that namespace for convenience.
    user_namespace = ROOT.PWGJE.EMCALJetTasks
    # First retrieve the function
    efficiency_function = user_namespace.AliAnalysisTaskEmcalJetHUtils.DetermineTrackingEfficiency
    # Then the period by string
    efficiency_period = getattr(user_namespace.AliAnalysisTaskEmcalJetHUtils, "k" + period)

    # Provide full access to the rest of the attributes by inheriting and making the members public.
    # The winding route to get there is due to namespace issues in PyROOT.
    code = """
    // Tell ROOT where to find AliRoot headers
    R__ADD_INCLUDE_PATH($ALICE_ROOT)
    // Tell ROOT where to find AliPhysics headers
    R__ADD_INCLUDE_PATH($ALICE_PHYSICS)
    #include <AliAnalysisTaskEmcalJetHUtils.h>
    class JetHUtilsPublic : public PWGJE::EMCALJetTasks::AliAnalysisTaskEmcalJetHUtils {
      public:
        // Can't just use (appears to be because of a PyROOT limitation...):
        //using PWGJE::EMCALJetTasks::AliAnalysisTaskEmcalJetHUtils::LHC11aTrackingEfficiency;
        // So we define the function by hand.
        static double LHC11aTrackingEfficiency(const double trackPt, const double trackEta, const int centralityBin,
                                               const std::string& taskName)
        {
            return PWGJE::EMCALJetTasks::AliAnalysisTaskEmcalJetHUtils::LHC11aTrackingEfficiency(trackPt, trackEta, centralityBin, taskName);
        }
        // LHC11h
        static double LHC11hTrackingEfficiency(const double trackPt, const double trackEta, const int centralityBin,
                                               const std::string& taskName)
        {
            return PWGJE::EMCALJetTasks::AliAnalysisTaskEmcalJetHUtils::LHC11hTrackingEfficiency(trackPt, trackEta, centralityBin, taskName);
        }
        // LHC15o
        static double LHC15oTrackingEfficiency(const double trackPt, const double trackEta, const int centralityBin,
                                               const std::string& taskName)
        {
            return PWGJE::EMCALJetTasks::AliAnalysisTaskEmcalJetHUtils::LHC15oTrackingEfficiency(trackPt, trackEta, centralityBin, taskName);
        }
        static double LHC15oPtEfficiency(const double trackPt, const int centBin)
        {
            std::map<int, const double*> centMap = {
                { 0, PWGJE::EMCALJetTasks::AliAnalysisTaskEmcalJetHUtils::LHC15oParam_0_10_pt },
                { 1, PWGJE::EMCALJetTasks::AliAnalysisTaskEmcalJetHUtils::LHC15oParam_10_30_pt },
                { 2, PWGJE::EMCALJetTasks::AliAnalysisTaskEmcalJetHUtils::LHC15oParam_30_50_pt },
                { 3, PWGJE::EMCALJetTasks::AliAnalysisTaskEmcalJetHUtils::LHC15oParam_50_90_pt },
            };
            return PWGJE::EMCALJetTasks::AliAnalysisTaskEmcalJetHUtils::LHC15oPtEfficiency(trackPt, centMap[centBin]);
        }
        static double LHC15oEtaEfficiency(const double trackEta, const int centBin)
        {
            std::map<int, const double*> centMap = {
                { 0, PWGJE::EMCALJetTasks::AliAnalysisTaskEmcalJetHUtils::LHC15oParam_0_10_eta },
                { 1, PWGJE::EMCALJetTasks::AliAnalysisTaskEmcalJetHUtils::LHC15oParam_10_30_eta },
                { 2, PWGJE::EMCALJetTasks::AliAnalysisTaskEmcalJetHUtils::LHC15oParam_30_50_eta },
                { 3, PWGJE::EMCALJetTasks::AliAnalysisTaskEmcalJetHUtils::LHC15oParam_50_90_eta },
            };
            return PWGJE::EMCALJetTasks::AliAnalysisTaskEmcalJetHUtils::LHC15oEtaEfficiency(trackEta, centMap[centBin]);
        }
        static double LHC15oEtaEfficiencyNormalization(const double centralityBin)
        {
            std::map<int, const double*> centMap = {
                { 0, PWGJE::EMCALJetTasks::AliAnalysisTaskEmcalJetHUtils::LHC15oParam_0_10_eta },
                { 1, PWGJE::EMCALJetTasks::AliAnalysisTaskEmcalJetHUtils::LHC15oParam_10_30_eta },
                { 2, PWGJE::EMCALJetTasks::AliAnalysisTaskEmcalJetHUtils::LHC15oParam_30_50_eta },
                { 3, PWGJE::EMCALJetTasks::AliAnalysisTaskEmcalJetHUtils::LHC15oParam_50_90_eta },
            };
            return centMap[centralityBin][12];
        }
    };
    """
    # Create the class
    # We check for the attribute to ensure that we don't doubly define it.
    try:
        PublicUtils = ROOT.JetHUtilsPublic
    except AttributeError:
        ROOT.gInterpreter.ProcessLine(code)
        # Now that it's defined, we retrieve it.
        PublicUtils = ROOT.JetHUtilsPublic

    return cast(Callable[..., float], efficiency_function), efficiency_period, PublicUtils

def get_15o_eta_max_efficiency(PublicUtils: T_PublicUtils, n_cent_bins: int) -> List[List[float]]:
    """ Get the max of the eta efficiencies for properly normalizing the functions.

    Takes advantage of the public interface to the ``AliAnalysisTaskEmcalJetHUtils`` defiend in the setup.

    Args:
        PublicUtils: Jet-H public utils class.
        n_cent_bins: Number of centrality bins
    Returns:
        Max value in the eta efficiency.
    """
    efficiency_function = PublicUtils.LHC15oEtaEfficiency
    max_values = []
    for centrality_bin in range(n_cent_bins):
        # Then determine the max eta values.
        values_left = []
        values_right = []
        # Record the values with high granularity
        for eta in np.linspace(-0.9, 0, 1000):
            values_left.append(efficiency_function(eta, centrality_bin))
        for eta in np.linspace(-0.06, 0.9, 1000):
            values_right.append(efficiency_function(eta, centrality_bin))

        # Store the result
        max_values.append([np.max(values_left), np.max(values_right)])

    return max_values

def jetH_task_efficiency_for_comparison(period: str, system: params.CollisionSystem,
                                        pt_values: np.ndarray, eta_values: np.ndarray) -> np.ndarray:
    """ Provide the efficiencies from the Jet-hadron correlations task for comparison.

    Requires a hack of the jet-hadron correlations task to make ``fCentBin`` and ``fCurrentRunNumber``.
    One way to do so is:

    .. code-block:: cpp

        void SetCentralityBin(int centBin) { fCentBin = centBin; }
        void SetCurrentRun(int run) { fCurrentRunNumber = run; }

    The comparison only works for LHC11h or LHC11a.

    Args:
        period: Period for calculating the efficiencies.
        system: Collision system.
        pt_values: pT values where the efficiency should be evaluated.
        eta_values: eta values where the efficiency should be evaluated.
    Returns:
        Calculated efficiencies for the same range as the JetHUtils implementation.
    """
    # Setup
    n_cent_bins = 4 if system != params.CollisionSystem.pp else 1
    import ROOT
    jetHTask = ROOT.PWGJE.EMCALJetTasks.AliAnalysisTaskEmcalJetHCorrelations()

    # Corresponds to good runs
    jetHTask.SetCurrentRun(167902)
    # Enable efficiency correction
    efficiency = jetHTask.kEffPP if period == "LHC11a" else jetHTask.kEffAutomaticConfiguration
    system_enum = ROOT.AliAnalysisTaskEmcal.kpp if system == params.CollisionSystem.pp else ROOT.AliAnalysisTaskEmcal.kAA
    jetHTask.SetSingleTrackEfficiencyType(efficiency)

    # Loop over values
    efficiencies = np.zeros(shape = (n_cent_bins, len(pt_values), len(eta_values)))
    for centrality_bin in range(n_cent_bins):
        jetHTask.SetCentralityBin(centrality_bin)
        for pt_index, pt in enumerate(pt_values):
            for eta_index, eta in enumerate(eta_values):
                efficiencies[centrality_bin, pt_index, eta_index] = jetHTask.EffCorrection(
                    eta, pt, system_enum
                )

    return efficiencies

def check_for_accidentally_repeated_parameters(n_cent_bins: int, efficiencies: np.ndarray) -> bool:
    """ Check for accidnetally repeated parameters.

    Namely, it checks whether the efficiencies for one centrality bin is the same as another.
    If they are the same, it indicates that parameters are accidentally repeated.

    Args:
        n_cent_bins: Number of centrality bins.
        efficiencies: Calculated efficiencies.
    Returns:
        True if there are duplicated parameters
    """
    for centrality_bin in range(n_cent_bins):
        for other_centrality_bin in range(n_cent_bins):
            # True if the indices are the same, false if not
            comparison_value = (centrality_bin == other_centrality_bin)
            assert np.allclose(efficiencies[centrality_bin], efficiencies[other_centrality_bin]) is comparison_value

    return False

def calculate_efficiencies(n_cent_bins: int, pt_values: np.ndarray, eta_values: np.ndarray, efficiency_period: Any, efficiency_function: Callable[..., float]) -> np.ndarray:
    """ Caluclate the efficiency given the parameters.

    Args:
        n_cent_bins: Number of centrality bins.
        pt_values: Pt values to evaluate.
        eta_values: Eta values to evaluate.
        efficiency_period: Enum value correspoonding to the period.
        efficiency_function: The actual efficiency function.
    Returns:
        Efficiency evaluated at all centrality bins, pt values, and eta values, indexed as
            [centrality_bin, pt_value, eta_value].
    """
    # Centrality, pt, eta
    efficiencies = np.zeros(shape = (n_cent_bins, len(pt_values), len(eta_values)))
    for centrality_bin in range(n_cent_bins):
        for pt_index, pt in enumerate(pt_values):
            for eta_index, eta in enumerate(eta_values):
                efficiencies[centrality_bin, pt_index, eta_index] = efficiency_function(
                    pt, eta, centrality_bin, efficiency_period, "task_name"
                )

    return efficiencies

def efficiency_properties(n_cent_bins: int, efficiencies: np.ndarray, PublicUtils: T_PublicUtils) -> bool:
    """ Determine and check a variety of efficiency properties.

    Args:
        n_cent_bins: Number of centrality bins.
        efficiencies: Calculated efficiencies.
        PublicUtils: Jet-H public utils class.
    Returns:
        True if the properties were calculated and checked.
    """
    # Determine maximum
    for centrality_bin in range(n_cent_bins):
        m = np.max(efficiencies[centrality_bin])
        logger.debug(f"centrality bin: {centrality_bin}, max: {m}")

    # Find maxima for eta normalization
    if period == "LHC15o":
        logger.info("Checking max eta efficiency normalization")
        max_eta_efficiencies = get_15o_eta_max_efficiency(
            PublicUtils = PublicUtils, n_cent_bins = n_cent_bins,
        )

        for centrality_bin, (left, right) in enumerate(max_eta_efficiencies):
            if np.isclose(left, 1.0, atol = 1e-4) or np.isclose(right, 1.0, atol = 1e-4):
                logger.info(f"Eta efficiency for LHC15o centrality bin {centrality_bin} is properly normalized.")
            else:
                logger.warning(
                    f"Eta efficiency for LHC15o centrality bin {centrality_bin} appears not to be normalized."
                    " Check if this is expected!"
                )
        logger.info(f"Max eta efficiencies: {max_eta_efficiencies}")

    # Check that the parameters are all set correctly.
    logger.info("Checking for accidentally repeated parameters")
    duplicated = check_for_accidentally_repeated_parameters(n_cent_bins = n_cent_bins, efficiencies = efficiencies)
    if duplicated:
        raise ValueError("Some parameters appear to be accidentally repeated.")
    logger.info("No repeated parameters!")

    return True

def plot_tracking_efficiency_parametrization(efficiency: np.ndarray, centrality_range: params.SelectedRange,
                                             period: str, system: params.CollisionSystem,
                                             output_info: analysis_objects.PlottingOutputWrapper) -> None:
    """ Plot the given tracking efficiencies parametrization.

    Args:
        efficiency: Calculated tracking efficiencies.
        centrality_range: Associated centrality range.
        period: Data taking period.
        system: Collision system.
        output_info: Output info for saving figures.
    Returns:
        None.
    """
    # Get the parameters
    pt_values, eta_values, n_cent_bins, centrality_ranges = generate_parameters(system)

    logger.debug(fr"Plotting efficiencies for {centrality_range.min}--{centrality_range.max}%")
    fig, ax = plt.subplots(figsize = (8, 6))
    im = ax.imshow(
        efficiency.T,
        extent = [np.min(pt_values), np.max(pt_values), np.min(eta_values), np.max(eta_values)],
        interpolation = "nearest", aspect = "auto", origin = "lower",
        norm = matplotlib.colors.Normalize(vmin = 0.5, vmax = 1), cmap = "viridis",
    )

    # Add the colorbar
    fig.colorbar(im, ax = ax)

    # Labels
    ax.set_xlabel(fr"${labels.pt_display_label()}\:({labels.momentum_units_label_gev()})$")
    ax.set_ylabel(r"$\eta$")
    title = f"{period} tracking efficiency parametrization"
    if system != params.CollisionSystem.pp:
        title += rf", ${centrality_range.min} \textendash {centrality_range.max}\%$"
    ax.set_title(title, size = 16)

    # Final adjustments
    fig.tight_layout()
    name = f"efficiency_{period}"
    if system != params.CollisionSystem.pp:
        name += f"_centrality_parametrization_{centrality_range.min}_{centrality_range.max}"
    plot_base.save_plot(output_info, fig, name)

    # Cleanup
    plt.close(fig)

def retrieve_efficiency_data(n_cent_bins: int, centrality_ranges: Dict[int, params.SelectedRange]) -> Tuple[List[Hist], List[Hist], List[Hist]]:
    """ Retrieve efficiency data.

    Args:
        n_cent_bins: Number of centrality bins.
        centrality_ranges: Map from centrality bin numbers to centrality ranges.
    Returns:
        (2D efficiency data, 1D pt efficiency data, 1D eta efficiency data)
    """
    # Retrieve histograms
    hists = histogram.get_histograms_in_list(
        filename = "trains/PbPbMC/55/AnalysisResults.root",
        list_name = "AliAnalysisTaskPWGJEQA_tracks_caloClusters_emcalCells_histos"
    )
    matched_sparse = hists["tracks_Matched"]
    generator_sparse = hists["tracks_PhysPrim"]

    # Retrieve the centrality dependent data
    efficiency_data_1D_pt = []
    efficiency_data_1D_eta = []
    efficiency_data_2D = []
    for dimension in ["1D", "2D"]:
        for centrality_bin, centrality_range in centrality_ranges.items():
            # Select in centrality
            matched_sparse.GetAxis(0).SetRangeUser(centrality_range.min + epsilon, centrality_range.max - epsilon)
            generator_sparse.GetAxis(0).SetRangeUser(centrality_range.min + epsilon, centrality_range.max - epsilon)
            # Restrict pt range to < 10 GeV
            matched_sparse.GetAxis(1).SetRangeUser(0.15, 10)
            generator_sparse.GetAxis(1).SetRangeUser(0.15, 10)

            if dimension == "2D":
                # (pt_gen, eta_gen) - order is reversed because 2D API is backwards...
                pt_gen_matched = matched_sparse.Projection(2, 1)
                pt_gen_matched.SetName(f"pt_gen_matched_cent_{centrality_bin}")
                # (pt_gen, eta_gen, findable)
                pt_gen_2d = generator_sparse.Projection(1, 2, 4)
                pt_gen_2d.SetName(f"pt_gen_2D_cent_{centrality_bin}")
                # Select only findable particles and use that efficiency.
                pt_gen_2d.GetZaxis().SetRange(2, 2)
                pt_gen_findable = pt_gen_2d.Project3D("yx")
                logger.debug(f"pt_gen_matched: {pt_gen_matched}, pt_gen_findable: {pt_gen_findable}")

                efficiency_hist = pt_gen_findable.Clone()
                efficiency_hist.Divide(pt_gen_matched, pt_gen_findable, 1.0, 1.0, "B")
                efficiency_data_2D.append(efficiency_hist)

            elif dimension == "1D":
                # pT 1D efficiency
                # NOTE: We can't just project from the 2D efficiency. Integrating over eta will get the wrong
                #       wrong values. Here, we project before we divide to get the right answer.
                pt_gen_matched1D = matched_sparse.Projection(1)
                pt_gen_matched1D.SetName(f"pt_gen_matched_1D_cent_{centrality_bin}")
                pt_gen_1d = generator_sparse.Projection(4, 1)
                pt_gen_1d.SetName(f"pt_gen_1D_cent_{centrality_bin}")
                # Select only findable particles and use that efficiency.
                pt_gen_1d.GetYaxis().SetRange(2, 2)
                pt_gen_findable = pt_gen_1d.ProjectionX()
                logger.debug(f"pt_gen_matched1D: {pt_gen_matched1D}, pt_gen_findable: {pt_gen_findable}")

                efficiency_1D = pt_gen_findable.Clone()
                efficiency_1D.Divide(pt_gen_matched1D, pt_gen_findable, 1.0, 1.0, "B")
                efficiency_data_1D_pt.append(efficiency_1D)

                # Eta 1D
                eta_gen_matched_1D = matched_sparse.Projection(2)
                eta_gen_matched_1D.SetName(f"eta_gen_matched_1D_cent_{centrality_bin}")
                eta_gen_1D = generator_sparse.Projection(4, 2)
                eta_gen_1D.SetName(f"eta_gen_1D_cent_{centrality_bin}")
                # Select only findable particles and use that efficiency.
                eta_gen_1D.GetYaxis().SetRange(2, 2)
                eta_gen_findable = eta_gen_1D.ProjectionX()
                logger.debug(f"eta_gen_matched_1D: {eta_gen_matched_1D}, eta_gen_findable: {eta_gen_findable}")

                efficiency_1D = eta_gen_findable.Clone()
                efficiency_1D.Divide(eta_gen_matched_1D, eta_gen_findable, 1.0, 1.0, "B")
                efficiency_data_1D_eta.append(efficiency_1D)
            else:
                # Shouldn't ever really happen, but just for sanity.
                raise RuntimeError(f"Invalid dimension {dimension}")

    return efficiency_data_2D, efficiency_data_1D_pt, efficiency_data_1D_eta

def calculate_residual_2D(efficiency_data: Hist, efficiency_function: Callable[..., float],
                          efficiency_period: Any, centrality_bin: int) -> Tuple[np.ndarray, List[float], List[float]]:
    """ Calculate residual for 2D tracking efficiency.

    There is a separate 1D and 2D function for convenience. If there is no entries for a particular
    bin, we set the value to NaN so that it can be ignored later when plotting.

    Args:
        efficiency_data: 2D efficiency data.
        efficiency_function: Efficiency function.
        efficiency_period: Efficiency period.
        centrality_bin: Centrality bin.
    Returns:
        Calculated residual, pt values where it was evaluated, eta values where it was evaluated.
    """
    pts = [efficiency_data.GetXaxis().GetBinCenter(x) for x in range(1, efficiency_data.GetXaxis().GetNbins() + 1)]
    etas = [efficiency_data.GetYaxis().GetBinCenter(y) for y in range(1, efficiency_data.GetYaxis().GetNbins() + 1)]
    residual = np.zeros(shape = (efficiency_data.GetXaxis().GetNbins(),
                                 efficiency_data.GetYaxis().GetNbins()))
    # Loop over all of the bins in the data histogram.
    chi_2 = []
    for pt_index, pt in enumerate(pts):
        for eta_index, eta in enumerate(etas):
            x = pt_index + 1
            y = eta_index + 1
            # Calculate the efficiency. It's calculated again here to ensure that it's evaluated at exactly
            # the same location as in the data histogram.
            efficiency_at_value = efficiency_function(pt, eta, centrality_bin, efficiency_period, "task_name")

            # Determine the histogram value, setting it to NaN if there's no entries.
            if np.abs(efficiency_data.GetBinContent(x, y)) < epsilon:
                value = np.nan
            else:
                value = (efficiency_data.GetBinContent(x, y) - efficiency_at_value) / efficiency_at_value * 100.
                # The points around the edges aren't super reliable for calcuating chi squared
                if pt > 1 and np.abs(eta) < 0.8:
                    chi_2.append(np.power(efficiency_data.GetBinContent(x, y) - efficiency_at_value, 2) / np.power(efficiency_data.GetBinError(x, y), 2))

            residual[pt_index, eta_index] = value

    # Check max values
    logger.debug(f"min efficiency_data: {efficiency_data.GetMinimum()}, "
                 f"max efficiency_data: {efficiency_data.GetMaximum()}")
    logger.debug(f"min residual: {np.nanmin(residual)}, max residual: {np.nanmax(residual)}")
    logger.debug(f"standard mean: {np.nanmean(residual)}")
    logger.debug(f"restricted mean: {np.nanmean(residual[:,np.abs(etas) < 0.8])}")
    logger.debug(f"len(pts): {len(pts)}, len(etas): {len(etas)}")

    # Check chi squared
    chi_squared = np.sum(chi_2)
    # 23 is the number of parameters (10 + 13) at any given point
    ndf = len(chi_2) - 23
    logger.warning("NOTE: The restricted chi squared value calculated here may not be super reliable.")
    logger.info(f"Chi squared: {chi_squared}")
    logger.info(f"NDF: {ndf}")
    logger.info(f"chi2/ndf: {chi_squared / ndf}")

    return residual, pts, etas

def plot_residual(residual: np.ndarray, pts: List[float], etas: List[float],
                  period: str, centrality_bin: int, centrality_ranges: Dict[int, params.SelectedRange],
                  output_info: analysis_objects.PlottingOutputWrapper) -> None:
    """ Plot the residual between the data and the parametrization.

    Args:
        residual: Calculated residual.
        pts: Pt values where the residual was evaluated.
        etas: Eta values where the residual was evaluated.
        period: Name of the data taking period.
        centrality_bin: Centrality bin.
        centraliy_ranges: Map of centrality bins to ranges.
        output_info: Output info for saving figures.
    Returns:
        None.
    """
    fig, ax = plt.subplots(figsize = (8, 6))
    im = ax.imshow(
        residual.T, extent = [np.nanmin(pts), np.nanmax(pts), np.nanmin(etas), np.nanmax(etas)],
        interpolation = "nearest", aspect = "auto", origin = "lower",
        # An even normalization is better for the colorscheme.
        # NOTE: This causes clipping at the lowest pt values, but I don't think this is a big problem.
        norm = matplotlib.colors.Normalize(
            #vmin = np.nanmin(residuals[centrality_bin]), vmax = np.nanmax(residuals[centrality_bin])
            vmin = -40, vmax = 40
        ),
        # This is a good diverging color scheme when it's centered at 0.
        cmap = "RdBu",
    )

    # Add the colorbar
    color_bar = fig.colorbar(im, ax = ax)
    color_bar.set_label(r"(data - fit)/fit (\%)")

    # Labels
    ax.set_xlabel(fr"${labels.pt_display_label()}\:({labels.momentum_units_label_gev()})$")
    ax.set_ylabel(r"$\eta$")
    title = f"{period} tracking efficiency residuals"
    if system != params.CollisionSystem.pp:
        centrality_range = centrality_ranges[centrality_bin]
        title += rf", ${centrality_range.min} \textendash {centrality_range.max}\%$"
    ax.set_title(title, size = 16)

    # Final adjustments
    fig.tight_layout()
    name = f"efficiency_residuals_{period}"
    if system != params.CollisionSystem.pp:
        centrality_range = centrality_ranges[centrality_bin]
        name += f"_centrality_{centrality_range.min}_{centrality_range.max}"
    plot_base.save_plot(output_info, fig, name)

    # Cleanup
    plt.close(fig)

def plot_2D_efficiency_data(efficiency_hist: Hist, centrality_range: params.SelectedRange, output_info: analysis_objects.PlottingOutputWrapper) -> None:
    """ Plot the 2D efficiency data

    Args:
        efficiency_hist: Efficiecny histogram.
        centrality_range: Centrality range.
        output_info: Output info for saving figures.
    Returns:
        None.
    """
    X, Y, efficiency_data = histogram.get_array_from_hist2D(hist = efficiency_hist)
    logger.debug(f"efficiency data min: {np.nanmin(efficiency_data)}, max: {np.nanmax(efficiency_data)}")
    fig, ax = plt.subplots(figsize = (8, 6))
    im = ax.imshow(
        efficiency_data.T, extent = [np.nanmin(X), np.nanmax(X), np.nanmin(Y), np.nanmax(Y)],
        interpolation = "nearest", aspect = "auto", origin = "lower",
        norm = matplotlib.colors.Normalize(
            vmin = np.nanmin(efficiency_data), vmax = np.nanmax(efficiency_data)
            #vmin = 0.5, vmax = 1,
        ),
        cmap = "viridis",
    )

    # Add the colorbar
    color_bar = fig.colorbar(im, ax = ax)
    color_bar.set_label("Efficiency")

    # Labels
    ax.set_xlabel(fr"${labels.pt_display_label()}\:({labels.momentum_units_label_gev()})$")
    ax.set_ylabel(r"$\eta$")
    title = f"{period} tracking efficiency data"
    if system != params.CollisionSystem.pp:
        title += rf", ${centrality_range.min} \textendash {centrality_range.max}\%$"
    ax.set_title(title, size = 16)

    # Final adjustments
    fig.tight_layout()
    name = f"efficiency_{period}"
    if system != params.CollisionSystem.pp:
        name += f"_centrality_{centrality_range.min}_{centrality_range.max}"
    plot_base.save_plot(output_info, fig, name)

    # Cleanup
    plt.close(fig)

def plot_1D_pt_efficiency(efficiency: Hist, PublicUtils: T_PublicUtils, efficiency_period: Any,
                          centrality_bin: int, centrality_range: params.SelectedRange,
                          output_info: analysis_objects.PlottingOutputWrapper) -> None:
    """ Plot 1D pt efficiency.

    Args:
        efficiency: Pt efficiency hist.
        PublicUtils: Jet-H public utils class.
        efficiency_period: Data taking period in the efficiency enum.
        centrality_bin: int
        centrality_range: Centrality range.
        output_info: Output info for saving figures.
    Returns:
        None.
    """
    # 1D efficiency as a function of pt
    logger.debug(f"max efficiency_1D: {efficiency.GetMaximum()}")
    h = histogram.Histogram1D.from_existing_hist(efficiency)
    fig, ax = plt.subplots(figsize = (8, 6))
    ax.errorbar(
        h.x, h.y, yerr = h.errors,
        label = "${labels.pt_display_label()}$",
        color = "black", marker = ".", linestyle = "",
    )

    # Efficiency function
    parametrization = []
    for x in h.x:
        parametrization.append(PublicUtils.LHC15oPtEfficiency(x, centrality_bin))
    ax.plot(
        h.x, parametrization,
        label = "${labels.pt_display_label()}$ param.",
        color = "red",
    )

    # Ensure that it's on a consistent axis
    ax.set_ylim(0.6, 1)

    # Labels
    ax.set_xlabel(fr"${labels.pt_display_label()}\:({labels.momentum_units_label_gev()})$")
    ax.set_ylabel(r"Efficiency")
    title = f"{period} ${labels.pt_display_label()}$ tracking efficiency"
    if system != params.CollisionSystem.pp:
        title += rf", ${centrality_range.min} \textendash {centrality_range.max}\%$"
    ax.set_title(title, size = 16)

    # Final adjustments
    fig.tight_layout()
    name = f"efficiency_pt_{period}"
    if system != params.CollisionSystem.pp:
        name += f"_centrality_{centrality_range.min}_{centrality_range.max}"
    plot_base.save_plot(output_info, fig, name)

def plot_1D_eta_efficiency(efficiency: Hist, PublicUtils: T_PublicUtils, efficiency_period: Any,
                           centrality_bin: int, centrality_range: params.SelectedRange,
                           output_info: analysis_objects.PlottingOutputWrapper) -> None:
    """ Plot 1D eta efficiency.

    Args:
        efficiency: Eta efficiency hist.
        PublicUtils: Jet-H public utils class.
        efficiency_period: Data taking period in the efficiency enum.
        centrality_bin: int
        centrality_range: Centrality range.
        output_info: Output info for saving figures.
    Returns:
        None.
    """
    # 1D efficiency as a function of eta
    logger.debug(f"max efficiency_1D: {efficiency.GetMaximum()}")
    h = histogram.Histogram1D.from_existing_hist(efficiency)
    fig, ax = plt.subplots(figsize = (8, 6))
    ax.errorbar(
        h.x, h.y, yerr = h.errors,
        label = r"$\eta$",
        color = "black", marker = ".", linestyle = "",
    )

    # Efficiency function
    parametrization = []
    for x in h.x:
        # Only evaluate for eta < 0.9 - otherwise the function goes back up again.
        if np.abs(x) < 0.9:
            # Need to undo the scale factor here to get something that's reasonably close..
            value = PublicUtils.LHC15oEtaEfficiency(x, centrality_bin) * \
                PublicUtils.LHC15oEtaEfficiencyNormalization(centrality_bin)
        else:
            # Just set to 0 since it's not meaningful
            value = 0
        parametrization.append(value)

    # And plot it
    ax.plot(
        h.x, parametrization,
        label = r"$\eta$ param.",
        color = "red",
    )

    # Ensure that it's on a consistent axis
    ax.set_ylim(0.6, 1)

    # Labels
    ax.set_xlabel(fr"$\eta$")
    ax.set_ylabel(r"Efficiency")
    title = fr"{period} $\eta$ tracking efficiency"
    if system != params.CollisionSystem.pp:
        title += rf", ${centrality_range.min} \textendash {centrality_range.max}\%$"
    ax.set_title(title, size = 16)

    # Final adjustments
    fig.tight_layout()
    name = f"efficiency_eta_{period}"
    if system != params.CollisionSystem.pp:
        name += f"_centrality_{centrality_range.min}_{centrality_range.max}"
    plot_base.save_plot(output_info, fig, name)

def characterize_tracking_efficiency(period: str, system: params.CollisionSystem) -> None:
    """ Characterize the tracking efficiency.

    Args:
        period: Period for calculating the efficiencies.
        system: Collision system.
    Returns:
        Calculated efficiencies for the same range as the JetHUtils implementation.
    """
    # Setup
    logger.warning(f"Plotting efficiencies for {period}, system {system}")
    efficiency_function, efficiency_period, PublicUtils = setup_AliPhysics(period)
    # Setting up evaluation ranges.
    pt_values, eta_values, n_cent_bins, centrality_ranges = generate_parameters(system)
    # Plotting output location
    output_info = analysis_objects.PlottingOutputWrapper(
        output_prefix = f"output/{system}/{period}/trackingEfficiency",
        printing_extensions = ["pdf"],
    )

    # Calculate the efficiency.
    efficiencies = calculate_efficiencies(
        n_cent_bins = n_cent_bins, pt_values = pt_values, eta_values = eta_values,
        efficiency_period = efficiency_period, efficiency_function = efficiency_function,
    )

    # Calculate and check efficiency properties.
    result = efficiency_properties(
        n_cent_bins = n_cent_bins, efficiencies = efficiencies, PublicUtils = PublicUtils
    )
    if not result:
        raise RuntimeError("Failed to calculate all efficiency properties. Check the logs!")

    # Comparison to previous task
    if period in ["LHC11a", "LHC11h"]:
        try:
            comparison_efficiencies = jetH_task_efficiency_for_comparison(period, system, pt_values, eta_values)
            np.testing.assert_allclose(efficiencies, comparison_efficiencies)
            logger.info("Efficiencies agree!")
        except AttributeError as e:
            logger.warning(f"{e.args[0]}. Skipping!")
    else:
        logger.info("Skipping efficiencies comparison because no other implementation exists.")

    # Done checking the efficiency parametrization, so now plot it.
    for centrality_bin in range(n_cent_bins):
        plot_tracking_efficiency_parametrization(
            efficiencies[centrality_bin], centrality_ranges[centrality_bin],
            period, system, output_info
        )

    # Next, compare to the actual efficiency data.
    # We only have the LHC15o data easily available.
    if period == "LHC15o":
        # First, retrieve efficiency data
        efficiency_data_2D, efficiency_data_1D_pt, efficiency_data_1D_eta = \
            retrieve_efficiency_data(n_cent_bins, centrality_ranges)

        # Calculate residuals.
        logger.info("Calculating and plotting residuals")
        # Setup
        residuals: np.ndarray = None
        for centrality_bin, centrality_range in centrality_ranges.items():
            # Finish the setup
            if residuals is None:
                residuals = np.zeros(
                    shape = (n_cent_bins, efficiency_data_2D[centrality_bin].GetXaxis().GetNbins(),
                             efficiency_data_2D[centrality_bin].GetYaxis().GetNbins())
                )
            logger.debug(f"residuals shape: {residuals.shape}")

            # Calculate the residuals.
            residuals[centrality_bin], pts, etas, = calculate_residual_2D(
                efficiency_data_2D[centrality_bin], efficiency_function, efficiency_period, centrality_bin
            )

            # Plot the residuals.
            plot_residual(residuals[centrality_bin], pts, etas, period, centrality_bin, centrality_ranges, output_info)

            # Plot the 2D efficiency data
            plot_2D_efficiency_data(efficiency_data_2D[centrality_bin],
                                    centrality_ranges[centrality_bin], output_info)

            # 1D efficiency comparison
            # NOTE: Since we are using the 2D efficiency parametrization, the 1D fits aren't going to be quite right.
            #       This is basically because we can't integrate an efficiency (it's not well defined).
            #       It will be fairly close, but not exactly right.
            plot_1D_pt_efficiency(efficiency_data_1D_pt[centrality_bin],
                                  PublicUtils, efficiency_period,
                                  centrality_bin, centrality_ranges[centrality_bin],
                                  output_info)
            plot_1D_eta_efficiency(efficiency_data_1D_eta[centrality_bin],
                                   PublicUtils, efficiency_period,
                                   centrality_bin, centrality_ranges[centrality_bin],
                                   output_info)

if __name__ == "__main__":
    # Basic setup
    coloredlogs.install(
        level = logging.DEBUG,
        fmt = "%(asctime)s %(name)s:%(lineno)d %(levelname)s %(message)s"
    )
    # Quiet down the matplotlib logging
    logging.getLogger("matplotlib").setLevel(logging.INFO)

    # Run for each period
    for period, system in [("LHC11a", params.CollisionSystem.pp),
                           ("LHC11h", params.CollisionSystem.PbPb),
                           ("LHC15o", params.CollisionSystem.PbPb)]:
        characterize_tracking_efficiency(period, system)