from collections import OrderedDict
from chempy import Reaction
from chempy.kinetics.rates import MassAction, RadiolyticBase
from chempy.units import to_unitless, default_units as u


def jl_dict(od):
    return "Dict([%s])" % ", ".join(["(:%s, %.4g)" % (k, v) for k, v in od.items()])


def _r(r, p, substmap, parmap, *, unit_conc, unit_time, variables=None):
    """
    Parameters
    ==========
    ...
    variables: dict
        e.g. dict(doserate=99.9*u.Gy/u.s, density=998*u.kg/u.m3)
    """
    (pk,) = r.param.unique_keys
    if isinstance(r.param, MassAction):
        ratcoeff = to_unitless(p[pk], unit_conc ** (1 - r.order()) / unit_time)
        if not r.inact_reac:
            r_str = "{}, {}".format(
                parmap[pk],
                r.string(
                    substances=substmap,
                    with_param=False,
                    Reaction_arrow="-->",
                    Reaction_coeff_space="",
                ),
            )
        else:
            all_keys = r.keys()
            reac_stoichs = r.all_reac_stoich(all_keys)
            act_stoichs = r.active_reac_stoich(all_keys)
            rate = "*".join(
                [parmap[pk]]
                + [
                    ("%s^%d" % (substmap[k], v)) if v > 1 else substmap[k]
                    for k, v in zip(all_keys, act_stoichs)
                    if v > 0
                ]
            )
            r2 = Reaction(
                dict([(k, v) for k, v in zip(all_keys, reac_stoichs) if v]), r.prod
            )
            r_str = "{}, {}".format(
                rate,
                r2.string(
                    substances=substmap,
                    with_param=False,
                    Reaction_arrow="\u21D2",
                    Reaction_coeff_space="",
                ),
            )
    elif isinstance(r.param, RadiolyticBase):
        ratcoeff = to_unitless(
            p[pk] * variables["doserate"] * variables["density"], unit_conc / unit_time
        )
        assert not r.reac and not r.inact_reac and not r.inact_prod
        ((prod, n),) = r.prod.items()
        assert n == 1
        r_str = ("{}, 0 \u21D2 {}" if ratcoeff > 0 else "{}, {} \u21D2 0").format(
            parmap[pk], substmap[prod]
        )
    else:
        raise NotImplementedError("What's that?")
    return r_str, pk, abs(ratcoeff)


class DiffEqBioJl:
    _template_body = """\
{name} = @{crn_macro} begin
    {reactions}
end {parameters}
{post}
"""

    defaults = dict(unit_conc=u.molar, unit_time=u.second)

    def __init__(self, *, rxs, pars, substance_key_map, parmap, **kwargs):
        self.rxs = rxs
        self.pars = pars
        self.substance_key_map = substance_key_map
        self.parmap = parmap
        self.unit_conc = kwargs.get("unit_conc", self.defaults["unit_conc"])
        self.unit_time = kwargs.get("unit_time", self.defaults["unit_time"])

    @classmethod
    def from_rsystem(
        cls,
        rsys,
        par_vals,
        *,
        variables=None,
        substance_key_map=lambda i, sk: "y%d" % i,
        **kwargs
    ):
        if not isinstance(substance_key_map, dict):
            substance_key_map = {
                sk: substance_key_map(si, sk) for si, sk in enumerate(rsys.substances)
            }
        parmap = dict(
            [(r.param.unique_keys[0], "p%d" % i) for i, r in enumerate(rsys.rxns)]
        )
        rxs, pars = [], OrderedDict()
        for r in rsys.rxns:
            rs, pk, pv = _r(
                r,
                par_vals,
                substance_key_map,
                parmap,
                variables=variables,
                unit_conc=kwargs.get("unit_conc", cls.defaults["unit_conc"]),
                unit_time=kwargs.get("unit_time", cls.defaults["unit_time"]),
            )
            rxs.append(rs)
            if pk in pars:
                raise ValueError("Are you sure (sometimes intentional)?")
            pars[parmap[pk]] = pv
        return cls(
            rxs=rxs,
            pars=pars,
            substance_key_map=substance_key_map,
            parmap=parmap,
            **kwargs
        )

    def render_body(self, sparse_jac=False):
        name = "rn"
        return self._template_body.format(
            crn_macro="min_reaction_network" if sparse_jac else "reaction_network",
            name=name,
            reactions="\n    ".join(self.rxs),
            parameters=" ".join(self.pars),
            post="addodes!({}, sparse_jac=True)".format(name) if sparse_jac else "",
        )

    def render_setup(self, *, ics, atol, tex=True, tspan=None):
        export = ""
        export += "p = %s\n" % jl_dict(self.pars)
        export += "ics = %s\n" % jl_dict(
            OrderedDict(
                {
                    self.substance_key_map[k]: v
                    for k, v in to_unitless(ics, u.molar).items()
                }
            )
        )
        if atol:
            export += "abstol_d = %s\n" % jl_dict(
                {
                    self.substance_key_map[k]: v
                    for k, v in to_unitless(atol, u.molar).items()
                }
            )
            export += (
                "abstol = Array([get(abstol_d, k, 1e-10) for k=keys(speciesmap(rn))])"
            )
        if tex:
            export += "subst_tex = Dict([%s])\n" % ", ".join(
                '(:%s, ("%s", "%s"))' % (v, k, k.latex_name)
                for k, v in self.substance_key_map.items()
            )
        if tspan:
            export += """\
tspan = (0., %12.5g)
u0 = Array([get(ics, k, 1e-28) for k=keys(speciesmap(rn))])
parr = Array([p[k] for k=keys(paramsmap(rn))])
oprob = ODEProblem(rn, u0, tspan, parr)
"""
        return export

    def render_solve(self):
        return (
            "sol = solve(oprob, reltol=1e-9, abstol=abstol, Rodas5(),"
            " callback=PositiveDomain(ones(length(u0)), abstol=abstol))"
        )
