u"""
Enable the NASA 9 term polynomials to create chemical compound objects.

This module allows the user to use the NASA9 term polynomials (abbreviated as
NASA 9) to model chemical compounds and reactions. It is worth noting that all
the output is stripped of any phyisical unit, that is, results are returned as
numpy floats. *Therefore to end any form of ambiguity we reinstate that all the
results are in SI units using molar basis*. It is up to the user to beware of
any physical unit conversion concerning his problem.

Classes:

    Compound: chemical compound present in the NASA9 term polynomials
    database.

    CompoundIdealGas: chemical compound as an ideal gas present in the NASA9
    term polynomials database. It inherits from Compound.

    Reaction: model of a chemical reaction.

Example:
    >>> import thermopy
    >>> from thermopy import nasa9polynomials as nasa9
    >>> db = nasa9.Database()
    >>> caf2 = db.set_compound('caf2')
    >>> print(caf2.elements)
    [('C', 1), ('F', 2)]
    >>> print(caf2.inchikey)
    WUKWITHWXAAZEY-UHFFFAOYSA-L
    >>> print(caf2.enthalpy_of_formation)
    -790828.409
    >>> print(caf2.heat_capacity(300))
    51.2707324499
    >>> print(caf2.molecular_weight)
    0.0780748064
    >>> water = db.set_compound('h2o(l)')
    >>> print(water.entropy(300))
    69.633703
    >>> print(water.elements)
    [('H', 2), ('O', 1)]

References: [1] Bonnie J. McBride, Michael J. Zehe, and Sanford Gordon. NASA
Glenn Coefficients for Calculating Thermodynamic Properties of Individual
Species. September 2002.

"""

import os
import re
import xml.etree.ElementTree as ET
import numpy as np
from thermopy.constants import ideal_gas_constant
_R = ideal_gas_constant[0]

DATADIR = '/usr/share/thermopy'

class Compound(object):
    u"""
    Chemical compound present in the NASA9 term polynomials database.

    It is usually set by nasa9polyniomials.Database.set_compound('identifier').
    If set this way this method already instantiates either a Compound or a
    CompoundIdealGas based on the phase of the chemical compound specified by
    the identifier (see note below).

    Note: *The default phase for compounds in this database is gas.* Thus if
    instantiating with the name 'H2O' one would get steam. For liquid water and
    ice one would rather look for 'H2O(L)' or 'H2O(cr)'. The same is valid for
    a lot of compounds expected to be in the condensed form (such as NaCl,
    Tungsten, etc).

    It has all the thermodynamic functions listed in [1] as methods which take
    temperature as their sole argument. Those were expanded to include
    gibbs_energy that could be defined by the given functions.

    Attributes:
        canonical_smiles (str): Canonical SMILES
            (Simplified molecular-input line-entry system) of the compound.
        cas_number (str): CAS (Chemical Abstract Service) number of the
            compound.
        comment (str): Comment found in the xml database. Usually references.
        condensed (bool): True if the compound is condensed, False if not.
        xml_compounds (list): List containing tuples of two entries. The first
            is the xml_compound and the second is the proportion of the
            xml_compound in the molecule. Both values are strings.
        enthalpy_of_formation (float): Enthalpy of formation of the compound.
        inchikey (str): InChI (International Chemical Identifier) key
            for the compound.
        inp_name (str): Name as per the original 'inp' file.
        iupac_name (str): IUPAC (International Union of Pure and Applied
            Chemistry) name of the compound.
        molecular_weight (float): Molecular
            weight of the compound.
        reference (str): Reference for the compound. See [1] for details.

    Methods:
        enthalpy:calculates the enthalpy for a Compound object.
        entropy: calculates the entropy for a Compound object.
        gibbs_energy: calculates the Gibbs energy for a Compound object.
        heat_capacity: calculates the heat capacity for a Compound object.

    Subclasses:
        CompoundIdealGas: chemical compound as an ideal gas present in
        the NASA9 term polynomials database. It inherits from Compound.

    Examples:
        >>> import thermopy
        >>> from thermopy import nasa9polynomials as nasa9
        >>> db = nasa9.Database()
        >>> uf6 = db.set_compound('uf6(cr)')
        >>> print(uf6)
        hexafluorouranium: UF6(cr)

    """

    def __init__(self, xml_compound):
        u"""
        Instantiate a Compound object from xml info.

        Arguments:
            xml_compound: xml tree containing the relevant fields to
        characterize the attributes and boundaries of temperature for which
        calculations are valid.

        """
        self._xml_compound = xml_compound
        self.inp_name = xml_compound.attrib['inp_file_name']
        self.inchikey = xml_compound.find('identification').find(
            'InChIKey').text
        self.canonical_smiles = xml_compound.find(
            'identification').find('canonical_smiles').text
        self.cas_number = xml_compound.find('identification').find(
            'cas_number').text
        self.iupac_name = xml_compound.find('identification').find(
            'IUPAC_name').text
        self.comment = xml_compound.find('comment').text
        self.reference = xml_compound.find('reference').text
        self.elements = self._get_xml_compounds(xml_compound.find(
            'elements'))
        self.condensed = bool(
            xml_compound.find('condensed').text == 'True')
        self.molecular_weight = float(
            xml_compound.find('molecular_weight').text)
        self.enthalpy_of_formation = float(
            xml_compound.find('hf298.15').text)

    def _get_xml_compounds(self, xml_compound):
        u"""
        Return a list of tuples containing the elements and their
        proportion in the chemical compound.

        Arguments:
            xml_compound (list): list of Element Tree objects containing
        elements.

        Returns:
            list: list of tuples. Tuples are of the form (str, int) where
        str is the symbol of the element and int is its proportion in the
        molecule.

        """
        xml_compounds_list = []
        for one_xml_compound in xml_compound:
            xml_compounds_list.append(tuple(*map(
                lambda x, y: (x, int(y)),
                *one_xml_compound.items()[0])))
        return xml_compounds_list

    def _evaluate_temperature_interval(self, T):
        u"""
        Ouput temperature interval to be used by public methods.

        Helper method to ouput temperature interval order.

        Arguments:
            T (float): temperature.

        Returns:
            int: The order of the temperature range (0th, 1st, 2nd, etc).
        Some compounds have more than one temperature range with different
        corrisponding coefficients. Therefore the temperature range has to be
        specified.

        Example: The KI gas has two temperature intervals (thus two sets of
        coefficients to be used). The ranges are: [200, 1000] and [1000, 6000]
        as for most gases. Thus requiring a property to be measured at 1100 K
        the second interval should be used and this method shall return the
        number 1 (as opposed to zero).

        """
        for (i, Trange) in enumerate(self._xml_compound.findall('T_range')):
            if (float(Trange.attrib['Tlow']) <= T <= float(
                    Trange.attrib['Thigh'])):
                return i
        raise Exception('Temperature out of range for '
                        + self.iupac_name + '/' + self.inp_name
                        + ' with ' + str(T) + ' K')

    def heat_capacity(self, T):
        u"""
        Calculate molar heat capacity at constant pressure for standard
        state.

        Arguments:
            T (float): temperature.

        Returns:
            numpy_float: The molar heat capacity for the compound for a
        given temperature.

        Examples:
            >>> import thermopy
            >>> from thermopy import nasa9polynomials as nasa9
            >>> db = nasa9.Database()
            >>> libr = db.set_compound('LiBr')
            >>> print(libr, '-', libr.heat_capacity(2700))
            lithium;bromide: LiBr - 39.6593057506

        """
        coefficients = np.empty(9, dtype=np.float32)
        for (i, coef) in enumerate(self._xml_compound.findall(
                'T_range')[self._evaluate_temperature_interval(T)]):
            coefficients[i] = np.array(coef.text, dtype=np.float32)
        exponents = np.array([-2, -1, 0, 1, 2, 3, 4], dtype=np.signedinteger)
        return np.sum(
            np.multiply(
                np.power(T, exponents, dtype=np.float32),
                coefficients[0:7]), dtype=np.float32) * _R

    def enthalpy(self, T):
        u"""
        Calculate molar enthalpy at constant pressure for the compound for
        a given temperature.

        Arguments:
            T (float): temperature.

        Returns:
            numpy_float: The molar enthalpy for the compound for a given
            temperature.

        Examples:
            >>> import thermopy
            >>> from thermopy import nasa9polynomials as nasa9
            >>> db = nasa9.Database()
            >>> magnesium_hydroxide = db.set_compound('Mg(OH)2(cr)')
            >>> print(magnesium_hydroxide, '-',
            >>> magnesium_hydroxide.enthalpy(500))  # J/mol
            Mg(OH)2(cr) - -906097.801815

        """
        coefficients = np.empty(9, dtype=np.float32)
        for (i, coef) in enumerate(self._xml_compound.findall(
                'T_range')[self._evaluate_temperature_interval(T)]):
            coefficients[i] = np.array(coef.text, dtype=np.float32)
        exponents = np.array([-2, -1, 0, 1, 2, 3, 4, -1],
                             dtype=np.signedinteger)
        other_factors = np.array([-1, np.log(T), 1, 0.5, 1/3, 0.25, 0.2, 1],
                                 dtype=np.float32)
        return np.sum(
            np.multiply(
                np.multiply(
                    np.power(T, exponents, dtype=np.float32),
                    coefficients[0:8]),
                other_factors), dtype=np.float32) * _R * T

    def entropy(self, T):
        u"""
        Calculate molar entropy at constant pressure for the compound for
        a given temperature.

        Arguments:
            T (float): temperature.

        Returns:
            numpy_float: The molar entropy for the compound for a given
            temperature.

        Examples:
            >>> import thermopy
            >>> from thermopy import nasa9polynomials as nasa9
            >>> db = nasa9.Database()
            >>> argon = db.set_compound('ar')
            >>> print(argon, '-', argon.entropy(200))
            argon: Ar - 146.546470215

        """
        coefficients = np.empty(9, dtype=np.float32)
        for (i, coef) in enumerate(self._xml_compound.findall(
                'T_range')[self._evaluate_temperature_interval(T)]):
            coefficients[i] = np.array(coef.text, dtype=np.float32)
        exponents = np.array([-2, -1, 0, 1, 2, 3, 4, 0, 0],
                             dtype=np.signedinteger)
        other_factors = np.array([-0.5, -1, np.log(T), 1, 0.5, 1/3, 0.25,
                                  0, 1],
                                 dtype=np.float32)
        return np.sum(
            np.multiply(
                np.multiply(np.power(T, exponents, dtype=np.float32),
                            coefficients[:]),
                other_factors), dtype=np.float32) * _R

    def gibbs_energy(self, T):
        u"""
        Calculate molar Gibbs energy at constant pressure for the compound
        for a given temperature.

        Arguments:
            T (float): temperature.

        Returns:
            numpy_float: The molar entropy for the compound for a given
            temperature.

        Examples:
            >>> import thermopy
            >>> from thermopy import nasa9polynomials as nasa9
            >>> db = nasa9.Database()
            >>> csbr = db.set_compound('csbr(cr)')
            >>> print(csbr, '-', csbr.gibbs_energy(273.15))
            cesium;bromide: CsBr(cr) - -436504.410044

        """
        return self.enthalpy(T) - T * self.entropy(T)

    def __str__(self):
        if self.iupac_name != 'N/A':
            return str(self.iupac_name) + ': ' + self.inp_name
        elif self.canonical_smiles != 'N/A':
            return str(self.canonical_smiles)
        else:
            return str(self.inp_name)

class CompoundIdealGas(Compound):
    u"""
    Chemical compound as an ideal gas present in the NASA9 term
    polynomials database.

    It is usually set by nasa9polyniomials.Database.set_compound('identifier').
    If set this way this method already instantiates either a Compound or a
    CompoundIdealGas based on the phase of the chemical compound specified by
    the identifier (see note below).

    Note: *The default phase for compounds in this database is gas.* Thus if
    instantiating with the name 'H2O' one would get steam. For liquid water and
    ice one would rather look for 'H2O(L)' or 'H2O(cr)'. The same is valid for
    a lot of compounds expected to be in the condensed form (such as NaCl,
    Tungsten, etc).

    It adds two extra methods which come from the thermodynamics of Ideal
    Gases.

    Inherits from Compound.

    Methods:
        heat_capacity_constant_v: calculates the heat capacity at a
        constant volume for a CompoundIdealGas object.
        internal_energy: calculates
        the internal energy for a CompoundIdealGas object.

    Examples:
        >>> # Instantiating a Compound whose condensed attributed is False
        >>> # automatically sets it as an Ideal Gas:
        >>> import thermopy
        >>> from thermopy import nasa9polynomials as nasa9
        >>> db = nasa9.Database()
        >>> co2 = db.set_compound('CO2')
        >>> print(co2, type(co2))
        carbon dioxide: CO2 <class
        'thermopy.nasa9polynomials.CompoundIdealGas'>

    """

    def __init__(self, xml_compound):
        u"""
        Initialize an ideal gas Compound.

        Arguments:
            xml_compound: xml tree containing the relevant fields to
        characterize the attributes and boundaries of temperature for which
        calculations are valid.

        """
        Compound.__init__(self, xml_compound)

    def heat_capacity_constant_v(self, T):
        u"""
        Calculate molar heat capacity at constant volume for standard
        state.

        Arguments:
            T (float): temperature.

        Returns:
            numpy_float: The molar heat capacity for the compound for a
            given temperature.

        Examples:
            >>> import thermopy
            >>> from thermopy import nasa9polynomials as nasa9
            >>> db = nasa9.Database()
            >>> xenon = db.set_compound('Xe')
            >>> print(xenon, '-', xenon.heat_capacity(298.15))
            xenon: Xe - 20.78618
            >>> print(xenon, '-', xenon.heat_capacity_constant_v(298.15))
            xenon: Xe - 12.471708
            >>> print('subtracting both hc:',
            ...       xenon.heat_capacity(298.15)
            ...       - xenon.heat_capacity_constant_v(298.15))
            subtracting both hc: 8.314472

        """
        return self.heat_capacity(T) - _R

    def internal_energy(self, T):
        u"""
        Calculate molar internal energy at constant pressure for standard
        state.

        Arguments:
            T (float): temperature.

        Returns:
            numpy_float: The molar internal energy for the compound for a
            given temperature.

        Examples:
            >>> import thermopy
            >>> from thermopy import nasa9polynomials as nasa9
            >>> db = nasa9.Database()
            >>> hcn = db.set_compound('hcn')
            >>> print(hcn, '-', hcn.internal_energy(500))
            formonitrile: HCN - 136809.830971

        """
        return self.enthalpy(T) - _R * T

class Database(object):
    u"""
    Nasa 9 term polynomials database (see NASA/TP—2002-211556).

    The preferred method for identifying compounds is via *usual name* followed
    by an aggregation state if the compound is not a gas. E.g. '(L)', '(cr)',
    '(a)', '(b)' where '(a)' and '(b)' are for allotropes.
    Other methods are:
        1. InChIKey.
        2. CAS number.
        3. IUPAC name.

    Examples:
        >>> from thermopy import nasa9polynomials as nasa9
        >>> db = nasa9.Database()

    """

    def __init__(self):
        u"""Initializes the database."""
        xmlPath = os.path.join(DATADIR, 'databases/nasa9polynomials.xml')
        self._nasa9 = ET.parse(os.path.abspath(xmlPath))
        self._root = self._nasa9.getroot()

    def _search_database(self, x):
        u"""
        Search the database.

        Arguments:
            x (str): identifier for compound being searched.

        Returns:
            tuple: (inp file name, iupac name, ET.Element).

        Note:
            The preferred method for identifying compounds is via *usual name*
            followed by an aggregation state if the compound is not a gas. E.g.
            '(L)', '(cr)', '(a)', '(b)' where '(a)' and '(b)' are for
            allotropes.
            Other methods are:
                1. InChIKey.
                2. CAS number.
                3. IUPAC name.

        """
        result_list = []
        inchikey_re = re.compile('[A-Z]{14}-[A-Z]{10}-[A-Z]')
        cas_re = re.compile('[0-9]{2,7}-[0-9][0-9]-[0-9]')
        # InChIKey search
        if re.match(inchikey_re, x):  # is an inchikey
            for specie in self._root:
                identification = specie.find('identification')
                inchikey = identification.find('InChIKey')
                if x == inchikey.text:
                    result_list.append(specie)
        # CAS search
        elif re.match(cas_re, x):
            for specie in self._root:
                identification = specie.find('identification')
                cas_number = identification.find('cas_number')
                if x == cas_number.text:
                    result_list.append(specie)
        # usual name search
        else:
            # tries exact match first
            for specie in self._root:
                iupac_name = specie.find('identification').find('IUPAC_name')
                if (x.lower() == specie.find('identification').find(
                        'IUPAC_name').text.lower() or
                        x.lower() == specie.attrib['inp_file_name'].lower()):
                    result_list.append(specie)
            if len(result_list) == 1:
                pass
            else:  # exact match was not sucessfull go to loose match
                result_list = []
                # if not found tries loose match
                for specie in self._root:
                    iupac_name = specie.find('identification').find(
                        'IUPAC_name')
                    augmented_namespace = (specie.attrib['inp_file_name']
                                           + ' ' + specie.find('comment').text
                                           + ' ' + iupac_name.text)
                    if x.lower() in augmented_namespace.lower():
                        result_list.append(specie)
        for specie in result_list:
            return [(y.attrib['inp_file_name'],
                     y.find('identification').find('IUPAC_name').text,
                     y) for y in result_list]

    def list_compound(self, x):
        u"""
        List the compounds for a given input.
        It is intended to be used in interactive mode.

        Arguments:
            x (str): identifier for compound being searched.

        Returns:
            list: list of tuples containing (str, str) being the 'inp name' and
            the IUPAC name respectively.

        Examples:
            >>> from thermopy import nasa9polynomials as nasa9
            >>> db = nasa9.Database()
            >>> for i in db.list_compound('fes'):
            ...     print(i)
            ('FeS(a)', 'sulfanylideneiron')
            ('FeS(b)', 'sulfanylideneiron')
            ('FeS(c)', 'sulfanylideneiron')
            ('FeS(L)', 'sulfanylideneiron')
            ('FeSO4(cr)', 'iron(2+);sulfate')
            ('FeS2(cr)', 'N/A')

        Note:
            The preferred method for identifying compounds is via *usual name*
            followed by an aggregation state if the compound is not a gas. E.g.
            '(L)', '(cr)', '(a)', '(b)' where '(a)' and '(b)' are for
            allotropes.
            Other methods are:
                1. InChIKey.
                2. CAS number.
                3. IUPAC name.

        """
        result_list = []
        for i in self._search_database(x):
            result_list.append((i[0], i[1]))
        return result_list

    def set_compound(self, x):
        u"""
        Set the compound if there is one entry specified on the database.

        It is important to notice that due to the nature of the work of this
        database, compounds are gases unless explicitly stated otherwise.

        Arguments:
            x (str): identifier for compound being searched.

        Returns:
            Compound: returns a Compound object if the phase is condensed.
            Returns a CompoundIdealGas otherwise.

        Example:
            >>> # Someone is looking for the element gallium but is not certain
            >>> # how to instantiate it. One would first list the compounds:
            >>> from thermopy import nasa9polynomials as nasa9
            >>> db = nasa9.Database()
            >>> for i in db.list_compound('gallium'):
            ...     print(i)
            ...
            ('Ga', 'gallium')
            ('Ga+', 'gallium')
            ('GaBr', 'bromogallium')
            ('GaBr2', 'dibromogallium')
            ('GaCl', 'chlorogallium')
            ('GaCl2', 'gallium;dichloride')
            ('GaF2', 'difluorogallium')
            ('GaI2', 'diiodogallium')
            ('GaO', 'oxogallium')
            ('GaOH', 'gallium;hydroxide')
            ('Ga2Cl4', 'gallium;gallium;tetrachloride')
            ('Ga2O', 'gallium;oxygen(2-)')
            ('Ga(cr)', 'gallium')
            ('Ga(L)', 'gallium')
            ('Ga2O3(cr)', 'digallium;oxygen(2-)')
            ('Ga2O3(L)', 'digallium;oxygen(2-)')
            >>> gallium = db.set_compound('Ga')
            >>> print(gallium)
            gallium: Ga

        Note:
            The preferred method for identifying compounds is via *usual name*
            followed by an aggregation state if the compound is not a gas. E.g.
            '(L)', '(cr)', '(a)', '(b)' where '(a)' and '(b)' are for
            allotropes.
            Other methods are:
                1. InChIKey.
                2. CAS number.
                3. IUPAC name.

        """
        result = self._search_database(x)
        if len(result) != 1:  # could not set component: give error messages
            if len(result) == 0:
                raise Exception('No compound found.')
            else:
                raise Exception('The compound \'' + str(x) + '\' you are '
                                'trying to set is not unique: ' + result[0][0],
                                result[1][0])
        if result[0][2].find('condensed') == 'True':
            return Compound(result[0][2])
        else:  # if it is an ideal gas
            return CompoundIdealGas(result[0][2])

class Reaction(object):
    u"""
    Model of a chemical reaction.

    Model of a chemical reaction using the NASA9 Compounds.

    Inherits from object.

    Methods:
        enthalpy_reaction: enthalpy of the reaction.
        entropy_reaction: entropy of the reaction.
        gibbs_energy_reaction: Gibbs energy of the reaction.
        equilibrium_constant: equilibrium constant of the reaction.

    Examples:
        >>> from thermopy import nasa9polynomials as nasa9
        >>> db = nasa9.Database()
        >>> na = db.set_compound('na(cr)')
        >>> # being careful to initilize solid compounds
        >>> water = db.set_compound('h2o(l)')
        >>> # being careful to initilize liquid compounds
        >>> sodium_hydroxide = db.set_compound('naoh(a)')
        >>> # being careful to initilize solid compounds
        >>> hydrogen = db.set_compound('h2')
        >>> reacts = (na, water)
        >>> prods = (sodium_hydroxide, hydrogen)
        >>> reacts_coefs = (2, 2)
        >>> prods_coefs = (2, 1)
        >>> reaction1 = nasa9.Reaction(300, reacts, prods, reacts_coefs,
        >>> prods_coefs)
        >>> print(reaction1)
        <reaction> +2 Na(cr) +2 H2O(L)  -> +2 NaOH(a) +1 H2
        >>> print(reaction1.entropy_reaction())
        149.097547531
        >>> print(reaction1.enthalpy_reaction())
        -279857.367433

    Notes:
        The Reaction class does not check for imbalances of the reaction (yet).

    """

    def __init__(self, T, reactants, products,
                 reactants_coefficients, product_coefficients):
        u"""
        Initializes a Reaction object.

        Arguments:
            T (float): temperature of the reaction.
            reactants (tuple): tuple of the reactants as Compounds.
            products (tuple): tuple of the products as Compounds.
            reactants_coefficients (tuple): tuple of the reactants
                coefficients.
            product_coefficients (tuple): tuple of the products coefficients.

        Examples:
            >>> import thermopy
            >>> from thermopy import nasa9polynomials as nasa9
            >>> db = nasa9.Database()
            >>> co = db.set_compound('carbon monoxide')
            >>> molybdenium_oxide = db.set_compound('MoO2(cr)')  # to set it a
            >>> solid
            >>> co2 = db.set_compound('co2')
            >>> molybdenium = db.set_compound('Mo(cr)')  # to set it a solid
            >>> reagents = (co, molybdenium_oxide)
            >>> products = (co2, molybdenium)
            >>> reactants_stoichometry = (2, 1)
            >>> prodcuts_stoichometry = (2, 1)
            >>> reaction1 = nasa9.Reaction(
            ...     298,
            ...     reagents,
            ...     products,
            ...     reactants_stoichometry,
            ...     prodcuts_stoichometry
            ...     )
            >>> print(reaction1, reaction1.enthalpy_reaction())
            <reaction> +2 CO +1 MoO2(cr)  -> +2 CO2 +1 Mo(cr)  23352.3949968

        Note:
            The tuple of the reactants and its coefficients should refer to the
            same compounds (follow the same order). See example.

        """
        self.T = T
        self._reactants = reactants
        self._products = products
        self._rcoefs = tuple(abs(z) for z in reactants_coefficients)
        self._pcoefs = tuple(abs(z) for z in product_coefficients)
        # error checking
        if (len(self._reactants) != len(self._rcoefs) or
                len(self._products) != len(self._pcoefs)):
            raise Exception('Number of reactants or products is different'
                            'from the number of coefficients given')

    def enthalpy_reaction(self, T=None):
        u"""
        Calculate the enthalpy of the reaction at the standard state.

        Arguments:
            T (float): temperature.

        Returns:
            numpy_float: The enthalpy of the reaction for a given temperature.

        Examples:
            >>> import thermopy
            >>> from thermopy import nasa9polynomials as nasa9
            >>> db = nasa9.Database()
            >>> nitric_acid = db.set_compound('hno3')
            >>> # the liquid phase is not present in the database
            >>> naoh = db.set_compound('naoh(a)')
            >>> sodium_nitrate = db.set_compound('nano3(a)')
            >>> water = db.set_compound('h2o(l)')
            >>> reagents = (nitric_acid, naoh)
            >>> products = (sodium_nitrate, water)
            >>> reactants_stoichometry = (1, 1)
            >>> prodcuts_stoichometry = (1, 1)
            >>> reaction1 = nasa9.Reaction(
            ...     298,
            ...     reagents,
            ...     products,
            ...     reactants_stoichometry,
            ...     prodcuts_stoichometry
            ...     )
            >>> print(reaction1, reaction1.enthalpy_reaction())
            <reaction> +1 HNO3 +1 NaOH(a)  -> +1 NaNO3(a) +1 H2O(L)
            -193773.133358

        """
        if T is not None:
            self.T = T
        deltah = 0
        for (coefficient, compound) in zip(self._rcoefs, self._reactants):
            deltah = deltah - coefficient * compound.enthalpy(self.T)
        for (coefficient, compound) in zip(self._rcoefs, self._products):
            deltah = deltah + coefficient * compound.enthalpy(self.T)
        return deltah

    def entropy_reaction(self, T=None):
        u"""
        Calculate the entropy of the reaction at the standard state.

        Arguments:
            T (float): temperature.

        Returns:
            numpy_float: The entropy of the reaction for a given temperature.

        Examples:
            >>> import thermopy
            >>> from thermopy import nasa9polynomials as nasa9
            >>> db = nasa9.Database()
            >>> nitric_acid = db.set_compound('hno3')
            >>> # the liquid phase is not present in the database
            >>> naoh = db.set_compound('naoh(a)')
            >>> sodium_nitrate = db.set_compound('nano3(a)')
            >>> water = db.set_compound('h2o(l)')
            >>> reagents = (nitric_acid, naoh)
            >>> products = (sodium_nitrate, water)
            >>> reactants_stoichometry = (1, 1)
            >>> prodcuts_stoichometry = (1, 1)
            >>> reaction1 = nasa9.Reaction(
            ...     298,
            ...     reagents,
            ...     products,
            ...     reactants_stoichometry,
            ...     prodcuts_stoichometry
            ...     )
            >>> print(reaction1, reaction1.entropy_reaction())
            <reaction> +1 HNO3 +1 NaOH(a)  -> +1 NaNO3(a) +1 H2O(L)
            -145.797754143

        """
        if T is not None:
            self.T = T
        deltas = 0
        for (coefficient, compound) in zip(self._rcoefs, self._reactants):
            deltas = deltas - coefficient * compound.entropy(self.T)
        for (coefficient, compound) in zip(self._rcoefs, self._products):
            deltas = deltas + coefficient * compound.entropy(self.T)
        return deltas

    def gibbs_energy_reaction(self, T=None):
        u"""
        Calculate the Gibbs energy of the reaction at the standard state.

        Arguments:
            T (float): temperature.

        Returns:
            numpy_float: The Gibbs energy of the reaction for a given
            temperature.

        Examples:
            >>> import thermopy
            >>> from thermopy import nasa9polynomials as nasa9
            >>> db = nasa9.Database()
            >>> pcl5 = db.set_compound('pcl5')
            >>> pcl3 = db.set_compound('pcl3')
            >>> chlorine = db.set_compound('cl2')
            >>> reagents = (pcl5,)
            >>> products = (pcl3, chlorine)
            >>> reactants_stoichometry = (1,)
            >>> prodcuts_stoichometry = (1, 1)
            >>> reaction1 = nasa9.Reaction(
            ...     298,
            ...     reagents,
            ...     products,
            ...     reactants_stoichometry,
            ...     prodcuts_stoichometry
            ...     )
            >>> print(reaction1, reaction1.gibbs_energy_reaction())
            <reaction> +1 PCl5  -> +1 PCl3 +1 Cl2  103038.712535

        """
        if T is not None:
            self.T = T
        deltag = 0
        for (coefficient, compound) in zip(self._rcoefs, self._reactants):
            deltag = deltag - coefficient * compound.gibbs_energy(self.T)
        for (coefficient, compound) in zip(self._rcoefs, self._products):
            deltag = deltag + coefficient * compound.gibbs_energy(self.T)
        return deltag

    def equilibrium_constant(self, T=None):
        u"""
        Calculate the equilibrium constant of the reaction at the standard
        state.

        Definition: K = exp(- deltaG / (R T))

        Arguments:
            T (float): temperature.

        Returns:
            numpy_float: The Gibbs energy of the reaction for a given
            temperature.

        Examples:
            >>> import thermopy
            >>> from thermopy import nasa9polynomials as nasa9
            >>> db = nasa9.Database()
            >>> pcl5 = db.set_compound('pcl5')
            >>> pcl3 = db.set_compound('pcl3')
            >>> chlorine = db.set_compound('cl2')
            >>> reagents = (pcl5,)
            >>> products = (pcl3, chlorine)
            >>> reactants_stoichometry = (1,)
            >>> prodcuts_stoichometry = (1, 1)
            >>> reaction1 = nasa9.Reaction(
            ...     500,
            ...     reagents,
            ...     products,
            ...     reactants_stoichometry,
            ...     prodcuts_stoichometry
            ...     )
            >>> print(reaction1, reaction1.equilibrium_constant())
            <reaction> +1 PCl5  -> +1 PCl3 +1 Cl2  6.39431126134e-13

        """
        if T is not None:
            self.T = T
        return np.exp(-1 * self.gibbs_energy_reaction(self.T) / (
            _R * self.T))

    def __repr__(self):
        u"""Define how a reaction should be print."""
        r = ''
        for (reag, coef) in zip(self._reactants, self._rcoefs):
            r = r + '+' + str(coef) + ' ' + reag.inp_name + ' '
        r = r + ' -> '
        for (reag, coef) in zip(self._products, self._pcoefs):
            r = r + '+' + str(coef) + ' ' + reag.inp_name + ' '
        return """<reaction> {0}""".format(r)


# = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = =
import felis
from astropy.io import fits
import scipy.signal
import pdb
import matplotlib.pyplot as plt
import datetime
import pickle
import sys
import os
import numpy as np
# = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = =
def match_templates2specs(templates,spectra,speczs,picklename,wavewindow=[50.0],wavecen_restframe=[1908.0],
                          vshift=None,min_template_level=1e-4,plotdir=None,plot_allCCresults=False,
                          subtract_spec_median=True,overwrite=False,verbose=True):
    """
    Wrapper around felis cross-correlation template matching, to match a list of spectra with a list of templtes.

    --- INPUT ---
    spectra               fits spectra to find a (cross-correlation) match to template for
    templates             fits templates to correlate with
    speczs                Spectroscopic redshifts to perform cross-correlation in rest-frame (shifting the spectrum).
    subtract_spec_median  Subtract median value of spectrum (approximating the continuum level)
    picklename            Name of pickle file to store final cross-correlation results in
    wavewindow            Window (wavecen_restframe * (1+speczs) +/- wavewindow) to perform template matching over.
    wavecen_restframe     Central rest-frame  wavelength of the region to match
    vshift                If a velcotiy shift is known, provide it here and it will be stored in output (not used)
    min_template_level    The template is interpolated to the wavelength grid of the spectrum and extrapolated
                          beyond it's edges if nescessary. In this extrapolation (assuming the template goes to ~0
                          at the edges), very small values (e.g., <1e-20) can be returned. To set these to 0.0
                          provide a level below which all values in the interpolated template are treated as 0s.
    plotdir               Directory to store plots to
    plot_allCCresults     To plot all the cross-correlation plots, set this to True
    overwrite             Overwrite existing pickle file if it already exists?
    verbose               Toggle verbosity

    --- EXAMPLE OF USE ---
    import felis
    import glob

    specdir  = '/Users/kschmidt/work/MUSE/uvEmissionlineSearch/felis_testing/'
    #specs    = glob.glob(specdir+'uves_felis_mock_MUSEspectrum_noisesigma*3p0.fits')
    specs    = glob.glob(specdir+'uves_felis_mock_MUSEspectrum_noisesigma*.fits')
    speczs   = [3.5]*len(specs)

    tempdir  = '/Users/kschmidt/work/MUSE/uvEmissionlineSearch/felis_testing/'
    #temps    = glob.glob(specdir+'uves_felis_template_CIIIdoublet_sig_0p25_fluxCIII1_4p0_flux*.fits')
    temps    = glob.glob(specdir+'uves_felis_template_CIIIdoublet_*fits')
    temps    = glob.glob(specdir+'uves_felis_template_CIVdoublet_*fits')

    plotdir  = '/Users/kschmidt/work/MUSE/uvEmissionlineSearch/felis_testing/plots_CCresults180615/'
    pickle   = '/Users/kschmidt/work/MUSE/uvEmissionlineSearch/felis_testing/CCresults180615_RENAME_.pkl'

    specs = ['/Volumes/DATABCKUP1/TDOSEextractions/171201_TDOSEextraction/Modelimg/tdose_spectra/tdose_spectrum_candels-cdfs-15_modelimg_0115003085-0115003085.fits']
    speczs = [3.2585198879241943]

    ccdic    = felis.match_templates2specs(temps,specs,speczs,pickle,wavewindow=[60]*len(specs),plotdir=plotdir,wavecen_restframe=[1549.0]*len(specs))

    --- OUTPUT ---

    This wrapper will collect all the cross-correlation results in a main dictionary.
    The dictionary will be returned directly but also saved to disk as a pickled filed.
    This file can be loaded with: felis.load_picklefile(picklefilename)

    The returned dictionary has the following format:

    dictionary.keys()  = the list of spectra that have been crossmatched (input = 'spectra')

    Each entry in the dictionary (dictionary[key]) contains the following entries:
        'wavelengths'               :   The wavelength vector used for the cross-correlation matching of
                                        each of the N templates
        'templatevec'               :   list of each of the N templates matched to spectrum
        'zspec'                     :   Spectroscopic redshift for spectrum
        'zCCmaxvec'                 :   the redshift corresponding to max S/N for each of the N templates matched
        'ccresultsarray_flux'       :   fluc vectors for each of the N templates matched
        'ccresultsarray_variance'   :   Variance vecotrs for each of the N templates matched
        'ccresultsarr_S2N'          :   S/N vector for each of the N templates matched
        'ccresultsarr_chi2'         :   chi squared values for the cross-correlation of each of the N templates matched
        'ccresultsarr_Ngoodent'     :   The number of good pixels used in the cross correlation for each of
                                        the N templates matched
        'S2NCCmaxvec'               :   vector with max(S/N) values for N templates matched
        'continuumlevel'            :   The 'continuum level' of the spectrum removed in the cross-correlation.
                                        Currently the value is simply the median of the spectrum.
        'vshift'                    :   If a velocity shift was provided for the template match this is stored here

    The picklefile can be used to assemble sub-sample results (e.g., S/N cuts) based on the
    template cross-correlations with
        felis.selection_from_picklefile()
    And individual entries can be plotted using
        felis.plot_picklefilecontent()

    """
    ccresultdic = {}

    if verbose: print(' - Starting cross-correlation of the '+str(len(spectra))+' spectra and '+
                      str(len(templates))+' templates')
    startstring = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")
    if verbose: print('   '+startstring+'\n')

    if len(spectra) == 0:
        sys.exit(' No spectra provided')

    if len(templates) == 0:
        sys.exit(' No templates provided')

    for ss, spec in enumerate(spectra):
        # Nwave = pyfits.open(spec)[1].header['NAXIS2']
        spec_namebase = spec.split('/')[-1].split('.fit')[0]
        for tt, temp in enumerate(templates):
            temp_namebase = temp.split('/')[-1].split('.fit')[0]

            wavecenter  = wavecen_restframe[ss]  * (1.0 + speczs[ss])
            waverange   = [wavecenter-wavewindow[ss],wavecenter+wavewindow[ss]]

            wave, ccresults, max_S2N, max_z, continuumlevel = \
                felis.cross_correlate_template(spec,temp,z_restframe=speczs[ss],spec_median_sub=subtract_spec_median,
                                               waverange=waverange,min_template_level=min_template_level,verbose=verbose)

            if tt == 0:
                ccresultsarr_flux      = np.array(ccresults[:,0])
                ccresultsarr_variance  = ccresults[:,1]
                ccresultsarr_S2N       = ccresults[:,2]
                ccresultsarr_chi2      = ccresults[:,3]
                ccresultsarr_Ngoodent  = ccresults[:,4]
                templatevec            = np.array([temp])
                S2NCCmaxvec            = np.array([max_S2N])
                zCCmaxvec              = np.array([max_z])
                #print('-------------------------->'+str(np.max(ccresultsarr_S2N))+'  '+str(max_S2N))
            else:
                ccresultsarr_flux      = np.vstack((ccresultsarr_flux,ccresults[:,0]))
                ccresultsarr_variance  = np.vstack((ccresultsarr_variance,ccresults[:,1]))
                ccresultsarr_S2N       = np.vstack((ccresultsarr_S2N,ccresults[:,2]))
                ccresultsarr_chi2      = np.vstack((ccresultsarr_chi2,ccresults[:,3]))
                ccresultsarr_Ngoodent  = np.vstack((ccresultsarr_Ngoodent,ccresults[:,4]))
                templatevec            = np.append(templatevec,temp)
                S2NCCmaxvec            = np.append(S2NCCmaxvec,max_S2N)
                zCCmaxvec              = np.append(zCCmaxvec,max_z)
                #print('-------------------------->'+str(np.max(ccresultsarr_S2N[tt,:]))+'  '+str(S2NCCmaxvec[tt]))

        ccresultdic[spec]  = {'wavelengths':wave, 'templatevec':templatevec, 'zspec':speczs[ss],
                              'zCCmaxvec':zCCmaxvec, 'ccresultsarray_flux':ccresultsarr_flux,
                              'ccresultsarray_variance':ccresultsarr_variance,'S2NCCmaxvec':S2NCCmaxvec,
                              'ccresultsarr_S2N':ccresultsarr_S2N,'ccresultsarr_chi2':ccresultsarr_chi2,
                              'ccresultsarr_Ngoodent':ccresultsarr_Ngoodent}

        ccresultdic[spec]['continuumlevel'] = continuumlevel
        if vshift is not None:
            ccresultdic[spec]['vshift'] = vshift[ss]

    if verbose: print('\n - Finished cross-correlation of the '+str(len(spectra))+' spectra and '+
                      str(len(templates))+' templates')
    if verbose: print('   '+datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")+'  (started at '+startstring+')')

    if verbose: print(' - Saving dictionary to '+picklename)
    if os.path.isfile(picklename) & (overwrite == False):
        print('\n   FELIS WARNING: The output file '+picklename+'exists but overwrite==False so returning "None" ')
        return None
    else:
        felis.save_dictionary(ccresultdic,picklename)

        if verbose: print(' - Will plot all cross-correlation results in saved dictionary as plot_allCCresults=True')
        if plot_allCCresults:
            for ss, spec in enumerate(spectra):
                spec_namebase = spec.split('/')[-1].split('.fit')[0]
                for tt, temp in enumerate(templates):
                    temp_namebase = temp.split('/')[-1].split('.fit')[0]
                    plotname      = plotdir+spec_namebase+'_CCwith_'+temp_namebase+'.pdf'

                    felis.plot_picklefilecontent([spec],picklename,plottemplates=[tt],showspecerr=False,
                                                 plotnames=[plotname],plotdir=plotdir,verbose=verbose)

        if verbose: print(' - Returning dictioniary with results ')
        loaddic = felis.load_picklefile(picklename)
        return  loaddic
# = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = =
def selection_from_picklefile(picklefile,S2Nmaxrange=[3,1e4],zspecrange=[0,10],voffsetrange=[-1e4,1e4],
                              zspecISzLya=False,verbose=True):
    """
    Function, returning the list of spectra (keys) from the FELIS tempalte match pickled dictionary satisfying
    a set of criteria on SN, redshift, velocity offset, etc.

    --- INPUT ---
    picklefile          The path and name to FELIS output pickelfile to select subset of matches from.
    S2Nmaxrange         The range of S/N template accepted for template matches returned.
    zspecrange          The range of redshifts accepted for template matches returned.
    voffsetrange        The range of velocity offsets in km/s (wrt. to the spectroscopic redshift in dictionary)
                        accepted for template matches returned.
    zspecISzLya         If picklefile contains zLya keyword insead of zspec (was generated before 180912), set this
                        keyword to True to enable proper handling of the dictionary keywords.
    verbose             Toggle the verbosity.

    --- EXAMPLE OF USE ---
    import felis

    picklepath = '/Users/kschmidt/work/MUSE/uvEmissionlineSearch/'
    picklefile = picklepath+'MUSEWideLAEs_CCresults180502_CIV_all_575specX12templates.pkl'

    goodkeys   = felis.selection_from_picklefile(picklefile,S2Nmaxrange=[3,5])

    """
    if verbose: print(' - Loading the picklefile \n   '+picklefile)
    loaddic  = felis.load_picklefile(picklefile)
    goodkeys = []
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    # Ensure compatibility with FELIS output dictionaries from before 180912
    zkey = 'zspec'
    if zspecISzLya: zkey = 'zLya'
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    if verbose: print(' - Looking for keys in the pickled dictionary satisfying the cuts:')
    if verbose: print('       '+zkey+'          :  ['+str(zspecrange[0])+','+str(zspecrange[1])+']')
    if verbose: print('       max(S/N)       :  ['+str(S2Nmaxrange[0])+','+str(S2Nmaxrange[1])+']')
    if verbose: print('       voffset[km/s]  :  ['+str(voffsetrange[0])+','+str(voffsetrange[1])+']')
    for key in loaddic.keys():
        keydic = loaddic[key]

        template, vshift_intr, vshift_match, flux, fluxerr, S2Nmax, Ngoodent, chi2, zspec, zS2Nmax =  \
            felis.getresult4maxS2N(loaddic,key,zspecISzLya=zspecISzLya)

        if ((keydic[zkey] > zspecrange[0]) & (keydic[zkey] < zspecrange[1])) & \
                ((S2Nmax > S2Nmaxrange[0]) & (S2Nmax < S2Nmaxrange[1])) & \
                ((vshift_match > voffsetrange[0]) & (vshift_match < voffsetrange[1])):
            goodkeys.append(key)
    if verbose: print(' - Found '+str(len(goodkeys))+' keys in the pickled dictionary satisfying the cuts:')
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    if verbose: print(' - Returning those')
    return goodkeys
# = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = =
def plot_picklefilecontent(specs2plot,picklefile,plotnames=None,plotdir=None,z_restframe=None,
                           plottemplates=None,showspecerr=True,zspecISzLya=False,verbose=True):
    """
    Function to plot individual results from a picklefile generated with felis.match_templates2specs()

    --- INPUT ---
    specs2plot          The spectra from the pickle file (the pickle dictinary keys) to plot. These can be
                        provide in a list by hand or be selected with felis.selection_from_picklefile()
    picklefile          The path and name to pickelfile to plot content of
    plotnames           The names of the plot(s) to generate. If 'None' the string '_templatematch.pdf' will
                        be appended that pickle dictionary key (spectrum name).
    plotdir             To chose a different directory for saving the plots (other than the directory
                        in which the spectrum provided is stored) provide this here.
    z_restframe         The redshift used to move spectrum to rest-frame
    plottemplates       By default the template match with highest S/N is plotted. To plot another template
                        provide the entry in the dictionary of this template. To see the templates, look at
                        the 'templatevec' entry:
                            picload = felis.load_picklefile(picklefile)
                            picload[specname]['templatevec']
    showspecerr         Show the error on the data spectrum? Can make the automatically set y-axis range less ideal.
    zspecISzLya         If picklefile contains zLya keyword insead of zspec (was generated before 180912), set this
                        keyword to True to enable proper handling of the dictionary keywords.
    verbose             Toggle the verbosity.

    --- EXAMPLE OF USE ---
    import felis

    picklepath = '/Users/kschmidt/work/MUSE/uvEmissionlineSearch/'
    picklefile = picklepath+'MUSEWideLAEs_CCresults180502_CIV_all_575specX12templates.pkl'

    picklefile = '/Users/kschmidt/work/MUSE/uvEmissionlineSearch/felis_testing/CCresults180615_RENAME_.pkl'

    plotdir    = '/Users/kschmidt/work/MUSE/uvEmissionlineSearch/MUSEwideLAE_FELISplots/'
    specs2plot = ['/Volumes/DATABCKUP1/TDOSEextractions/171201_TDOSEextraction/Modelimg/tdose_spectra/tdose_spectrum_candels-cdfs-15_modelimg_0115003085-0115003085.fits']

    felis.plot_picklefilecontent(specs2plot,picklefile,plotnames=None,plotdir=plotdir,verbose=True)

    """
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    # Ensure compatibility with FELIS output dictionaries from before 180912
    zkey = 'zspec'
    if zspecISzLya: zkey = 'zLya'
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    if verbose: print(' - loading the dictionary from the pickled file:\n   '+picklefile)
    loaddic  = felis.load_picklefile(picklefile)

    for ss, spec in enumerate(specs2plot):
        if plotnames is None:
            plotname = spec.replace('.fits','_templatematch.pdf')
        else:
            plotname = plotnames[ss]

        if plotdir is not None:
            specdir  = '/'.join(plotname.split('/')[:-1])
            plotname = plotname.replace(specdir,plotdir)

        spec_wave, spec_flux, spec_df, spec_s2n = felis.load_spectrum(spec,verbose=verbose)
        spec_dic  = loaddic[spec]

        if plottemplates is None:
            # Getting the entry of the template with maximum S/N in the CC results
            max_S2N = np.max(spec_dic['S2NCCmaxvec'])
            besttemplate_ent = np.where(spec_dic['ccresultsarr_S2N'] == max_S2N)[0][0] # Template of max S/N
        else:
            max_S2N          = spec_dic['S2NCCmaxvec'][plottemplates[ss]]
            besttemplate_ent = plottemplates[ss]

        # moving spectrum to restframe
        if z_restframe is None:
            z_spec=spec_dic[zkey]

            spec_wave, spec_flux, spec_df, spec_s2n = \
                spec_wave / (1+z_spec), spec_flux * (1.0+z_spec), spec_df * (1.0+z_spec), spec_s2n
        else:
            z_spec = 0.0

        # subtract continuum level from spectrum
        spec_flux = spec_flux - spec_dic['continuumlevel'] * (1.0+z_spec)

        # Limit spectrum to range of wavelengths cross-correlated with tempalte
        goodent = np.where( (spec_wave >= spec_dic['wavelengths'][0]) & (spec_wave <= spec_dic['wavelengths'][-1]) )[0]
        spec_wave, spec_flux, spec_df, spec_s2n = \
            spec_wave[goodent], spec_flux[goodent], spec_df[goodent], spec_s2n[goodent]

        template         = spec_dic['templatevec'][besttemplate_ent]
        max_z            = spec_dic['zCCmaxvec'][besttemplate_ent]

        t_wave_init, t_flux_init, t_df_init, t_s2n_init = felis.load_spectrum(template,verbose=verbose)
        func       = scipy.interpolate.interp1d(t_wave_init,t_flux_init,kind='linear',fill_value="extrapolate")
        t_flux     = func(spec_wave)

        # Getting the entry in the CC flux scalings vector for the given template where S/N is max
        max_S2N_ent      = np.where(spec_dic['ccresultsarr_S2N'][besttemplate_ent,:] == max_S2N)[0][0]

        Npix = len(spec_flux)
        template_triplelength            = np.zeros(3*Npix)
        template_triplelength[0:Npix]    = t_flux
        template_shift_S2Nmax            = np.roll(template_triplelength, int(max_S2N_ent+np.floor(Npix/2.)))[Npix:-Npix]
        template_shift_S2Nmax_normalized = template_shift_S2Nmax/np.trapz(template_shift_S2Nmax,spec_wave)

        flux_scale_S2Nmax     = spec_dic['ccresultsarray_flux'][besttemplate_ent,max_S2N_ent]
        max_wave              = spec_dic['wavelengths'][max_S2N_ent]

        if verbose: print(' - Setting up and generating plot:\n   '+plotname)
        fig = plt.figure(figsize=(6, 7))
        fig.subplots_adjust(wspace=0.1, hspace=0.5,left=0.1, right=0.99, bottom=0.07, top=0.91)
        Fsize    = 9
        lthick   = 2
        marksize = 4
        plt.rc('text', usetex=True)
        plt.rc('font', family='serif',size=Fsize)
        plt.rc('xtick', labelsize=Fsize)
        plt.rc('ytick', labelsize=Fsize)
        plt.clf()
        plt.ioff()

        # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
        ax = plt.subplot(4,1,1)

        plt.plot(spec_wave,template_shift_S2Nmax_normalized,'g.',lw=lthick+1, markersize=marksize,alpha=1.0,
                 label='Temp. at z(S/N$_\\textrm{max}$) = '+str(max_z))
        plt.plot(spec_wave,flux_scale_S2Nmax * template_shift_S2Nmax_normalized,
                 'g-',lw=lthick+2, markersize=marksize,alpha=1.0,
                 label='Temp. flux$_\\textrm{tot}$ scale $\\alpha$ = '+str("%.4f" % flux_scale_S2Nmax)+'')

        Ftot_trapz = np.trapz(flux_scale_S2Nmax * template_shift_S2Nmax_normalized,spec_wave)

        if spec_dic['continuumlevel'] != 0.0:
            plt.plot(spec_wave,spec_flux+spec_dic['continuumlevel']*(1+z_spec),'k-',lw=lthick,
                     markersize=marksize,alpha=0.5, label='Spectrum (w/ cont.)')
            speclabel = 'Spectrum (w/o cont.)'
        else:
            speclabel = 'Spectrum (w/ cont.)'
        plt.plot(spec_wave,spec_flux,'k-',lw=lthick, markersize=marksize,alpha=1.0,label=speclabel)

        #what about...
        #  plt.step(specdat['wave'],specdat['flux'],'r',where='mid')
        #  plt.fill_between(specdat['wave'],specdat['flux']-specdat['fluxerror'],specdat['flux']+specdat['fluxerror'],step='mid',color='red',alpha=0.5)

        if showspecerr:
            plt.fill_between(spec_wave, spec_flux-spec_df, spec_flux+spec_df,color='black',alpha=0.2,label='Spectrum err')

            SNlineYrange = [np.min(spec_flux-spec_df),np.max(spec_flux+spec_df)]
        else:
            SNlineYrange = [np.min(spec_flux),np.max(spec_flux)]

        plt.plot([max_wave,max_wave],SNlineYrange,'--r',lw=lthick,
                 markersize=marksize,alpha=1.0,label='S/N$_\\textrm{max}$ = '+str("%.4f" % max_S2N)+'')

        plt.xlabel(' Wavelength [A]')
        plt.ylabel(' Flux ')
        #plt.ylim([-2,6])
        leg = plt.legend(fancybox=True, loc='upper center',prop={'size':Fsize/1.3},ncol=3,numpoints=1,
                         bbox_to_anchor=(0.45, 1.43))  # add the legend
        leg.get_frame().set_alpha(0.7)
        # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
        plt.text(0.45, 1.45, 'Template No. '+str(int(besttemplate_ent+1))+'/'+str(len(spec_dic['templatevec']))+
                 ': '+template.split('/')[-1].replace('_','\_')+
                 ' (Ftot\_trapz = '+str("%.2f" % Ftot_trapz)+')',
                 fontsize=Fsize/1.3, horizontalalignment='center', transform=ax.transAxes)

        # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
        plt.subplot(4,1,2)

        plt.plot(spec_wave,spec_s2n,'k-',lw=lthick, markersize=marksize,alpha=0.5)
        plt.plot(spec_wave,spec_flux/spec_df,'k-',lw=lthick, markersize=marksize,alpha=1.0)

        plt.plot([max_wave,max_wave],[np.min(spec_flux/spec_df),np.max(spec_flux/spec_df)],'--r',lw=lthick,
                 markersize=marksize,alpha=1.0)
        plt.xlabel(' Wavelength [A]')
        plt.ylabel(' Spectrum S/N ')

        # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
        plt.subplot(4,1,3)
        plt.plot(spec_dic['wavelengths'],spec_dic['ccresultsarr_S2N'][besttemplate_ent,:],'-r',lw=lthick, markersize=marksize,alpha=1.0)
        plt.plot([max_wave,max_wave],[np.min(spec_dic['ccresultsarr_S2N'][besttemplate_ent,:]),
                                      np.max(spec_dic['ccresultsarr_S2N'])],'--r',lw=lthick,
                 markersize=marksize,alpha=1.0)
        plt.xlabel(' Wavelength [A]')
        plt.ylabel(' Cross-Correlation S/N')

        # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
        plt.subplot(4,1,4)
        plt.plot(spec_dic['wavelengths'],spec_dic['ccresultsarr_chi2'][besttemplate_ent,:],'-r',lw=lthick, markersize=marksize,alpha=1.0)
        plt.plot([max_wave,max_wave],[np.min(spec_dic['ccresultsarr_chi2'][besttemplate_ent,:]),np.max(spec_dic['ccresultsarr_chi2'][besttemplate_ent,:])],'--r',lw=lthick,
                 markersize=marksize,alpha=1.0)
        plt.xlabel(' Wavelength [A]')
        plt.ylabel(' $\chi^2_\\textrm{min}$')

        # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
        if verbose: print('   Saving plot to '+plotname)
        plt.savefig(plotname)
        plt.clf()
        plt.close('all')

# = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = =
def minimize_chi2(data,data_sigma,template,temp_wave=None,verbose=True):
    """
    Minimizing chi2 for a 1D template matching and returning the corresponding (flux) scaling, alpha.

    Done using that
        chi**2         = Sum_i (d_i - alpha * t_i)**2 / sigma_i**2
    and that dchi**2/dalpha = 0 implies
        alpha          = ( Sum_i d_i*t_i/sigma_i**2 ) / ( Sum_i t_i**2 / sigma_i**2 )
        sigma_alpha**2 = 1 / Sum_ix( t_i**2 / sigma_i**2 ) = alpha_variance

    Here, d is the data (flux from spectrum), sigma is the uncertainty on the data, t is the template,
    alpha is the flux scaling of the template, sigma_alpha is the uncertainty on alpha, and i runs
    over the pixels in the spectrum.

    --- INPUT ---

    data            Data to match template to
    data_sigma      Uncertainty on data, i.e., sqrt(variance)
    template        Normalized template to search for in data. Should be of the same lenght as data
    temp_wave       Provide template wavelengths for normalization check if dlam is not 1.
    verbose         Toggle verbosity

    """
    normlim = 1e-10
    if temp_wave is not None:
        temp_int = np.trapz(template,temp_wave)
    else:
        temp_int = np.trapz(template)

    if temp_int == 0: # all pixels are 0
        pass
    else:
        sumdiff = np.abs(temp_int-1.0)
        if sumdiff > normlim: #checking if provided template is normalized
            raise ValueError('FELIS WARNING: Template is not normalized: |np.trapz(template,dwave)-1.0| = '+
                             str(sumdiff)+' > '+str(normlim))

    if len(data) == len(template):
        Npix = len(data)
    else:
        sys.exit('The length of the data and template should be the same; it is not.')
    if verbose: print(' - Will calculate and minimize chi**2 for data and template of length '+str(Npix))
    goodent  = np.where((data_sigma > 0) & (np.isfinite(data)) & (template != 0))[0]
    Ngoodent = len(goodent)

    if Ngoodent == 0:
        if verbose: print(' - No entries left where data_sigma > 0 & data is finite & template != 0, so returning 0s')
        chi2_min        = 0.0
        alpha           = 0.0
        alpha_variance  = 0.0
        S2N             = 0.0
    else:
        t2err2         = template[goodent]**2 / data_sigma[goodent]**2
        alpha          = np.sum( data[goodent] * template[goodent] / data_sigma[goodent]**2 ) / np.sum(t2err2)
        alpha_variance = 1.0 / np.sum(t2err2)

        S2N            = alpha / np.sqrt(alpha_variance)
        chi2_min       = np.sum( ( data[goodent]-alpha*template[goodent] )**2 / data_sigma[goodent]**2 )

    return alpha, alpha_variance, S2N, chi2_min, Ngoodent
# = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = =
def cross_correlate_template(spectrum,template,z_restframe=None,waverange=None,
                             min_template_level=1e-2,spec_median_sub=False,verbose=True):
    """
    Function to cross-correlate a spectrum with a template using felis.minimize_chi2().
    It can be run for multiple spectra with the wrapper match_templates2specs which also enables
    easy plotting of the results.

    --- INPUT ---
    spectrum            fits spectrum to find a (cross-correlation) match to template for (flux in f_lambda)
    template            fits template to correlate with
    z_restframe         To perform cross-correlation in rest-frame (shifting the spectrum) provide a redshift.
    waverange           To prevent running the cross-correlation on the full spectrum, provide the wavelength
                        range (based on wavelength-vector in spectrum) to restrict the correlation to.
    min_template_level  The template is interpolated to the wavelength grid of the spectrum and extrapolated
                        beyond it's edges if nescessary. In this extrapolation (assuming the template goes to ~0
                        at the edges), very small values (e.g., <1e-20) can be returned. To set these to 0.0
                        provide a level below which all values in the interpolated template are treated as 0s.
    spec_median_sub     If True, the median level (continuum?) of the spectrum is subtracted prior to matching.
    verbose             Toggle verbosity

    --- EXAMPLE OF USE ---
    import felis

    specdir  = '/Users/kschmidt/work/MUSE/uvEmissionlineSearch/felis_testing/'
    spectrum = specdir+'uves_felis_mock_MUSEspectrum_noisesigma1p0.fits'
    spectrum = specdir+'uves_felis_mock_MUSEspectrum_noisesigma0p05.fits'
    template = specdir+'uves_felis_template_CIIIdoublet_sig_0p25_fluxCIII1_2p0_fluxratio_0p5.fits'

    s_wave, ccresults, max_S2N, max_z, continuumlevel = felis.cross_correlate_template(spectrum,template,z_restframe=3.5,waverange=[8560,8610])

    """
    if verbose: print(' - Loading spectrum and template to cross-correlate ')
    s_wave, s_flux, s_df, s_s2n = felis.load_spectrum(spectrum,verbose=verbose)

    if waverange is not None:
        if verbose: print(' - Limiting the wavelength range to cross-correlate over ')
        goodent = np.where( (s_wave > waverange[0]) & (s_wave < waverange[1]) )
        s_wave, s_flux, s_df, s_s2n = s_wave[goodent], s_flux[goodent], s_df[goodent], s_s2n[goodent]

    if verbose: print(' - Estimating continuum from spectrum (median value) in observed units')
    if spec_median_sub:
        continuumval = np.median(s_flux)
    else:
        continuumval = 0.0

    if z_restframe is not None:
        if verbose: print(' - Shifting spectrum to rest-frame using the redshift '+str(z_restframe))
        s_wave, s_flux, s_df, s_s2n = s_wave / (1+z_restframe), s_flux * (1+z_restframe), s_df * (1+z_restframe), s_s2n
    else:
        z_restframe=0.0

    if verbose: print(' - Removing continuum from spectrum ')
    s_flux       = s_flux - continuumval * (1.0+z_restframe)

    if verbose: print(' - Interpolate template to spectrums wavelength resolution before cross-correlating')
    t_wave_init, t_flux_init, t_df_init, t_s2n_init = felis.load_spectrum(template,verbose=verbose)
    func       = scipy.interpolate.interp1d(t_wave_init,t_flux_init,kind='linear',fill_value="extrapolate")
    t_flux     = func(s_wave)
    t_flux[t_flux < min_template_level] = 0.0

    if verbose: print(' - Normalizing total template flux to 1')
    if len(t_flux[t_flux != 0]) == 0:
        if verbose: print('   FELIS WARNING All interpolated template pixels are 0.0')
    else:
        temp_int = np.trapz(t_flux,s_wave)
        t_flux   = t_flux / temp_int

    Npix = len(s_flux)

    template_triplelength = np.zeros(3*Npix)
    template_triplelength[0:Npix] = t_flux

    ccresults = np.zeros([Npix,5])

    for ii in np.arange(Npix):
        rollsize       = int(ii+np.floor(Npix/2.))
        template_shift = np.roll(template_triplelength, rollsize)[Npix:-Npix]

        normlim   = 1e-10
        temp_int  = np.trapz(template_shift,s_wave)
        sumdiff   = np.abs(temp_int-1.0)
        if sumdiff > normlim: # make sure "half-template" shifts are also normalized
            if temp_int == 0: # only attempt normalization if template is not just zeros
                pass
            else:
                template_shift   = template_shift / temp_int

        try:
            flux_scale, flux_scale_variance, S2N, chi2_min, NgoodentChi2 = \
                felis.minimize_chi2(s_flux,s_df,template_shift,temp_wave=s_wave,verbose=False)

            ccresults[ii,:] = flux_scale, flux_scale_variance, S2N, chi2_min, NgoodentChi2
        except:
            print(' ERROR: Problems in minimizing Chi**2 with felis.minimize_chi2() while cross-correlating. '
                  'Stopping for further investigation')

            pdb.set_trace()

    max_S2N     = np.max(ccresults[:,2])
    max_S2N_ent = np.where(ccresults[:,2] == max_S2N)[0][0]
    max_wave    = s_wave[max_S2N_ent]

    dlam     = np.median(np.diff(s_wave))
    t_wave   = np.arange(np.min(t_wave_init),np.max(t_wave_init),dlam)
    t_cenent = int(np.ceil(len(t_wave)/2))
    max_z    = (max_wave*(z_restframe+1.0) / t_wave[t_cenent]) - 1.0

    return  s_wave, ccresults, max_S2N, max_z, continuumval

# = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = =
def load_spectrum(specfile,verbose=True):
    """
    Loading spectrum generated with felis.save_spectrum() or TDOSE.

    --- INPUT ---

    --- EXAMPLE OF USE ---
    import felis
    specfile  = './spectrum_output.fits'
    w, f, df, s2n = felis.load_spectrum(specfile)

    """
    if verbose: print(' - Loading SPEC1D extension (spectrum) of \n   '+specfile)
    dat = fits.open(specfile)['SPEC1D'].data

    wave    = dat['wave']
    flux    = dat['flux']
    fluxerr = dat['fluxerror']
    s2n     = dat['s2n']

    return wave, flux, fluxerr, s2n
# = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = =
def save_spectrum(outfile,wave,flux,fluxerr,headerinfo=None,waveunits='ANGSTROMS',fluxunits='',overwrite=False,verbose=True):
    """
    Function saving a spectrum (or template) to a fits file

    --- INPUT ---
    outfile         Name of output file to generate
    wave            Wavelength in units of Angstrom
    flux            Flux to store to fits file
    fluxerr         Error un flux
    headerinfo      To add info to the header provide it to this keyword as a dictionary on the format:
                       headerinfo = {'KEYNAME1':[VALUE1,INFOCOMMENT1], 'KEYNAME2':[VALUE2,INFOCOMMENT2]}
    waveunits       Units of wave vector to store in fits header
    fluxunits       Units of flux to store in fits header
    overwrite       Overwrite existing file?
    verbose         Toggle verbosity

    --- EXAMPLE OF USE ---
    import felis
    import numpy as np
    wave        = np.arange(1,100,1)
    flux        = wave**0.5 + (wave*0.0+2.3)
    fluxerr     = np.sqrt(flux)
    headerinfo  = {'FEL_EL1':[1548,'Wave of FELIS line 1 in template'], 'FEL_EL2':[1551,'Wave of FELIS line 2 in template']}
    outfile     = './spectrum_output.fits'

    felis.save_spectrum(outfile,wave,flux,fluxerr,headerinfo=headerinfo,overwrite=False,verbose=True)

    """
    S2N = flux/fluxerr

    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    if verbose: print(' - Saving wavelength and flux values to \n   '+outfile)
    mainHDU = fits.PrimaryHDU()       # primary HDU
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    c1 = fits.Column(name='wave',      format='D', unit=waveunits, array=wave)
    c2 = fits.Column(name='flux',      format='D', unit=fluxunits, array=flux)
    c3 = fits.Column(name='fluxerror', format='D', unit=fluxunits, array=fluxerr)
    c4 = fits.Column(name='s2n',       format='D', unit='', array=S2N)

    coldefs = fits.ColDefs([c1,c2,c3,c4])
    tbHDU   = fits.BinTableHDU.from_columns(coldefs) # creating default header

    # writing hdrkeys:'---KEY--',                             '----------------MAX LENGTH COMMENT-------------'
    tbHDU.header.append(('EXTNAME ','SPEC1D'                     ,'cube containing source'),end=True)
    if headerinfo is not None:
        for key in headerinfo.keys():
            tbHDU.header.append((key,headerinfo[key][0],headerinfo[key][1]),end=True)

    hdulist = fits.HDUList([mainHDU,tbHDU])
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

    hdulist.writeto(outfile, overwrite=overwrite)
# = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = =
def save_dictionary(dictionary, output='./saveddictionary_RENAME_.pkl'):
    with open(output, 'wb') as f:
        pickle.dump(dictionary, f, pickle.HIGHEST_PROTOCOL)
# = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = =
def load_picklefile(picklefile):
    pversion = sys.version_info[0]
    with open(picklefile, 'rb') as f:
        if pversion == 2:
            return pickle.load(f) # for python 2.X
        elif pversion == 3:
            return pickle.load(f, encoding='latin1') # for python 3.X
        else:
            sys.exit('   felis.load_picklefile() saw unknown version of python: version = '+str(pversion))
# = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = =
def getresult4maxS2N(dictionary,dictionarykey,zspecISzLya=False):
    """
    Function returning results at maxmimum S/N for the template with maximum S/N overall, for a given
    dictionary keyword of the pickled FELIS output.

    --- INPUT ---
    dictionary      = Dictionary containing FELIS outputs, i.e. the output from felis.load_picklefile()
    dictionarykey   = The object/spectrum (dictionary key) to return results for
    zspecISzLya         If picklefile contains zLya keyword insead of zspec (was generated before 180912), set this
                        keyword to True to enable proper handling of the dictionary keywords.

    """
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    # Ensure compatibility with FELIS output dictionaries from before 180912
    zkey = 'zspec'
    if zspecISzLya: zkey = 'zLya'
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    keys = dictionary.keys()
    if dictionarykey in keys:
        keydic        = dictionary[dictionarykey]
        S2Nmaxvec     = keydic['S2NCCmaxvec']
        S2Nmax        = np.max(S2Nmaxvec)
        templateent   = np.where(S2Nmaxvec == S2Nmax)[0]
        if len(templateent) > 1:
            print('    FELIS WARNING felis.getresult4maxS2N(): \n    Multiple templates ('+str(len(templateent))+' of '+
                  str(len(S2Nmaxvec))+' templates) share the same maximum S/N ('+str(S2Nmax)+
                  ') of cross-correlation for \n   '+dictionarykey)
            print('   Will return the results for the first template')
            templateent = np.array([templateent[0]])

        template      = keydic['templatevec'][templateent][0]
        zS2Nmax       = keydic['zCCmaxvec'][templateent][0]
        zspec         = keydic[zkey]
        vshift_match  = 299792.458 * (zspec - zS2Nmax)/(1.0+zS2Nmax) # cf. Erb+2014

        if 'vshift' in keydic.keys():
            vshift_intr   = keydic['vshift'] # velocity shift stored in dictionary with felis.match_templates2specs()
        else:
            vshift_intr   = -99

        S2Nvec        = keydic['ccresultsarr_S2N'][templateent][0]
        S2Nvecmax     = np.where(S2Nvec == np.max(S2Nvec))[0]

        flux          = keydic['ccresultsarray_flux'][templateent,S2Nvecmax][0]
        fluxerr       = np.sqrt(keydic['ccresultsarray_variance'][templateent,S2Nvecmax][0])
        Ngoodent      = keydic['ccresultsarr_Ngoodent'][templateent,S2Nvecmax][0]
        chi2          = keydic['ccresultsarr_chi2'][templateent,S2Nvecmax][0]
    else:
        template      = 'None'
        vshift_intr   = -99.0
        vshift_match  = -99.0
        S2Nmax        = -99.0
        flux          = -99.0
        fluxerr       = -99.0
        Ngoodent      = -99.0
        chi2          = -99.0

    return template, vshift_intr, vshift_match, flux, fluxerr, S2Nmax, Ngoodent, chi2, zspec, zS2Nmax
# = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = =