"""
Reaction diffusion solver for two-component 1D systems.

Justin Bois, February 2015
"""

import numpy as np
from scipy.integrate import odeint

# #################################
def asdm_rxn(a, s, mu):
    """
    Reaction expression for activator-substrate depletion model.

    Returns the rate of production of activator and substrate, respectively.

    r_a = a**2 * s - a
    r_s = mu * (1 - a**2 * s)
    """
    a2s = a**2 * s
    return a2s - a, mu * (1.0 - a2s)


# #################################
def gierer_meinhardt_rxn(a, h, rho, sigma=0.0):
    """
    Reaction expression for the Gierer-Meinhardt model.

    Returns the rate of production of activator and inhibitor, respectively.

    r_a = a**2 / h - a + sigma
    r_h = rho_h * (a**2 - h)
    """
    return a * (a / h - 1.0) + sigma, rho * (a**2 - h)


# #################################
def dc_dt(c, t, diffusion_coeffs, rxn_fun, rxn_params, h2):
    """
    Returns the time derivative of concentrations for a
    reaction-diffusion system.

    The concentration of the two species are interleaved in the input
    array c.  The interleaving allows us to take advantage of the
    banded structure of the Jacobian (when applicable) when using the
    Hindmarsh algorithm for integrating in time.

    rxn_params is a tuple containing additional arguments to be passed to
    rxn_fun, which gives the rate of production of the respective species.

    h2 is the square of the grid spacing.

    """

    # Views of concentrations for convenience (they are interleaved to give
    # nice structure to Jacobian for odeint to use)
    a = c[::2]
    b = c[1::2]

    # Return array with views
    dc_dt = np.empty_like(c)
    da_dt = dc_dt[::2]
    db_dt = dc_dt[1::2]

    # Compute diffusion terms (central differencing w/ Neumann BCs)
    da_dt[0] = diffusion_coeffs[0] * 2.0 * (a[1] - a[0]) / h2
    da_dt[1:-1] = diffusion_coeffs[0] * np.diff(a, 2) / h2
    da_dt[-1] = diffusion_coeffs[0] * 2.0 * (a[-2] - a[-1]) / h2
    db_dt[0] = diffusion_coeffs[1] * 2.0 * (b[1] - b[0]) / h2
    db_dt[1:-1] = diffusion_coeffs[1] * np.diff(b, 2) / h2
    db_dt[-1] = diffusion_coeffs[1] * 2.0 * (b[-2] - b[-1]) / h2

    # Compute reaction terms
    r_a, r_b = rxn_fun(a, b, *rxn_params)

    # Add to diffusion terms
    da_dt += r_a
    db_dt += r_b

    return dc_dt


# #################################
def solve(a_0, b_0, t, len_x, diffusion_coeffs, rxn_fun, rxn_params,
          banded=False):
    """
    Solves RD system.
    a_0 = initial condition (n_grid_points in length)
    b_0 = initial condition (n_grid_points in length)
    t = time points
    n = number of grid points
    len_x = physical length of the x-domain
    diffusion_coefficients = array of diffusion coeffs of species
    rxn_fun = function returning chemical reaction contributions to dynamics
    rxn_params = tuple of parameters to be passed to rxn_fun

    Returns an array that is len(t) by n_grid_points containing
    the solution to the reaction-diffusion PDEs for each species.

    Uniform grid point spacing is assumed.
    """

    # Number of grid points
    n = len(a_0)

    # Square of grid point spacing
    h2 = (len_x / (n - 1.0))**2

    # Set up parameters to be passed in to dc_dt
    params = (diffusion_coeffs, rxn_fun, rxn_params, h2)

    # Set up initial condition
    c0 = np.empty(2*n)
    c0[::2] = a_0
    c0[1::2] = b_0

    # Solve using odeint, taking advantage of banded structure
    if banded:
        c = odeint(dc_dt, c0, t, args=params, ml=2, mu=2)
    else:
        c = odeint(dc_dt, c0, t, args=params)
        
    # Return results, concentrations of a and b through time
    return c[:,::2], c[:,1::2]