import numpy   as np
import cantera as ct
from scipy.optimize  import fsolve, linprog
from scipy.integrate import quad

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

class streams():

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

    def __init__(self, x0, x1, T0, T1, P, mechFile, Γ=None):
        """
        x0, input, dictionary, species name:mole (or mole fraction) for ξ=0
        x1, input, dictionary, species name:mole (or mole fraction) for ξ=1
        T0, input, double, temperature for ξ=0 (K)
        T1, input, double, temperature for ξ=1 (K)
        P,  input, double, system pressure (Pa)
        mechFile, input, string, mechanism name ("gri30.yaml")
        """
        s = self
        s.gas = ct.Solution(mechFile)
        s.P = P

        s.gas.TPX = T0, P, x0
        s.T0 = s.gas.T
        s.x0 = s.gas.X
        s.y0 = s.gas.Y
        s.h0 = s.gas.enthalpy_mass

        s.gas.TPX = T1, P, x1
        s.T1 = s.gas.T
        s.x1 = s.gas.X
        s.y1 = s.gas.Y
        s.h1 = s.gas.enthalpy_mass

        s._set_Γ(Γ)
        s._set_β01()
        s._set_ξst()

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

    def _set_Γ(self, Γ):
        s = self
        if Γ is None:
            s._Γ    = np.zeros(4)     # CHON
            s._Γ[0] = 2.0 /s.gas.atomic_weight(s.gas.element_index("C"))
            s._Γ[1] = 0.5 /s.gas.atomic_weight(s.gas.element_index("H"))
            s._Γ[2] = -1.0/s.gas.atomic_weight(s.gas.element_index("O"))
            s._Γ[3] = 0.0
        else:
            s._Γ    = Γ


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

    def _set_β01(self):
        s = self

        s.gas.Y = s.y0
        yCHON = [s.gas.elemental_mass_fraction("C"),
                 s.gas.elemental_mass_fraction("H"),
                 s.gas.elemental_mass_fraction("O"),
                 s.gas.elemental_mass_fraction("N")]
        s._β0 = np.sum(s._Γ*yCHON)

        s.gas.Y = s.y1
        yCHON = [s.gas.elemental_mass_fraction("C"),
                 s.gas.elemental_mass_fraction("H"),
                 s.gas.elemental_mass_fraction("O"),
                 s.gas.elemental_mass_fraction("N")]
        s._β1 = np.sum(s._Γ*yCHON)

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

    def _set_ξst(self):
        s = self
        s.ξst = (0.0 - s._β0)/(s._β1 - s._β0)

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

    def get_ξ(self, x_or_y, moleOrMass="mass"):
        s = self

        if moleOrMass.lower() == "mass":
            s.gas.Y = x_or_y
        else:
            s.gas.X = x_or_y
        yCHON = [s.gas.elemental_mass_fraction("C"),
                 s.gas.elemental_mass_fraction("H"),
                 s.gas.elemental_mass_fraction("O"),
                 s.gas.elemental_mass_fraction("N")]
        β = np.sum(s._Γ*yCHON)

        return (β - s._β0)/(s._β1 - s._β0)

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

    def set_gas_mixing_state(self, ξ):
        s = self

        y,h = s.get_mixing_state(ξ)
        s.gas.HPY = h,s.P,y

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

    def get_mixing_state(self, ξ):
        s = self

        h = s.h0*(1-ξ) + s.h1*(ξ)
        y = s.y0*(1-ξ) + s.y1*(ξ)
        
        s.gas.HPY = h,s.P,y

        return y, h

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

    def get_pCC(self, ξ, getYorX='y'):
        s = self
        # CxHyOzNw + (x+y/4-z/2)O2 --> (x)CO2 + (y/2)H2O  + (w/2)N2

        s.gas.Y = s.get_mixing_state(s.ξst)[0]

        xC = s.gas.elemental_mole_fraction("C")
        xH = s.gas.elemental_mole_fraction("H")
        xO = s.gas.elemental_mole_fraction("O")
        xN = s.gas.elemental_mole_fraction("N")

        x = np.zeros(s.gas.n_species)
        x[s.gas.species_index("CO2")] = xC
        x[s.gas.species_index("H2O")] = xH/2
        x[s.gas.species_index("N2")]  = xN/2

        s.gas.X = x
        yst = s.gas.Y

        if ξ < s.ξst:
            y = s.y0 + (yst - s.y0) * ξ/s.ξst
        else:
            y = yst + (s.y1 - yst)*(ξ-s.ξst)/(1-s.ξst)
        
        if getYorX == 'y' or getYorX == 'Y':
            return y
        else:
            s.gas.Y = y
            return s.gas.X

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

    def get_pWGS(self, ξ, getYorX='y'):
        '''
        Water gas shift equilibrium constraint. 
        CO + H2O = CO2 + H2
        Three regions: lean, rich but ξ<ξstCO, and ξ>ξstCO, where ξstCO is the stoic mixf for fuel --> CO+H2
        Just PCC in region 1.
        WGS equilibrium in region 2: products are CO2, CO, H2O, H2, N2
        Linear profiles in region 3 between state at ξstCO and ξ=1.
        For methane, this works well up to ξ=0.15.
        '''
        s = self

        #-------------- Case 1: lean

        if ξ < s.ξst:                             
            s.gas.Y = s.get_mixing_state(s.ξst)[0]

            xC = s.gas.elemental_mole_fraction("C")
            xH = s.gas.elemental_mole_fraction("H")
            xO = s.gas.elemental_mole_fraction("O")
            xN = s.gas.elemental_mole_fraction("N")

            x = np.zeros(s.gas.n_species)
            x[s.gas.species_index("CO2")] = xC
            x[s.gas.species_index("H2O")] = xH/2
            x[s.gas.species_index("N2")]  = xN/2

            s.gas.X = x
            yst = s.gas.Y

            y = s.y0 + (yst - s.y0) * ξ/s.ξst

        #-------------- Case 2: rich

        else:

            #-------------- get ξstCO   # stoic mix frac for fuel --> CO, H2
            #------- using a "Bilger" mixture fraction for CO, H2 products
            #------- set Γ's so that β=0 at the CO stoichiometric point ξstCO

            ΓCO =  np.array([1/s.gas.atomic_weight(s.gas.element_index("C")),
                    -1/s.gas.atomic_weight(s.gas.element_index("O"))])

            s.gas.Y = s.y0
            yCO = [s.gas.elemental_mass_fraction("C"),
                   s.gas.elemental_mass_fraction("O")]
            β0 = np.sum(ΓCO*yCO)

            s.gas.Y = s.y1
            yCO = [s.gas.elemental_mass_fraction("C"),
                   s.gas.elemental_mass_fraction("O")]
            β1 = np.sum(ΓCO*yCO)

            ξstCO = (0.0 - β0)/(β1 - β0)

            #--------------- Case 2a: rich, but not too rich
            # do water gas shift equilibrium products
            # CxHyOzNw  --> bCO2 + cCO + dH2O + eH2 + (w/2)N2
            # balance equations: C: b+c=x
            #                    H: 2d+2e=y
            #                    O: 2b+c+d=z
            # Keq = (b*e)/(c*d)
            #       c = x-b
            #       d = z-x-b
            #       e = y/2 - z + x + b
            # Keq is quadratic in b. Solve for b, then c d e, then composition


            if ξ < ξstCO:
            
                s.gas.Y = s.get_mixing_state(ξ)[0]

                x = s.gas.elemental_mole_fraction("C")
                y = s.gas.elemental_mole_fraction("H")
                z = s.gas.elemental_mole_fraction("O")
                w = s.gas.elemental_mole_fraction("N")

                #------------

                def get_Keq(T):
                    s.gas.TPX = T, 101325, s.gas.X
                    Gsp_RT = s.gas.standard_gibbs_RT 
                    ΔGrxn_RT = Gsp_RT[s.gas.species_index("CO2")] + Gsp_RT[s.gas.species_index("H2")] - \
                               Gsp_RT[s.gas.species_index("CO")]  - Gsp_RT[s.gas.species_index("H2O")]
                    return np.exp(-ΔGrxn_RT)

                #------------

                h = s.h0*(1-ξ) + s.h1*(ξ)

                #------------

                def F(T, yMF):         # yMF is along for the ride: convenience in getting final y 

                    Keq = get_Keq(T)
                    aa = Keq-1
                    bb = z - Keq*z - x - y/2
                    cc = Keq*x*z - Keq*x*x

                    b = ( -bb - np.sqrt(bb*bb - 4*aa*cc) )/(2*aa)
                    c = x-b
                    d = z-x-b
                    e = y/2 - z + x + b

                    xMF = np.zeros(s.gas.n_species)
                    xMF[s.gas.species_index("CO2")] = b
                    xMF[s.gas.species_index("CO")]  = c
                    xMF[s.gas.species_index("H2O")] = d
                    xMF[s.gas.species_index("H2")]  = e
                    xMF[s.gas.species_index("N2")]  = w/2

                    s.gas.HPX = h, s.P, xMF
                    yMF[:] = s.gas.Y
                    #return 0        # just use equilibrium temperature
                    return s.gas.enthalpy_mass - h

                #------------

                Tg = s.get_Tad_for_EQ(ξ)       # guess value
                yMF  = np.zeros(s.gas.n_species)
                T = fsolve(F, Tg, factor=0.1, maxfev=10000, xtol=1E-3, args=(yMF,))[0]
                y = yMF

            #--------------- Case 2b: too rich, have leftover fuel
            # just a linear coupling between ξstCO and ξ=1

            else:
                s.gas.Y = s.get_mixing_state(ξstCO)[0]

                xC = s.gas.elemental_mole_fraction("C")
                xH = s.gas.elemental_mole_fraction("H")
                xO = s.gas.elemental_mole_fraction("O")
                xN = s.gas.elemental_mole_fraction("N")

                x = np.zeros(s.gas.n_species)
                x[s.gas.species_index("CO")] = xC
                x[s.gas.species_index("H2")] = xH/2
                x[s.gas.species_index("N2")] = xN/2

                s.gas.X = x
                ystCO = s.gas.Y

                y = ystCO + (s.y1 - ystCO)*(ξ-ξstCO)/(1-ξstCO)

        #--------------- Done: return results

        if getYorX == 'y' or getYorX == 'Y':
            return y
        else:
            s.gas.Y = y
            return s.gas.X
    

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

    def get_pEQ(self, ξ, getYorX='y'):
        s = self

        y,h = s.get_mixing_state(ξ)
        s.gas.HPY = h, s.P, y
        s.gas.equilibrate("HP")
        y = s.gas.Y
        if getYorX == 'y' or getYorX == 'Y':
            return s.gas.Y
        else:
            return s.gas.X

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

    def get_Tad_for_pCC(self, ξ):
        s = self

        h = s.get_mixing_state(ξ)[1]
        y = s.get_pCC(ξ)
        s.gas.HPY = h, s.P, y
        return s.gas.T

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

    def get_Tad_for_pWGS(self, ξ):
        s = self

        h = s.get_mixing_state(ξ)[1]
        y = s.get_pWGS(ξ)
        s.gas.HPY = h, s.P, y
        return s.gas.T

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

    def get_Tad_for_EQ(self, ξ):
        s = self

        y,h = s.get_mixing_state(ξ)
        s.gas.HPY = h, s.P, y
        s.gas.equilibrate("HP")
        return s.gas.T

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

    def get_Φ_from_ξ(self, ξ):
        s = self

        if ξ == 1.0:
            return 1E20
        else:
            return ξ*(1-s.ξst)/(s.ξst*(1-ξ))

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

    def get_ξ_from_φ(self, φ):
        s = self
        return s.ξst*φ/(1 - s.ξst + s.ξst*φ)

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

    def get_y_from_x(self, x):
        s = self
        s.gas.X = x
        return s.gas.Y

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

    def get_x_from_y(self, y):
        s = self
        s.gas.Y = y
        return s.gas.X

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

    def set_gas_state_adiabatic_compression_expansion2(self, V2V1):
        """
        V2V1, input, double, V2/V1.
        Ideal gas
        Assume gas is already set at state 1
        """
        s = self

        V1 = s.gas.volume_mole
        T1 = s.gas.T
        P1 = s.gas.P

        V2 = V1*V2V1

        lnV2V1 = np.log(V2V1)
        def F(T2):
            def f(T2):
                s.gas.TPX = T2, ct.gas_constant*T2/V2, s.gas.X
                return s.gas.cv_mole/ct.gas_constant/T2
            return quad(f, T1, T2) + lnV2V1
        T2g = T1*(1/V2V1)**(s.gas.cp/s.gas.cv - 1)
        T2 = fsolve(F, T2g)[0]
        s.gas.TPX = T2, ct.gas_constant*T2/V2, s.gas.X

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

    def set_gas_state_adiabatic_compression_expansion(self, V2V1):
        """
        V2V1, input, double, V2/V1.
        Ideal gas
        Assume gas is already set at state 1
        """
        s = self

        V1 = s.gas.volume_mole
        T1 = s.gas.T
        P1 = s.gas.P
        S1 = s.gas.entropy_mass

        V2 = V1*V2V1

        s.gas.SV = s.gas.entropy_mass, s.gas.volume_mass*V2V1

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

    def get_LHV_pCC(self):
        s = self
        s.set_gas_mixing_state(s.ξst)
        s.TPX = 298.15, 101325, s.gas.X
        hR = s.gas.enthalpy_mass
        yPCC =  s.get_pCC(s.ξst, getYorX='y')
        s.gas.TPY = 298.15,101325,yPCC
        hP = s.gas.enthalpy_mass
        LHV = (hR - hP)/s.ξst
        return LHV                   # J/(kg fuel)

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

    def get_Aik(self):
        '''
        Get the matrix of element moles in species.
        Matrix is Nsp x Nel
        '''
        s = self
        nsp = s.gas.n_species
        nel = s.gas.n_elements
        A = np.zeros((nsp, nel))
        for i in range(nsp):
            for k in range(nel):
                A[i,k] = s.gas.n_atoms(i,k)
        return A

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

    def get_n_maxmin(self, α):
        '''
        Get maxmin composition as moles. 
        That is, the composition satisfying element balances whose smallest mole fraction is maximized.
        See Pope Appendix A: https://tcg.mae.cornell.edu/pubs/Pope_CUR_03.pdf
        input: α is the array of element moles: corresponding to gas.element_names
            This can be computed as Aik dot x for some mole fraction vector x, where Aik is from get_Aik.
            α is Pope's c
        '''
        s = self

        nsp = s.gas.n_species
        nel = s.gas.n_elements
        n   = nsp + 1

        c = np.zeros(n)           # Pope's f
        c[-1] = -1

        Aik = s.get_Aik()         # Pope's B

        A = np.diag(np.full(nsp, -1))
        A = np.column_stack((A, np.ones(nsp)))

        Ae = np.column_stack((Aik.T, np.zeros(s.gas.n_elements)))

        b = np.zeros(nsp)

        res = linprog(c, A, b, Ae, α)
        if not res['success']:
            print('WARNING: get_x_maxmin linprog call failed to find optimum')

        ni_mm = res['x'][:-1]    # array of moles of species
        return ni_mm
        

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

#strm = streams({"O2":1, "N2":3.76}, {"CH4":1}, 298.15, 298.15, 101325, "gri30.yaml")
#print(strm.ξst)
#print(strm.get_pCC(strm.ξst))
#print(strm.get_Tad_for_pCC(strm.ξst))