#!/usr/bin/env python
# encoding: utf-8 -*-

"""
This module contains unit tests of the rmgpy.reaction module.
"""

import numpy
import unittest
from external.wip import work_in_progress

from rmgpy.species import Species, TransitionState
from rmgpy.reaction import Reaction
from rmgpy.statmech.translation import Translation, IdealGasTranslation
from rmgpy.statmech.rotation import Rotation, LinearRotor, NonlinearRotor, KRotor, SphericalTopRotor
from rmgpy.statmech.vibration import Vibration, HarmonicOscillator
from rmgpy.statmech.torsion import Torsion, HinderedRotor
from rmgpy.statmech.conformer import Conformer
from rmgpy.kinetics import Arrhenius
from rmgpy.thermo import Wilhoit
import rmgpy.constants as constants

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

class PseudoSpecies:
    """
    Can be used in place of a :class:`rmg.species.Species` for isomorphism checks.
    
    PseudoSpecies('a') is isomorphic with PseudoSpecies('A')
    but nothing else.
    """
    def __init__(self, label):
        self.label = label
    def __repr__(self):
        return "PseudoSpecies('{0}')".format(self.label)
    def __str__(self):
        return self.label
    def isIsomorphic(self, other):
        return self.label.lower() == other.label.lower()

class TestReactionIsomorphism(unittest.TestCase):
    """
    Contains unit tests of the isomorphism testing of the  Reaction class.
    """

    def makeReaction(self,reaction_string):
        """"
        Make a Reaction (containing PseudoSpecies) of from a string like 'Ab=CD'
        """
        reactants, products = reaction_string.split('=')
        reactants = [PseudoSpecies(i) for i in reactants]
        products = [PseudoSpecies(i) for i in products]
        return Reaction(reactants=reactants, products=products)
    def test1to1(self):
        r1 = self.makeReaction('A=B')
        self.assertTrue(r1.isIsomorphic(self.makeReaction('a=B')))
        self.assertTrue(r1.isIsomorphic(self.makeReaction('b=A')))
        self.assertFalse(r1.isIsomorphic(self.makeReaction('B=a'),eitherDirection=False))
        self.assertFalse(r1.isIsomorphic(self.makeReaction('A=C')))
        self.assertFalse(r1.isIsomorphic(self.makeReaction('A=BB')))
    def test1to2(self):
        r1 = self.makeReaction('A=BC')
        self.assertTrue(r1.isIsomorphic(self.makeReaction('a=Bc')))
        self.assertTrue(r1.isIsomorphic(self.makeReaction('cb=a')))
        self.assertTrue(r1.isIsomorphic(self.makeReaction('a=cb'),eitherDirection=False))
        self.assertFalse(r1.isIsomorphic(self.makeReaction('bc=a'),eitherDirection=False))
        self.assertFalse(r1.isIsomorphic(self.makeReaction('a=c')))
        self.assertFalse(r1.isIsomorphic(self.makeReaction('ab=c')))
    def test2to2(self):
        r1 = self.makeReaction('AB=CD')
        self.assertTrue(r1.isIsomorphic(self.makeReaction('ab=cd')))
        self.assertTrue(r1.isIsomorphic(self.makeReaction('ab=dc'),eitherDirection=False))
        self.assertTrue(r1.isIsomorphic(self.makeReaction('dc=ba')))
        self.assertFalse(r1.isIsomorphic(self.makeReaction('cd=ab'),eitherDirection=False))
        self.assertFalse(r1.isIsomorphic(self.makeReaction('ab=ab')))
        self.assertFalse(r1.isIsomorphic(self.makeReaction('ab=cde')))
    def test2to3(self):
        r1 = self.makeReaction('AB=CDE')
        self.assertTrue(r1.isIsomorphic(self.makeReaction('ab=cde')))
        self.assertTrue(r1.isIsomorphic(self.makeReaction('ba=edc'),eitherDirection=False))
        self.assertTrue(r1.isIsomorphic(self.makeReaction('dec=ba')))
        self.assertFalse(r1.isIsomorphic(self.makeReaction('cde=ab'),eitherDirection=False))
        self.assertFalse(r1.isIsomorphic(self.makeReaction('ab=abc')))
        self.assertFalse(r1.isIsomorphic(self.makeReaction('abe=cde')))


class TestReaction(unittest.TestCase):
    """
    Contains unit tests of the Reaction class.
    """
    
    def setUp(self):
        """
        A method that is called prior to each unit test in this class.
        """
        ethylene = Species(
            label = 'C2H4',
            conformer = Conformer(
                E0 = (44.7127, 'kJ/mol'),
                modes = [
                    IdealGasTranslation(
                        mass = (28.0313, 'amu'),
                    ),
                    NonlinearRotor(
                        inertia = (
                            [3.41526, 16.6498, 20.065],
                            'amu*angstrom^2',
                        ),
                        symmetry = 4,
                    ),
                    HarmonicOscillator(
                        frequencies = (
                            [828.397, 970.652, 977.223, 1052.93, 1233.55, 1367.56, 1465.09, 1672.25, 3098.46, 3111.7, 3165.79, 3193.54],
                            'cm^-1',
                        ),
                    ),
                ],
                spinMultiplicity = 1,
                opticalIsomers = 1,
            ),
        )
        
        hydrogen = Species(          
            label = 'H',
            conformer = Conformer(
                E0 = (211.794, 'kJ/mol'),
                modes = [
                    IdealGasTranslation(
                        mass = (1.00783, 'amu'),
                    ),
                ],
                spinMultiplicity = 2,
                opticalIsomers = 1,
            ),
        )
        
        ethyl = Species(
            label = 'C2H5',
            conformer = Conformer(
                E0 = (111.603, 'kJ/mol'),
                modes = [
                    IdealGasTranslation(
                        mass = (29.0391, 'amu'),
                    ),
                    NonlinearRotor(
                        inertia = (
                            [4.8709, 22.2353, 23.9925],
                            'amu*angstrom^2',
                        ),
                        symmetry = 1,
                    ),
                    HarmonicOscillator(
                        frequencies = (
                            [482.224, 791.876, 974.355, 1051.48, 1183.21, 1361.36, 1448.65, 1455.07, 1465.48, 2688.22, 2954.51, 3033.39, 3101.54, 3204.73],
                            'cm^-1',
                        ),
                    ),
                    HinderedRotor(
                        inertia = (1.11481, 'amu*angstrom^2'),
                        symmetry = 6,
                        barrier = (0.244029, 'kJ/mol'),
                        semiclassical = None,
                    ),
                ],
                spinMultiplicity = 2,
                opticalIsomers = 1,
            ),
        )
        
        TS = TransitionState(
            label = 'TS',
            conformer = Conformer(
                E0 = (266.694, 'kJ/mol'),
                modes = [
                    IdealGasTranslation(
                        mass = (29.0391, 'amu'),
                    ),
                    NonlinearRotor(
                        inertia = (
                            [6.78512, 22.1437, 22.2114],
                            'amu*angstrom^2',
                        ),
                        symmetry = 1,
                    ),
                    HarmonicOscillator(
                        frequencies = (
                            [412.75, 415.206, 821.495, 924.44, 982.714, 1024.16, 1224.21, 1326.36, 1455.06, 1600.35, 3101.46, 3110.55, 3175.34, 3201.88],
                            'cm^-1',
                        ),
                    ),
                ],
                spinMultiplicity = 2,
                opticalIsomers = 1,
            ),
            frequency = (-750.232, 'cm^-1'),
        )
        
        self.reaction = Reaction(
            reactants = [hydrogen, ethylene],
            products = [ethyl], 
            kinetics = Arrhenius(
                A = (501366000.0, 'cm^3/(mol*s)'),
                n = 1.637,
                Ea = (4.32508, 'kJ/mol'),
                T0 = (1, 'K'),
                Tmin = (300, 'K'),
                Tmax = (2500, 'K'),
            ),
            transitionState = TS,
        )
    
        # CC(=O)O[O]
        acetylperoxy = Species(
            label='acetylperoxy',
            thermo=Wilhoit(Cp0=(4.0*constants.R,"J/(mol*K)"), CpInf=(21.0*constants.R,"J/(mol*K)"), a0=-3.95, a1=9.26, a2=-15.6, a3=8.55, B=(500.0,"K"), H0=(-6.151e+04,"J/mol"), S0=(-790.2,"J/(mol*K)")),
        )

        # C[C]=O
        acetyl = Species(
            label='acetyl',
            thermo=Wilhoit(Cp0=(4.0*constants.R,"J/(mol*K)"), CpInf=(15.5*constants.R,"J/(mol*K)"), a0=0.2541, a1=-0.4712, a2=-4.434, a3=2.25, B=(500.0,"K"), H0=(-1.439e+05,"J/mol"), S0=(-524.6,"J/(mol*K)")),
        )

        # [O][O]
        oxygen = Species(
            label='oxygen',
            thermo=Wilhoit(Cp0=(3.5*constants.R,"J/(mol*K)"), CpInf=(4.5*constants.R,"J/(mol*K)"), a0=-0.9324, a1=26.18, a2=-70.47, a3=44.12, B=(500.0,"K"), H0=(1.453e+04,"J/mol"), S0=(-12.19,"J/(mol*K)")),
        )
        
        self.reaction2 = Reaction(
            reactants=[acetyl, oxygen], 
            products=[acetylperoxy], 
            kinetics = Arrhenius(
                A = (2.65e12, 'cm^3/(mol*s)'),
                n = 0.0,
                Ea = (0.0, 'kJ/mol'),
                T0 = (1, 'K'),
                Tmin = (300, 'K'),
                Tmax = (2000, 'K'),
            ),
        )
        
    def testIsIsomerization(self):
        """
        Test the Reaction.isIsomerization() method.
        """
        isomerization = Reaction(reactants=[Species()], products=[Species()])
        association = Reaction(reactants=[Species(),Species()], products=[Species()])
        dissociation = Reaction(reactants=[Species()], products=[Species(),Species()])
        bimolecular = Reaction(reactants=[Species(),Species()], products=[Species(),Species()])
        self.assertTrue(isomerization.isIsomerization())
        self.assertFalse(association.isIsomerization())
        self.assertFalse(dissociation.isIsomerization())
        self.assertFalse(bimolecular.isIsomerization())
        
    def testIsAssociation(self):
        """
        Test the Reaction.isAssociation() method.
        """
        isomerization = Reaction(reactants=[Species()], products=[Species()])
        association = Reaction(reactants=[Species(),Species()], products=[Species()])
        dissociation = Reaction(reactants=[Species()], products=[Species(),Species()])
        bimolecular = Reaction(reactants=[Species(),Species()], products=[Species(),Species()])
        self.assertFalse(isomerization.isAssociation())
        self.assertTrue(association.isAssociation())
        self.assertFalse(dissociation.isAssociation())
        self.assertFalse(bimolecular.isAssociation())

    def testIsDissociation(self):
        """
        Test the Reaction.isDissociation() method.
        """
        isomerization = Reaction(reactants=[Species()], products=[Species()])
        association = Reaction(reactants=[Species(),Species()], products=[Species()])
        dissociation = Reaction(reactants=[Species()], products=[Species(),Species()])
        bimolecular = Reaction(reactants=[Species(),Species()], products=[Species(),Species()])
        self.assertFalse(isomerization.isDissociation())
        self.assertFalse(association.isDissociation())
        self.assertTrue(dissociation.isDissociation())
        self.assertFalse(bimolecular.isDissociation())

    def testHasTemplate(self):
        """
        Test the Reaction.hasTemplate() method.
        """
        reactants = self.reaction.reactants[:]
        products = self.reaction.products[:]
        self.assertTrue(self.reaction.hasTemplate(reactants, products))
        self.assertTrue(self.reaction.hasTemplate(products, reactants))
        self.assertFalse(self.reaction2.hasTemplate(reactants, products))
        self.assertFalse(self.reaction2.hasTemplate(products, reactants))
        
        reactants.reverse()
        products.reverse()
        self.assertTrue(self.reaction.hasTemplate(reactants, products))
        self.assertTrue(self.reaction.hasTemplate(products, reactants))
        self.assertFalse(self.reaction2.hasTemplate(reactants, products))
        self.assertFalse(self.reaction2.hasTemplate(products, reactants))
        
        reactants = self.reaction2.reactants[:]
        products = self.reaction2.products[:]
        self.assertFalse(self.reaction.hasTemplate(reactants, products))
        self.assertFalse(self.reaction.hasTemplate(products, reactants))
        self.assertTrue(self.reaction2.hasTemplate(reactants, products))
        self.assertTrue(self.reaction2.hasTemplate(products, reactants))
        
        reactants.reverse()
        products.reverse()
        self.assertFalse(self.reaction.hasTemplate(reactants, products))
        self.assertFalse(self.reaction.hasTemplate(products, reactants))
        self.assertTrue(self.reaction2.hasTemplate(reactants, products))
        self.assertTrue(self.reaction2.hasTemplate(products, reactants))

    def testEnthalpyOfReaction(self):
        """
        Test the Reaction.getEnthalpyOfReaction() method.
        """
        Tlist = numpy.arange(200.0, 2001.0, 200.0, numpy.float64)
        Hlist0 = [float(v) for v in ['-146007', '-145886', '-144195', '-141973', '-139633', '-137341', '-135155', '-133093', '-131150', '-129316']]
        Hlist = self.reaction2.getEnthalpiesOfReaction(Tlist)
        for i in range(len(Tlist)):
            self.assertAlmostEqual(Hlist[i] / 1000., Hlist0[i] / 1000., 2)

    def testEntropyOfReaction(self):
        """
        Test the Reaction.getEntropyOfReaction() method.
        """
        Tlist = 
numpy.arange(200.0, 2001.0, 200.0, numpy.float64)