# models.py
# some pre-defined models

import jax.numpy as jnp
from jsmfsb import Spn


def bd(th=[1, 1.1]):
    """Create a birth-death model

    Create and return a Spn object representing a discrete stochastic
    birth-death model.

    Parameters
    ----------
    th: array
        array of length 2 containing the birth and death rates

    Returns
    -------
    Spn model object with rates `th`

    Examples
    --------
    >>> import jsmfsb
    >>> import jax
    >>> bd = jsmfsb.models.bd()
    >>> step = bd.step_gillespie()
    >>> k = jax.random.key(42)
    >>> jsmfsb.sim_time_series(k, bd.m, 0, 50, 0.1, step)
    """
    return Spn(
        ["X"],
        ["Birth", "Death"],
        [[1], [1]],
        [[2], [0]],
        lambda x, t: jnp.array([th[0] * x[0], th[1] * x[0]]),
        [100],
    )


def dimer(th=[0.00166, 0.2]):
    """Create a dimerisation kinetics model

    Create and return a Spn object representing a discrete stochastic
    dimerisation kinetics model.

    Parameters
    ----------
    th: array
        array of length 2 containing the rates of the bind and unbind reactions

    Returns
    -------
    Spn model object with rates `th`

    Examples
    --------
    >>> import jsmfsb
    >>> import jax
    >>> dimer = jsmfsb.models.dimer()
    >>> step = dimer.step_gillespie()
    >>> k = jax.random.key(42)
    >>> jsmfsb.sim_time_series(k, dimer.m, 0, 50, 0.1, step)
    """
    return Spn(
        ["P", "P2"],
        ["Dim", "Diss"],
        [[2, 0], [0, 1]],
        [[0, 1], [2, 0]],
        lambda x, t: jnp.array([th[0] * x[0] * (x[0] - 1) / 2, th[1] * x[1]]),
        [301, 0],
    )


def id(th=[1, 0.1]):
    """Create an immigration-death model

    Create and return a Spn object representing a discrete stochastic
    immigration-death model.

    Parameters
    ----------
    th: array
        array of length 2 containing the immigration and death rates

    Returns
    -------
    Spn model object with rates `th`

    Examples
    --------
    >>> import jsmfsb
    >>> import jax
    >>> id = jsmfsb.models.id()
    >>> step = id.step_gillespie()
    >>> k = jax.random.key(42)
    >>> jsmfsb.sim_time_series(k, id.m, 0, 50, 0.1, step)
    """
    return Spn(
        ["X"],
        ["Immigration", "Death"],
        [[0], [1]],
        [[1], [0]],
        lambda x, t: jnp.array([th[0], th[1] * x[0]]),
        [0],
    )


def lv(th=[1, 0.005, 0.6]):
    """Create a Lotka-Volterra model

    Create and return a Spn object representing a discrete stochastic
    Lotka-Volterra model.

    Parameters
    ----------
    th: array
        array of length 3 containing the rates of the three governing reactions,
        prey reproduction, predator-prey interaction, and predator death

    Returns
    -------
    Spn model object with rates `th`

    Examples
    --------
    >>> import jsmfsb
    >>> import jax
    >>> lv = jsmfsb.models.lv()
    >>> step = lv.step_gillespie()
    >>> k = jax.random.key(42)
    >>> jsmfsb.sim_time_series(k, lv.m, 0, 50, 0.1, step)
    """
    return Spn(
        ["Prey", "Predator"],
        ["Prey rep", "Inter", "Pred death"],
        [[1, 0], [1, 1], [0, 1]],
        [[2, 0], [0, 2], [0, 0]],
        lambda x, t: jnp.array([th[0] * x[0], th[1] * x[0] * x[1], th[2] * x[1]]),
        [50, 100],
    )


def mm(th=[0.00166, 1e-4, 0.1]):
    """Create a Michaelis-Menten enzyme kinetic model

    Create and return a Spn object representing a discrete stochastic
    Michaelis-Menten enzyme kinetic model.

    Parameters
    ----------
    th: array
        array of length 3 containing the binding, unbinding and production rates

    Returns
    -------
    Spn model object with rates `th`

    Examples
    --------
    >>> import jsmfsb
    >>> import jax
    >>> mm = jsmfsb.models.mm()
    >>> step = mm.step_gillespie()
    >>> k = jax.random.key(42)
    >>> jsmfsb.sim_time_series(k, mm.m, 0, 50, 0.1, step)
    """
    return Spn(
        ["S", "E", "SE", "P"],
        ["Bind", "Unbind", "Produce"],
        [[1, 1, 0, 0], [0, 0, 1, 0], [0, 0, 1, 0]],
        [[0, 0, 1, 0], [1, 1, 0, 0], [0, 1, 0, 1]],
        lambda x, t: jnp.array([th[0] * x[0] * x[1], th[1] * x[2], th[2] * x[2]]),
        [301, 120, 0, 0],
    )


def sir(th=[0.0015, 0.1]):
    """Create a basic SIR compartmental epidemic model

    Create and return a Spn object representing a discrete stochastic
    SIR model.

    Parameters
    ----------
    th: array
        array of length 2 containing the rates of the two governing transitions

    Returns
    -------
    Spn model object with rates `th`

    Examples
    --------
    >>> import jsmfsb
    >>> import jax
    >>> sir = jsmfsb.models.sir()
    >>> step = sir.step_gillespie()
    >>> k = jax.random.key(42)
    >>> jsmfsb.sim_time_series(k, sir.m, 0, 50, 0.1, step)
    """
    return Spn(
        ["S", "I", "R"],
        ["S->I", "I->R"],
        [[1, 1, 0], [0, 1, 0]],
        [[0, 2, 0], [0, 0, 1]],
        lambda x, t: jnp.array([th[0] * x[0] * x[1], th[1] * x[1]]),
        [197, 3, 0],
    )


# eof
