<commit_before>#!/usr/bin/env python
# encoding: utf-8

################################################################################
#
#   RMG - Reaction Mechanism Generator
#
#   Copyright (c) 2009-2011 by the RMG Team (<EMAIL>)
#
#   Permission is hereby granted, free of charge, to any person obtaining a
#   copy of this software and associated documentation files (the 'Software'),
#   to deal in the Software without restriction, including without limitation
#   the rights to use, copy, modify, merge, publish, distribute, sublicense,
#   and/or sell copies of the Software, and to permit persons to whom the
#   Software is furnished to do so, subject to the following conditions:
#
#   The above copyright notice and this permission notice shall be included in
#   all copies or substantial portions of the Software.
#
#   THE SOFTWARE IS PROVIDED 'AS IS', WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
#   IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
#   FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
#   AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
#   LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
#   FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
#   DEALINGS IN THE SOFTWARE.
#
################################################################################

"""
This module contains functions for writing of Chemkin input files.
"""
import shutil
import math
import re
import logging
import textwrap
import os.path
import numpy

import rmgpy.kinetics as _kinetics
from rmgpy.reaction import Reaction
#from species import Species
from rmgpy.rmg.model import Species
from rmgpy.rmg.pdep import PDepReaction
from rmgpy.thermo import NASAPolynomial, NASA
import rmgpy.constants as constants
from rmgpy.quantity import Quantity
from rmgpy.data.base import Entry 
from rmgpy.data.kinetics.library import LibraryReaction
from rmgpy.data.kinetics.family import TemplateReaction
from rmgpy.rmg.pdep import PDepNetwork
from rmgpy.molecule import Molecule
from rmgpy.transport import TransportData

__chemkin_reaction_count = None
    
from rmgpy.util import makeOutputSubdirectory

################################################################################

class ChemkinError(Exception):
    """
    An exception class for exceptional behavior involving Chemkin files. Pass a
    string describing the circumstances that caused the exceptional behavior.
    """
    pass

################################################################################

def Ffloat(string):
    """
    Parse a Fortran-ish string into a float, like "1.00D 03" or "1.00d+03"
    """
    return float(string.replace("e", "E").replace("d", "D").replace("D", "E").replace("E ", "E+"))
    
def readThermoEntry(entry, Tmin=0, Tint=0, Tmax=0):
    """
    Read a thermodynamics `entry` for one species in a Chemkin file. Returns
    the label of the species and the thermodynamics model as a :class:`NASA`
    object.
    
    Format specification at http://www2.galcit.caltech.edu/EDL/public/formats/chemkin.html
    """
    lines = entry.splitlines()
    species = str(lines[0][0:18].split()[0].strip())
    
    comment = lines[0][len(species):24].strip()
    formula = {}
    for i in [24, 29, 34, 39, 74]:
        element,count = lines[0][i:i+2].strip(), lines[0][i+2:i+5].strip()
        if element:
            try:
                count = int(count)
            except ValueError:
                # Chemkin allows float values for the number of atoms, so try this next.
                try:
                    count = int(float(count))
                except ValueError:
                    logging.info("Trouble reading line '{0}' element segment '{1}'".format(lines[0].strip(),lines[0][i:i+5]))
                    if count == '' and re.match('\.?0*', element):
                        if i == 74: logging.info("Assuming it's spillover from Tint, and ignoring.")
                        else: logging.warning("Assuming it's not meant to be there, although it would be good to fix the chemkin file.")
                        count = 0
                    else:
                        raise
            if count != 0: # Some people put garbage elements in, with zero count. Ignore these. Allow negative counts though (eg. negative one electron)
                formula[element] = count
    phase = lines[0][44]
    if phase.upper() != 'G':
        logging.warning("Was expecting gas phase thermo data for {0}. Skipping thermo data.".format(species))
        return species, None, None
    
    # Extract the NASA polynomial coefficients
    # Remember that the high-T polynomial comes first!
    try:
        try:
            Tmin = float(lines[0][45:55].strip())
        except ValueError:
            pass
        try:
            Tmax = float(lines[0][55:65].strip())
        except ValueError:
            pass
        try:
            Tint = float(lines[0][65:73].strip())
        except ValueError:
            pass
        a0_high = Ffloat(lines[1][0:15].strip())
        a1_high = Ffloat(lines[1][15:30].strip())
        a2_high = Ffloat(lines[1][30:45].strip())
        a3_high = Ffloat(lines[1][45:60].strip())
        a4_high = Ffloat(lines[1][60:75].strip())
    
        a5_high = Ffloat(lines[2][0:15].strip())
        a6_high = Ffloat(lines[2][15:30].strip())
        a0_low = Ffloat(lines[2][30:45].strip())
        a1_low = Ffloat(lines[2][45:60].strip())
        a2_low = Ffloat(lines[2][60:75].strip())
    
        a3_low = Ffloat(lines[3][0:15].strip())
        a4_low = Ffloat(lines[3][15:30].strip())
        a5_low = Ffloat(lines[3][30:45].strip())
        a6_low = Ffloat(lines[3][45:60].strip())
    except (IndexError, ValueError), e:
        logging.warning('Error while reading thermo entry for species {0}'.format(species))
        logging.warning(e.message)
        return species, None, None
    
    # Construct and return the thermodynamics model
    thermo = NASA(
        polynomials = [
            NASAPolynomial(Tmin=(Tmin,"K"), Tmax=(Tint,"K"), coeffs=[a0_low, a1_low, a2_low, a3_low, a4_low, a5_low, a6_low]),
            NASAPolynomial(Tmin=(Tint,"K"), Tmax=(Tmax,"K"), coeffs=[a0_high, a1_high, a2_high, a3_high, a4_high, a5_high, a6_high])
        ],
        Tmin = (Tmin,"K"),
        Tmax = (Tmax,"K"),
    )
    if comment:
        thermo.comment = comment

    return species, thermo, formula

################################################################################

def readKineticsEntry(entry, speciesDict, Aunits, Eunits):
    """
    Read a kinetics `entry` for a single reaction as loaded from a Chemkin
    file. The associated mapping of labels to species `speciesDict` should also
    be provided. Returns a :class:`Reaction` object with the reaction and its
    associated kinetics.
    """
    Afactor = {
        'cm^3/(mol*s)': 1.0e6,
        'cm^3/(molecule*s)': 1.0e6,
        'm^3/(mol*s)': 1.0,
        'm^3/(molecule*s)': 1.0,
    }[Aunits[2]]
    
    lines = entry.strip().splitlines()
    
    # The first line contains the reaction equation and a set of
    # modified Arrhenius parameters
    reaction, thirdBody, kinetics, kunits, klow_units = _readKineticsReaction(
        line=lines[0], speciesDict=speciesDict, Aunits=Aunits, Eunits=Eunits)

    if len(lines) == 1 and not thirdBody:
        # If there's only one line then we know to use the high-P limit kinetics as-is
        reaction.kinetics = kinetics['arrhenius high']
    else:
        # There's more kinetics information to be read
        kinetics.update({
            'chebyshev coefficients': [],
            'efficiencies': {},
        })

        # Note that the subsequent lines could be in any order
        for line in lines[1:]:
            kinetics = _readKineticsLine(
                line=line, reaction=reaction, speciesDict=speciesDict, Eunits=Eunits,
                kunits=kunits, klow_units=klow_units,
                kinetics=kinetics)

        # Decide which kinetics to keep and store them on the reaction object
        # Only one of these should be true at a time!
        if 'chebyshev' in kinetics:
            chebyshev = kinetics['chebyshev']
            if chebyshev.Tmin is None or chebyshev.Tmax is None:
                raise ChemkinError('Missing TCHEB line for reaction {0}'.format(reaction))
            if chebyshev.Pmin is None or chebyshev.Pmax is None:
                raise ChemkinError('Missing PCHEB line for reaction {0}'.format(reaction))
            index = 0
            for t in range(chebyshev.degreeT):
                for p in range(chebyshev.degreeP):
                    chebyshev.coeffs.value_si[t,p] = kinetics[
                        'chebyshev coefficients'][index]
                    index += 1
            # Don't forget to convert the Chebyshev coefficients to SI units!
            # This assumes that s^-1, cm^3/mol*s, etc. are compulsory
            chebyshev.coeffs.value_si[0,0] -= (len(reaction.reactants) - 1) * math.log10(Afactor)
            reaction.kinetics = chebyshev
        elif 'pressure-dependent arrhenius' in kinetics:
            pdepArrhenius = kinetics['pressure-dependent arrhenius']
            # Check for duplicates and combine them to MultiArrhenius objects
            duplicatesToRemove = []
            duplicatesToAdd = []
            for index1 in range(len(pdepArrhenius)):
                reaction1 = pdepArrhenius[index1]
                P1, kinetics1 = reaction1
                if reaction1 in duplicatesToRemove:
                    continue
                for index2 in range(index1+1, len(pdepArrhenius)):
                    reaction2 = pdepArrhenius[index2]
                    P2, kinetics2 = reaction2
                    if P1 == P2:
                        if reaction1 not in duplicatesToRemove:
                            new_kinetics = _kinetics.MultiArrhenius()
                            duplicatesToAdd.append([P1,new_kinetics])
                            new_kinetics.arrhenius = [kinetics1]
                            duplicatesToRemove.append(reaction1)
                        new_kinetics.arrhenius.append(kinetics2)
                        duplicatesToRemove.append(reaction2)
            for item in duplicatesToRemove:
                pdepArrhenius.remove(item)
            pdepArrhenius.extend(duplicatesToAdd)
            
            pdepArrhenius = sorted(pdepArrhenius, key=lambda reaction: reaction[0])  # sort by ascending pressures

            reaction.kinetics = _kinetics.PDepArrhenius(
                pressures = ([P for P, arrh in pdepArrhenius],"atm"),
                arrhenius = [arrh for P, arrh in pdepArrhenius],
            )
        elif 'troe' in kinetics:
            troe = kinetics['troe']
            troe.arrheniusHigh = kinetics['arrhenius high']
            troe.arrheniusLow = kinetics['arrhenius low']
            troe.efficiencies = kinetics['efficiencies']
            reaction.kinetics = troe
        elif thirdBody:
            reaction.kinetics = _kinetics.ThirdBody(
                arrheniusLow=kinetics['arrhenius low'])
            reaction.kinetics.efficiencies = kinetics['efficiencies']
        elif 'arrhenius low' in kinetics:
            reaction.kinetics = _kinetics.Lindemann(
                arrheniusHigh=kinetics['arrhenius high'],
                arrheniusLow=kinetics['arrhenius low'])
            reaction.kinetics.efficiencies = kinetics['efficiencies']
        elif 'explicit reverse' in kinetics or reaction.duplicate:
            # it's a normal high-P reaction - the extra lines were only either REV (explicit reverse) or DUP (duplicate)
            reaction.kinetics = kinetics['arrhenius high']
        else:
            raise ChemkinError(
                'Unable to understand all additional information lines for reaction {0}.'.format(entry))
        
        # These things may *also* be true
        if 'sri' in kinetics:
            reaction.kinetics.comment += "Warning: SRI parameters from chemkin file ignored on import. "

        if 'explicit reverse' in kinetics:
            reaction.kinetics.comment += (
                "Chemkin file stated explicit reverse rate: {0}"
                ).format(kinetics['explicit reverse'])

    return reaction


def _readKineticsReaction(line, speciesDict, Aunits, Eunits):
    """
    Parse the first line of of a Chemkin reaction entry.
    """
    tokens = line.split()
    
    rmg = True
    try:
        float(tokens[-6])
    except (ValueError, IndexError):
        rmg = False
    AuncertaintyType = '+|-'
    if rmg:
        A = float(tokens[-6])
        n = float(tokens[-5])
        Ea = float(tokens[-4])
        try:
            dA = float(tokens[-3])
        except ValueError:
            AuncertaintyType = '*|/'
            dA = float(tokens[-3][1:])
        dn = float(tokens[-2])
        dEa = float(tokens[-1])
        reaction = ''.join(tokens[:-6])
    else:
        A = Ffloat(tokens[-3])
        n = Ffloat(tokens[-2])
        Ea = Ffloat(tokens[-1])
        dA = 0.0
        dn = 0.0
        dEa = 0.0
        reaction = ''.join(tokens[:-3])
    thirdBody = False
    
    # Split the reaction equation into reactants and products
    reversible = True
    reactants, products = reaction.split('=')
    if '<=>' in reaction:
        reactants = reactants[:-1]
        products = products[1:]
    elif '=>' in reaction:
        products = products[1:]
        reversible = False
    if '(+M)' in reactants: reactants = reactants.replace('(+M)','')
    if '(+m)' in reactants: reactants = reactants.replace('(+m)','')
    if '(+M)' in products:  products = products.replace('(+M)','')
    if '(+m)' in products:  products = products.replace('(+m)','')
    
    # Create a new Reaction object for this reaction
    reaction = Reaction(reactants=[], products=[], reversible=reversible)
    
    # Convert the reactants and products to Species objects using the speciesDict
    for reactant in reactants.split('+'):
        reactant = reactant.strip()
        stoichiometry = 1
        if reactant[0].isdigit():
            # This allows for reactions to be of the form 2A=B+C instead of A+A=B+C
            # The implementation below assumes an integer between 0 and 9, inclusive
            stoichiometry = int(reactant[0])
            reactant = reactant[1:]               
        if reactant.upper() == 'M':
            thirdBody = True
        elif reactant not in speciesDict:
            raise ChemkinError('Unexpected reactant "{0}" in reaction {1}.'.format(reactant, reaction))
        else:
            for i in range(stoichiometry):
                reaction.reactants.append(speciesDict[reactant])
    for product in products.split('+'):
        product = product.strip()
        stoichiometry = 1
        if product[0].isdigit():
            # This allows for reactions to be of the form A+B=2C instead of A+B=C+C
            # The implementation below assumes an integer between 0 and 9, inclusive
            stoichiometry = int(product[0])
            product = product[1:]
        if product.upper() == 'M':
            pass
        elif product not in speciesDict:
            if re.match('[0-9.]+',product):
                logging.warning("Looks like reaction {0!r} has fractional stoichiometry, which RMG cannot handle. Ignoring".format(line))
                raise ChemkinError('Skip reaction!')
            raise ChemkinError('Unexpected product "{0}" in reaction {1}.'.format(product, reaction))
        else:
            for i in range(stoichiometry):
                reaction.products.append(speciesDict[product])
    
    # Determine the appropriate units for k(T) and k(T,P) based on the number of reactants
    # This assumes elementary kinetics for all reactions
    try:
        Nreac = len(reaction.reactants) + (1 if thirdBody else 0)
        kunits = Aunits[Nreac]
        klow_units = Aunits[Nreac+1]
    except IndexError:
        raise ChemkinError('Invalid number of reactant species for reaction {0}.'.format(reaction))
    
    key = 'arrhenius low' if thirdBody else 'arrhenius high'
    
    kinetics = {
        key: _kinetics.Arrhenius(
            A = (A,kunits,AuncertaintyType,dA),
            n = (n,'','+|-',dn),
            Ea = (Ea,Eunits,'+|-',dEa),
            T0 = (1,"K"),
        ),
    }
    return (reaction, thirdBody, kinetics, kunits, klow_units)


def _readKineticsLine(line, reaction, speciesDict, Eunits, kunits, klow_units, kinetics):
    """
    Parse the subsequent lines of of a Chemkin reaction entry.
    """
    case_preserved_tokens = line.split('/')
    line = line.upper()
    tokens = line.split('/')

    if 'DUP' in line:
        # Duplicate reaction
        reaction.duplicate = True

    elif 'LOW' in line:
        # Low-pressure-limit Arrhenius parameters
        tokens = tokens[1].split()
        kinetics['arrhenius low'] = _kinetics.Arrhenius(
            A = (float(tokens[0].strip()),klow_units),
            n = float(tokens[1].strip()),
            Ea = (float(tokens[2].strip()),Eunits),
            T0 = (1,"K"),
        )

    elif 'HIGH' in line:
        # What we thought was high, was in fact low-pressure
        kinetics['arrhenius low'] = kinetics['arrhenius high']
        kinetics['arrhenius low'].A = (
            kinetics['arrhenius low'].A.value, klow_units)
        # High-pressure-limit Arrhenius parameters
        tokens = tokens[1].split()
        kinetics['arrhenius high'] = _kinetics.Arrhenius(
            A = (float(tokens[0].strip()),kunits),
            n = float(tokens[1].strip()),
            Ea = (float(tokens[2].strip()),Eunits),
            T0 = (1,"K"),
        )

    elif 'TROE' in line:
        # Troe falloff parameters
        tokens = tokens[1].split()
        alpha = float(tokens[0].strip())
        T3 = float(tokens[1].strip())
        T1 = float(tokens[2].strip())
        try:
            T2 = float(tokens[3].strip())
        except (IndexError, ValueError):
            T2 = None

        kinetics['troe'] = _kinetics.Troe(
            alpha = alpha,
            T3 = (T3,"K"),
            T1 = (T1,"K"),
            T2 = (T2,"K") if T2 is not None else None,
        )
    elif line.strip().startswith('SRI'):
        kinetics['sri'] = True
        """To define an SRI pressure-dependent reaction, in addition to the LOW or HIGH parameters, the
        keyword SRI followed by three or five parameters must be included in the following order: a, b,
        c, d, and e [Eq. (74)]. The fourth and fifth parameters are options. If only the first three are
        stated, then by default d = 1 and e = 0.
        """
        # see eg. http://www.dipic.unipd.it/faculty/canu/files/Comb/Docs/chemkinCK.pdf

    elif 'CHEB' in line:
        # Chebyshev parameters
        chebyshev = kinetics.get(
            'chebyshev', _kinetics.Chebyshev(kunits=kunits))
        kinetics['chebyshev'] = chebyshev
        tokens = [t.strip() for t in tokens]
        if 'TCHEB' in line:
            index = tokens.index('TCHEB')
            tokens2 = tokens[index+1].split()
            chebyshev.Tmin = Quantity(float(tokens2[0].strip()),"K")
            chebyshev.Tmax = Quantity(float(tokens2[1].strip()),"K")
        if 'PCHEB' in line:
            index = tokens.index('PCHEB')
            tokens2 = tokens[index+1].split()
            chebyshev.Pmin = Quantity(float(tokens2[0].strip()),"atm")
            chebyshev.Pmax = Quantity(float(tokens2[1].strip()),"atm")
        if 'TCHEB' in line or 'PCHEB' in line:
            pass
        elif chebyshev.degreeT == 0 or chebyshev.degreeP == 0:
            tokens2 = tokens[1].split()
            chebyshev.degreeT = int(float(tokens2[0].strip()))
            chebyshev.degreeP = int(float(tokens2[1].strip()))
            chebyshev.coeffs = numpy.zeros((chebyshev.degreeT,chebyshev.degreeP), numpy.float64)
        else:
            tokens2 = tokens[1].split()
            kinetics['chebyshev coefficients'].extend(
                [float(t.strip()) for t in tokens2])

    elif 'PLOG' in line:
        pdepArrhenius = kinetics.get('pressure-dependent arrhenius', [])
        kinetics['pressure-dependent arrhenius'] = pdepArrhenius
        tokens = tokens[1].split()
        pdepArrhenius.append([float(tokens[0].strip()),
            _kinetics.Arrhenius(
                A = (float(tokens[1].strip()),kunits),
                n = float(tokens[2].strip()),
                Ea = (float(tokens[3].strip()),Eunits),
                T0 = (1,"K"),
        )])
    elif tokens[0].startswith('REV'):
        reverseA = float(tokens[1].split()[0])
        kinetics['explicit reverse'] = line.strip()
        if reverseA == 0:
            logging.info("Reverse rate is 0 so making irreversible for reaction {0}".format(reaction))
            reaction.reversible = False
        else:
            logging.info("Ignoring explicit reverse rate for reaction {0}".format(reaction))

    else:
        # Assume a list of collider efficiencies
        try:
            for collider, efficiency in zip(case_preserved_tokens[0::2], case_preserved_tokens[1::2]):
                try:
                    efficiency = float(efficiency.strip())
                except ValueError:
                    error_msg = "{0!r} doesn't look like a collision efficiency for species {1} in line {2!r}".format(efficiency,collider.strip(),line)
                    logging.error(error_msg)
                    raise ChemkinError(error_msg)
                if collider.strip() in speciesDict:
                    kinetics['efficiencies'][speciesDict[collider.strip()].molecule[0]] = efficiency
                else: # try it with capital letters? Not sure whose malformed chemkin files this is needed for.
                    kinetics['efficiencies'][speciesDict[collider.strip().upper()].molecule[0]] = efficiency
        except IndexError:
            error_msg = 'Could not read collider efficiencies for reaction: {0}.\n'.format(reaction)
            error_msg += 'The following line was parsed incorrectly:\n{0}'.format(line)
            error_msg += "\n(Case-preserved tokens: {0!r} )".format(case_preserved_tokens)
            raise ChemkinError(error_msg)
    return kinetics


def readReactionComments(reaction, comments, read = True):
    """
    Parse the `comments` associated with a given `reaction`. If the comments
    come from RMG (Py or Java), parse them and extract the useful information.
    Return the reaction object based on the information parsed from these
    comments. If `read` if False, the reaction is returned as an "Unclassified"
    LibraryReaction.
    """
    
    if read == False:
        # The chemkin file was not generated by either RMG-Py or RMG-Java, thus, there should be no parsing
        # of the comments.  Instead, return as an unclassified LibraryReaction.
        reaction = LibraryReaction(
            index = reaction.index,
            reactants = reaction.reactants, 
            products = reaction.products, 
            kinetics = reaction.kinetics,
            reversible = reaction.reversible,
            duplicate = reaction.duplicate,
            library = 'Unclassified',
        )        
        
        return reaction  
    
    atKineticsComments = False
    lines = comments.strip().splitlines()
        
    for line in lines:
        
        
        tokens = line.split()
        if 'Reaction index:' in line:
            # Don't store the reaction indices
            pass
        
        elif 'Template reaction:' in line:
            label = str(tokens[-2])
            template = tokens[-1][1:-1].split(',')
            reaction = TemplateReaction(
                index = reaction.index,
                reactants = reaction.reactants, 
                products = reaction.products, 
                kinetics = reaction.kinetics,
                reversible = reaction.reversible,
                duplicate = reaction.duplicate,
                family = label,
                template = [Entry(label=g) for g in template],
            )
            
        elif 'Library reaction:' in line or 'Seed mechanism:' in line:
            label = str(tokens[-1])
            reaction = LibraryReaction(
                index = reaction.index,
                reactants = reaction.reactants, 
                products = reaction.products, 
                kinetics = reaction.kinetics,
                reversible = reaction.reversible,
                duplicate = reaction.duplicate,
                library = label,
            )   
            
        elif 'PDep reaction:' in line:
            networkIndex = int(tokens[-1][1:])
            reaction = PDepReaction(
                index = reaction.index,
                reactants = reaction.reactants, 
                products = reaction.products, 
                kinetics = reaction.kinetics, <commit_msg>the reaction family name is the last token instead of last second one<commit_after>#!/usr/bin/env python
# encoding: utf-8

################################################################################
#
#   RMG - Reaction Mechanism Generator
#
#   Copyright (c) 2009-2011 by the RMG Team (<EMAIL>)
#
#   Permission is hereby granted, free of charge, to any person obtaining a
#   copy of this software and associated documentation files (the 'Software'),
#   to deal in the Software without restriction, including without limitation
#   the rights to use, copy, modify, merge, publish, distribute, sublicense,
#   and/or sell copies of the Software, and to permit persons to whom the
#   Software is furnished to do so, subject to the following conditions:
#
#   The above copyright notice and this permission notice shall be included in
#   all copies or substantial portions of the Software.
#
#   THE SOFTWARE IS PROVIDED 'AS IS', WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
#   IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
#   FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
#   AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
#   LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
#   FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
#   DEALINGS IN THE SOFTWARE.
#
################################################################################

"""
This module contains functions for writing of Chemkin input files.
"""
import shutil
import math
import re
import logging
import textwrap
import os.path
import numpy

import rmgpy.kinetics as _kinetics
from rmgpy.reaction import Reaction
#from species import Species
from rmgpy.rmg.model import Species
from rmgpy.rmg.pdep import PDepReaction
from rmgpy.thermo import NASAPolynomial, NASA
import rmgpy.constants as constants
from rmgpy.quantity import Quantity
from rmgpy.data.base import Entry 
from rmgpy.data.kinetics.library import LibraryReaction
from rmgpy.data.kinetics.family import TemplateReaction
from rmgpy.rmg.pdep import PDepNetwork
from rmgpy.molecule import Molecule
from rmgpy.transport import TransportData

__chemkin_reaction_count = None
    
from rmgpy.util import makeOutputSubdirectory

################################################################################

class ChemkinError(Exception):
    """
    An exception class for exceptional behavior involving Chemkin files. Pass a
    string describing the circumstances that caused the exceptional behavior.
    """
    pass

################################################################################

def Ffloat(string):
    """
    Parse a Fortran-ish string into a float, like "1.00D 03" or "1.00d+03"
    """
    return float(string.replace("e", "E").replace("d", "D").replace("D", "E").replace("E ", "E+"))
    
def readThermoEntry(entry, Tmin=0, Tint=0, Tmax=0):
    """
    Read a thermodynamics `entry` for one species in a Chemkin file. Returns
    the label of the species and the thermodynamics model as a :class:`NASA`
    object.
    
    Format specification at http://www2.galcit.caltech.edu/EDL/public/formats/chemkin.html
    """
    lines = entry.splitlines()
    species = str(lines[0][0:18].split()[0].strip())
    
    comment = lines[0][len(species):24].strip()
    formula = {}
    for i in [24, 29, 34, 39, 74]:
        element,count = lines[0][i:i+2].strip(), lines[0][i+2:i+5].strip()
        if element:
            try:
                count = int(count)
            except ValueError:
                # Chemkin allows float values for the number of atoms, so try this next.
                try:
                    count = int(float(count))
                except ValueError:
                    logging.info("Trouble reading line '{0}' element segment '{1}'".format(lines[0].strip(),lines[0][i:i+5]))
                    if count == '' and re.match('\.?0*', element):
                        if i == 74: logging.info("Assuming it's spillover from Tint, and ignoring.")
                        else: logging.warning("Assuming it's not meant to be there, although it would be good to fix the chemkin file.")
                        count = 0
                    else:
                        raise
            if count != 0: # Some people put garbage elements in, with zero count. Ignore these. Allow negative counts though (eg. negative one electron)
                formula[element] = count
    phase = lines[0][44]
    if phase.upper() != 'G':
        logging.warning("Was expecting gas phase thermo data for {0}. Skipping thermo data.".format(species))
        return species, None, None
    
    # Extract the NASA polynomial coefficients
    # Remember that the high-T polynomial comes first!
    try:
        try:
            Tmin = float(lines[0][45:55].strip())
        except ValueError:
            pass
        try:
            Tmax = float(lines[0][55:65].strip())
        except ValueError:
            pass
        try:
            Tint = float(lines[0][65:73].strip())
        except ValueError:
            pass
        a0_high = Ffloat(lines[1][0:15].strip())
        a1_high = Ffloat(lines[1][15:30].strip())
        a2_high = Ffloat(lines[1][30:45].strip())
        a3_high = Ffloat(lines[1][45:60].strip())
        a4_high = Ffloat(lines[1][60:75].strip())
    
        a5_high = Ffloat(lines[2][0:15].strip())
        a6_high = Ffloat(lines[2][15:30].strip())
        a0_low = Ffloat(lines[2][30:45].strip())
        a1_low = Ffloat(lines[2][45:60].strip())
        a2_low = Ffloat(lines[2][60:75].strip())
    
        a3_low = Ffloat(lines[3][0:15].strip())
        a4_low = Ffloat(lines[3][15:30].strip())
        a5_low = Ffloat(lines[3][30:45].strip())
        a6_low = Ffloat(lines[3][45:60].strip())
    except (IndexError, ValueError), e:
        logging.warning('Error while reading thermo entry for species {0}'.format(species))
        logging.warning(e.message)
        return species, None, None
    
    # Construct and return the thermodynamics model
    thermo = NASA(
        polynomials = [
            NASAPolynomial(Tmin=(Tmin,"K"), Tmax=(Tint,"K"), coeffs=[a0_low, a1_low, a2_low, a3_low, a4_low, a5_low, a6_low]),
            NASAPolynomial(Tmin=(Tint,"K"), Tmax=(Tmax,"K"), coeffs=[a0_high, a1_high, a2_high, a3_high, a4_high, a5_high, a6_high])
        ],
        Tmin = (Tmin,"K"),
        Tmax = (Tmax,"K"),
    )
    if comment:
        thermo.comment = comment

    return species, thermo, formula

################################################################################

def readKineticsEntry(entry, speciesDict, Aunits, Eunits):
    """
    Read a kinetics `entry` for a single reaction as loaded from a Chemkin
    file. The associated mapping of labels to species `speciesDict` should also
    be provided. Returns a :class:`Reaction` object with the reaction and its
    associated kinetics.
    """
    Afactor = {
        'cm^3/(mol*s)': 1.0e6,
        'cm^3/(molecule*s)': 1.0e6,
        'm^3/(mol*s)': 1.0,
        'm^3/(molecule*s)': 1.0,
    }[Aunits[2]]
    
    lines = entry.strip().splitlines()
    
    # The first line contains the reaction equation and a set of
    # modified Arrhenius parameters
    reaction, thirdBody, kinetics, kunits, klow_units = _readKineticsReaction(
        line=lines[0], speciesDict=speciesDict, Aunits=Aunits, Eunits=Eunits)

    if len(lines) == 1 and not thirdBody:
        # If there's only one line then we know to use the high-P limit kinetics as-is
        reaction.kinetics = kinetics['arrhenius high']
    else:
        # There's more kinetics information to be read
        kinetics.update({
            'chebyshev coefficients': [],
            'efficiencies': {},
        })

        # Note that the subsequent lines could be in any order
        for line in lines[1:]:
            kinetics = _readKineticsLine(
                line=line, reaction=reaction, speciesDict=speciesDict, Eunits=Eunits,
                kunits=kunits, klow_units=klow_units,
                kinetics=kinetics)

        # Decide which kinetics to keep and store them on the reaction object
        # Only one of these should be true at a time!
        if 'chebyshev' in kinetics:
            chebyshev = kinetics['chebyshev']
            if chebyshev.Tmin is None or chebyshev.Tmax is None:
                raise ChemkinError('Missing TCHEB line for reaction {0}'.format(reaction))
            if chebyshev.Pmin is None or chebyshev.Pmax is None:
                raise ChemkinError('Missing PCHEB line for reaction {0}'.format(reaction))
            index = 0
            for t in range(chebyshev.degreeT):
                for p in range(chebyshev.degreeP):
                    chebyshev.coeffs.value_si[t,p] = kinetics[
                        'chebyshev coefficients'][index]
                    index += 1
            # Don't forget to convert the Chebyshev coefficients to SI units!
            # This assumes that s^-1, cm^3/mol*s, etc. are compulsory
            chebyshev.coeffs.value_si[0,0] -= (len(reaction.reactants) - 1) * math.log10(Afactor)
            reaction.kinetics = chebyshev
        elif 'pressure-dependent arrhenius' in kinetics:
            pdepArrhenius = kinetics['pressure-dependent arrhenius']
            # Check for duplicates and combine them to MultiArrhenius objects
            duplicatesToRemove = []
            duplicatesToAdd = []
            for index1 in range(len(pdepArrhenius)):
                reaction1 = pdepArrhenius[index1]
                P1, kinetics1 = reaction1
                if reaction1 in duplicatesToRemove:
                    continue
                for index2 in range(index1+1, len(pdepArrhenius)):
                    reaction2 = pdepArrhenius[index2]
                    P2, kinetics2 = reaction2
                    if P1 == P2:
                        if reaction1 not in duplicatesToRemove:
                            new_kinetics = _kinetics.MultiArrhenius()
                            duplicatesToAdd.append([P1,new_kinetics])
                            new_kinetics.arrhenius = [kinetics1]
                            duplicatesToRemove.append(reaction1)
                        new_kinetics.arrhenius.append(kinetics2)
                        duplicatesToRemove.append(reaction2)
            for item in duplicatesToRemove:
                pdepArrhenius.remove(item)
            pdepArrhenius.extend(duplicatesToAdd)
            
            pdepArrhenius = sorted(pdepArrhenius, key=lambda reaction: reaction[0])  # sort by ascending pressures

            reaction.kinetics = _kinetics.PDepArrhenius(
                pressures = ([P for P, arrh in pdepArrhenius],"atm"),
                arrhenius = [arrh for P, arrh in pdepArrhenius],
            )
        elif 'troe' in kinetics:
            troe = kinetics['troe']
            troe.arrheniusHigh = kinetics['arrhenius high']
            troe.arrheniusLow = kinetics['arrhenius low']
            troe.efficiencies = kinetics['efficiencies']
            reaction.kinetics = troe
        elif thirdBody:
            reaction.kinetics = _kinetics.ThirdBody(
                arrheniusLow=kinetics['arrhenius low'])
            reaction.kinetics.efficiencies = kinetics['efficiencies']
        elif 'arrhenius low' in kinetics:
            reaction.kinetics = _kinetics.Lindemann(
                arrheniusHigh=kinetics['arrhenius high'],
                arrheniusLow=kinetics['arrhenius low'])
            reaction.kinetics.efficiencies = kinetics['efficiencies']
        elif 'explicit reverse' in kinetics or reaction.duplicate:
            # it's a normal high-P reaction - the extra lines were only either REV (explicit reverse) or DUP (duplicate)
            reaction.kinetics = kinetics['arrhenius high']
        else:
            raise ChemkinError(
                'Unable to understand all additional information lines for reaction {0}.'.format(entry))
        
        # These things may *also* be true
        if 'sri' in kinetics:
            reaction.kinetics.comment += "Warning: SRI parameters from chemkin file ignored on import. "

        if 'explicit reverse' in kinetics:
            reaction.kinetics.comment += (
                "Chemkin file stated explicit reverse rate: {0}"
                ).format(kinetics['explicit reverse'])

    return reaction


def _readKineticsReaction(line, speciesDict, Aunits, Eunits):
    """
    Parse the first line of of a Chemkin reaction entry.
    """
    tokens = line.split()
    
    rmg = True
    try:
        float(tokens[-6])
    except (ValueError, IndexError):
        rmg = False
    AuncertaintyType = '+|-'
    if rmg:
        A = float(tokens[-6])
        n = float(tokens[-5])
        Ea = float(tokens[-4])
        try:
            dA = float(tokens[-3])
        except ValueError:
            AuncertaintyType = '*|/'
            dA = float(tokens[-3][1:])
        dn = float(tokens[-2])
        dEa = float(tokens[-1])
        reaction = ''.join(tokens[:-6])
    else:
        A = Ffloat(tokens[-3])
        n = Ffloat(tokens[-2])
        Ea = Ffloat(tokens[-1])
        dA = 0.0
        dn = 0.0
        dEa = 0.0
        reaction = ''.join(tokens[:-3])
    thirdBody = False
    
    # Split the reaction equation into reactants and products
    reversible = True
    reactants, products = reaction.split('=')
    if '<=>' in reaction:
        reactants = reactants[:-1]
        products = products[1:]
    elif '=>' in reaction:
        products = products[1:]
        reversible = False
    if '(+M)' in reactants: reactants = reactants.replace('(+M)','')
    if '(+m)' in reactants: reactants = reactants.replace('(+m)','')
    if '(+M)' in products:  products = products.replace('(+M)','')
    if '(+m)' in products:  products = products.replace('(+m)','')
    
    # Create a new Reaction object for this reaction
    reaction = Reaction(reactants=[], products=[], reversible=reversible)
    
    # Convert the reactants and products to Species objects using the speciesDict
    for reactant in reactants.split('+'):
        reactant = reactant.strip()
        stoichiometry = 1
        if reactant[0].isdigit():
            # This allows for reactions to be of the form 2A=B+C instead of A+A=B+C
            # The implementation below assumes an integer between 0 and 9, inclusive
            stoichiometry = int(reactant[0])
            reactant = reactant[1:]               
        if reactant.upper() == 'M':
            thirdBody = True
        elif reactant not in speciesDict:
            raise ChemkinError('Unexpected reactant "{0}" in reaction {1}.'.format(reactant, reaction))
        else:
            for i in range(stoichiometry):
                reaction.reactants.append(speciesDict[reactant])
    for product in products.split('+'):
        product = product.strip()
        stoichiometry = 1
        if product[0].isdigit():
            # This allows for reactions to be of the form A+B=2C instead of A+B=C+C
            # The implementation below assumes an integer between 0 and 9, inclusive
            stoichiometry = int(product[0])
            product = product[1:]
        if product.upper() == 'M':
            pass
        elif product not in speciesDict:
            if re.match('[0-9.]+',product):
                logging.warning("Looks like reaction {0!r} has fractional stoichiometry, which RMG cannot handle. Ignoring".format(line))
                raise ChemkinError('Skip reaction!')
            raise ChemkinError('Unexpected product "{0}" in reaction {1}.'.format(product, reaction))
        else:
            for i in range(stoichiometry):
                reaction.products.append(speciesDict[product])
    
    # Determine the appropriate units for k(T) and k(T,P) based on the number of reactants
    # This assumes elementary kinetics for all reactions
    try:
        Nreac = len(reaction.reactants) + (1 if thirdBody else 0)
        kunits = Aunits[Nreac]
        klow_units = Aunits[Nreac+1]
    except IndexError:
        raise ChemkinError('Invalid number of reactant species for reaction {0}.'.format(reaction))
    
    key = 'arrhenius low' if thirdBody else 'arrhenius high'
    
    kinetics = {
        key: _kinetics.Arrhenius(
            A = (A,kunits,AuncertaintyType,dA),
            n = (n,'','+|-',dn),
            Ea = (Ea,Eunits,'+|-',dEa),
            T0 = (1,"K"),
        ),
    }
    return (reaction, thirdBody, kinetics, kunits, klow_units)


def _readKineticsLine(line, reaction, speciesDict, Eunits, kunits, klow_units, kinetics):
    """
    Parse the subsequent lines of of a Chemkin reaction entry.
    """
    case_preserved_tokens = line.split('/')
    line = line.upper()
    tokens = line.split('/')

    if 'DUP' in line:
        # Duplicate reaction
        reaction.duplicate = True

    elif 'LOW' in line:
        # Low-pressure-limit Arrhenius parameters
        tokens = tokens[1].split()
        kinetics['arrhenius low'] = _kinetics.Arrhenius(
            A = (float(tokens[0].strip()),klow_units),
            n = float(tokens[1].strip()),
            Ea = (float(tokens[2].strip()),Eunits),
            T0 = (1,"K"),
        )

    elif 'HIGH' in line:
        # What we thought was high, was in fact low-pressure
        kinetics['arrhenius low'] = kinetics['arrhenius high']
        kinetics['arrhenius low'].A = (
            kinetics['arrhenius low'].A.value, klow_units)
        # High-pressure-limit Arrhenius parameters
        tokens = tokens[1].split()
        kinetics['arrhenius high'] = _kinetics.Arrhenius(
            A = (float(tokens[0].strip()),kunits),
            n = float(tokens[1].strip()),
            Ea = (float(tokens[2].strip()),Eunits),
            T0 = (1,"K"),
        )

    elif 'TROE' in line:
        # Troe falloff parameters
        tokens = tokens[1].split()
        alpha = float(tokens[0].strip())
        T3 = float(tokens[1].strip())
        T1 = float(tokens[2].strip())
        try:
            T2 = float(tokens[3].strip())
        except (IndexError, ValueError):
            T2 = None

        kinetics['troe'] = _kinetics.Troe(
            alpha = alpha,
            T3 = (T3,"K"),
            T1 = (T1,"K"),
            T2 = (T2,"K") if T2 is not None else None,
        )
    elif line.strip().startswith('SRI'):
        kinetics['sri'] = True
        """To define an SRI pressure-dependent reaction, in addition to the LOW or HIGH parameters, the
        keyword SRI followed by three or five parameters must be included in the following order: a, b,
        c, d, and e [Eq. (74)]. The fourth and fifth parameters are options. If only the first three are
        stated, then by default d = 1 and e = 0.
        """
        # see eg. http://www.dipic.unipd.it/faculty/canu/files/Comb/Docs/chemkinCK.pdf

    elif 'CHEB' in line:
        # Chebyshev parameters
        chebyshev = kinetics.get(
            'chebyshev', _kinetics.Chebyshev(kunits=kunits))
        kinetics['chebyshev'] = chebyshev
        tokens = [t.strip() for t in tokens]
        if 'TCHEB' in line:
            index = tokens.index('TCHEB')
            tokens2 = tokens[index+1].split()
            chebyshev.Tmin = Quantity(float(tokens2[0].strip()),"K")
            chebyshev.Tmax = Quantity(float(tokens2[1].strip()),"K")
        if 'PCHEB' in line:
            index = tokens.index('PCHEB')
            tokens2 = tokens[index+1].split()
            chebyshev.Pmin = Quantity(float(tokens2[0].strip()),"atm")
            chebyshev.Pmax = Quantity(float(tokens2[1].strip()),"atm")
        if 'TCHEB' in line or 'PCHEB' in line:
            pass
        elif chebyshev.degreeT == 0 or chebyshev.degreeP == 0:
            tokens2 = tokens[1].split()
            chebyshev.degreeT = int(float(tokens2[0].strip()))
            chebyshev.degreeP = int(float(tokens2[1].strip()))
            chebyshev.coeffs = numpy.zeros((chebyshev.degreeT,chebyshev.degreeP), numpy.float64)
        else:
            tokens2 = tokens[1].split()
            kinetics['chebyshev coefficients'].extend(
                [float(t.strip()) for t in tokens2])

    elif 'PLOG' in line:
        pdepArrhenius = kinetics.get('pressure-dependent arrhenius', [])
        kinetics['pressure-dependent arrhenius'] = pdepArrhenius
        tokens = tokens[1].split()
        pdepArrhenius.append([float(tokens[0].strip()),
            _kinetics.Arrhenius(
                A = (float(tokens[1].strip()),kunits),
                n = float(tokens[2].strip()),
                Ea = (float(tokens[3].strip()),Eunits),
                T0 = (1,"K"),
        )])
    elif tokens[0].startswith('REV'):
        reverseA = float(tokens[1].split()[0])
        kinetics['explicit reverse'] = line.strip()
        if reverseA == 0:
            logging.info("Reverse rate is 0 so making irreversible for reaction {0}".format(reaction))
            reaction.reversible = False
        else:
            logging.info("Ignoring explicit reverse rate for reaction {0}".format(reaction))

    else:
        # Assume a list of collider efficiencies
        try:
            for collider, efficiency in zip(case_preserved_tokens[0::2], case_preserved_tokens[1::2]):
                try:
                    efficiency = float(efficiency.strip())
                except ValueError:
                    error_msg = "{0!r} doesn't look like a collision efficiency for species {1} in line {2!r}".format(efficiency,collider.strip(),line)
                    logging.error(error_msg)
                    raise ChemkinError(error_msg)
                if collider.strip() in speciesDict:
                    kinetics['efficiencies'][speciesDict[collider.strip()].molecule[0]] = efficiency
                else: # try it with capital letters? Not sure whose malformed chemkin files this is needed for.
                    kinetics['efficiencies'][speciesDict[collider.strip().upper()].molecule[0]] = efficiency
        except IndexError:
            error_msg = 'Could not read collider efficiencies for reaction: {0}.\n'.format(reaction)
            error_msg += 'The following line was parsed incorrectly:\n{0}'.format(line)
            error_msg += "\n(Case-preserved tokens: {0!r} )".format(case_preserved_tokens)
            raise ChemkinError(error_msg)
    return kinetics


def readReactionComments(reaction, comments, read = True):
    """
    Parse the `comments` associated with a given `reaction`. If the comments
    come from RMG (Py or Java), parse them and extract the useful information.
    Return the reaction object based on the information parsed from these
    comments. If `read` if False, the reaction is returned as an "Unclassified"
    LibraryReaction.
    """
    
    if read == False:
        # The chemkin file was not generated by either RMG-Py or RMG-Java, thus, there should be no parsing
        # of the comments.  Instead, return as an unclassified LibraryReaction.
        reaction = LibraryReaction(
            index = reaction.index,
            reactants = reaction.reactants, 
            products = reaction.products, 
            kinetics = reaction.kinetics,
            reversible = reaction.reversible,
            duplicate = reaction.duplicate,
            library = 'Unclassified',
        )        
        
        return reaction  
    
    atKineticsComments = False
    lines = comments.strip().splitlines()
        
    for line in lines:
        
        
        tokens = line.split()
        if 'Reaction index:' in line:
            # Don't store the reaction indices
            pass
        
        elif 'Template reaction:' in line:
            label = str(tokens[-1])
            template = tokens[-1][1:-1].split(',')
            reaction = TemplateReaction(
                index = reaction.index,
                reactants = reaction.reactants, 
                products = reaction.products, 
                kinetics = reaction.kinetics,
                reversible = reaction.reversible,
                duplicate = reaction.duplicate,
                family = label,
                template = [Entry(label=g) for g in template],
            )
            
        elif 'Library reaction:' in line or 'Seed mechanism:' in line:
            label = str(tokens[-1])
            reaction = LibraryReaction(
                index = reaction.index,
                reactants = reaction.reactants, 
                products = reaction.products, 
                kinetics = reaction.kinetics,
                reversible = reaction.reversible,
                duplicate = reaction.duplicate,
                library = label,
            )   
            
        elif 'PDep reaction:' in line:
            networkIndex = int(tokens[-1][1:])
            reaction = PDepReaction(
                index = reaction.index,
                reactants = reaction.reactants, 
                products = reaction.products, 
                kinetics = reaction.kinetics, 
{'fraction_non_alphanumeric': 0.12249889730943501, 'fraction_numerical': 0.014675808973896308, 'mean_word_length': 2.222782192931447, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 4, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2080360', 'n_tokens_mistral': 14778, 'n_tokens_neox': 13791, 'n_words': 3984}