from chempy.kinetics.arrhenius import ArrheniusParamWithUnits
from chempy.kinetics.rates import MassAction
from chempy.kinetics._rates import RTPoly
from chempy.reactionsystem import ReactionSystem
from chempy.units import patched_numpy, Backend, unit_of, simplified, default_constants as dc, default_units as u
from chempy.util._expr import Expr, Constant
from chempy.util.parsing import get_parsing_context

def log10recipoly(*args):
    return 10**RTPoly(*args)
rpkw = dict(globals_=dict(u=u, log10recipoly=log10recipoly, Constant=Constant, MassAction=MassAction))
lgRTp = "MassAction([Constant(1/u.M/u.s) * log10recipoly([x*u.K**i for i, x in enumerate([13.339, -2.22e3, 7.333e5, -1.065e8])])])"
rs1 = ReactionSystem.from_string("OH + OH- -> H2O + O-; %s; name='R27f'" % lgRTp, rxn_parse_kwargs=rpkw)
rs1

rs1.rxns[0].param.rate_coeff({'temperature': 298*u.K})

be = Backend()
variables = {'temperature': 298.15*u.kelvin}
rates = rs1['R27f'].rate(dict(variables, **{'OH': 1e-11*u.molar, 'OH-': 1e-7*u.molar}), backend=be)
for sk in rs1.substances:
    assert sk in rates
    assert unit_of(rates[sk]) == u.molar/u.second
rates

class GibbsExpr(Expr):
    nargs = 4
    parameter_keys = ('temperature',)

    def __call__(self, variables, backend=patched_numpy, **kwargs):
        dS_over_R, dCp_over_R, dH_over_R, Tref = map(simplified, self.all_args(variables, backend=backend))
        T, = self.all_params(variables, backend=backend)
        return backend.exp(dS_over_R)*(T/Tref)**dCp_over_R*backend.exp(-dH_over_R/T)

GeNH3 = GibbsExpr([
    18.8*u.cal/u.K/u.mol/dc.molar_gas_constant,
    52*u.cal/u.K/u.mol/dc.molar_gas_constant,
    -0.87e3*u.cal/u.mol/dc.molar_gas_constant,
    298.15*u.K
])

dCp_R = 18.8*u.cal/u.K/u.mol/dc.molar_gas_constant
dCp_R, dCp_R.simplified

str_rs2 = """
NH4+ + OH- -> NH3 + H2O; rs1['R27f'].param; name='ammonium_hydroxide'
NH3 + H2O -> NH4+ + OH-; MassAction(rs1['R27f'].param.args[0]/GeNH3); name='ammonia_water'
"""

globals_ = get_parsing_context()
globals_['rs1'] = rs1
globals_['GeNH3'] = GeNH3
globals_['Arrh'] = ArrheniusParamWithUnits
rs2 = ReactionSystem.from_string(str_rs2, rxn_parse_kwargs=dict(globals_=globals_))
rs2

rs2.rxns[-1].rate_expr()

rs2.rxns[-1].param.rate_coeff({'temperature': 298*u.K})

for rk, cd in dict(ammonium_hydroxide={'NH4+': 1e-3*u.molar, 'OH-': 1e-7*u.molar},
                   ammonia_water={'NH3': 0.42*u.molar, 'H2O': 55.4*u.molar}).items():
    print(rs2[rk].rate(dict(variables, **cd), backend=be))