"""
Calcium-induced Calcium release, coupled 40-state model from
Greenstein et al. (2012)
"""

import sympy as sp

# Subspace calcium states for individual cicr states
CASS_INDEX = [2 * int(i == 2) + int(j == 2) for i in range(4) for j in range(10)]
CASS_INDEX_INV = [[k for k, val in enumerate(CASS_INDEX) if val == i] for i in range(4)]

# Symbols names
simple_names = ["f", "g", "a", "b", "gamma_0", "omega", "alpha", "beta",
                "alpha_a", "beta_b", "y_inf", "tau_y", "k_b", "k_f", "omega_b", 
                "F", "R", "T", "t", "P_LCC", "r_RyR", "r_xfer", "V_SS"]

# symbols of RyR transition constants
for (i, j) in [(a, a + 1) for a in range(5)] + [(1, 4)]:
    simple_names.append(f"K0_[{i},{j}]")
    simple_names.append(f"K0_[{j},{i}]")

# Create symbols from names
syms = {n: sp.symbols(n, nonnegative=True) for n in simple_names}

# Indivisual names and symbols
vm = syms["V_m"] = sp.Symbol("V_m", real=True)

# Calcium symbols
for loc in "o in JSR".split():
    syms[f"ca_{loc}"] = sp.Symbol("[Ca^{2+}]_{%s}" % (loc), nonnegative=True)

syms["ca_ss"] = sp.Matrix(sp.symbols("[Ca^{2+}]_ss[0:4]", nonnegative=True))
p_cicr = sp.Matrix(sp.symbols("p_0:4", nonnegative=True))

# Algebrain relations of subspace Calcium'
v2FRT = 2 * vm * syms["F"] / syms["R"] / syms["T"]
e2vfrt = sp.E ** v2FRT
ca_out_weight = syms["P_LCC"] / syms["V_SS"] * v2FRT / (e2vfrt - 1)
ca_out_value = 0.341 * syms["ca_out"] * ca_out_weight
ca_out_weight *= e2vfrt
cain_weight = syms["R_XFER"]
ca_in_value = syms["ca_in"] * cain_weight
ca_jsr_weight = syms["R_RYR"]
ca_jsr_value = syms["ca_JSR"] * ca_jsr_weight

ca_ss = [ca_in_value / cain_weight,
         (ca_out_value + ca_in_value) / (ca_out_weight + cain_weight),
         (ca_jsr_value + ca_in_value) / (ca_jsr_weight + cain_weight),
         (ca_out_value + ca_in_value + ca_jsr_value) / (ca_out_weight + cain_weight + ca_jsr_weight)]

# Dependent relationships of symbols
deps = [(syms["alpha"], 0.78113 * sp.exp(0.08 * (vm - 8))),
        (syms["beta"],  0.34319 * sp.exp(-0.086676 * (vm - 8))),
        (syms["alpha_a"], syms["a"] * syms["alpha"]),
        (syms["beta_b"], syms["beta"] / syms["b"]),
        (syms["y_inf"], 0.25 + 0.75 / (1 + sp.exp((vm + 25) / 5))),
        (syms["tau_y"], 17 + 500 / ((1 + sp.exp((vm + 28) / 10)) * (1 + sp.exp(-(vm + 42) / 14)))),
        (syms["k_b"], syms["y_inf"] / syms["tau_y"]),
        (syms["k_f"], (1 - syms["y_inf"]) / syms["tau_y"]),
        (syms["omega_b"], syms["omega"] / syms["b"]),
        zip(syms["ca_ss"], ca_ss),
        (),
        ()]

# Equations for printing
equations = {str(lhs): sp.Eq(lhs, rhs, evaluate=False) for (lhs, rhs) in deps}


def add_rates(jac, i, j, kf, kb):
    """
    Adds 1st order reaction constants to
    the jacobian matrix of Markov-type ODE systems
    """
    jac[i, i] -= kf
    jac[j, j] -= kb
    jac[i, j] += kb
    jac[j, i] += kf
    return jac


# Jacobian matrix of the ODE system
jacobian = sp.zeros(40)

# Fill in the LCC part
for r_idx in range(0, 40, 10):
    g_cass = syms["gamma_0"] * syms["ca_ss"][CASS_INDEX[r_idx]]
    g_a_cass = syms["a"] * g_cass
    for i in range(r_idx, r_idx + 5):
        jacobian = add_rates(jacobian, i, i + 5, syms["k_f"], syms["k_b"])
    for i in (0 + r_idx, 5 + r_idx):
        jacobian = add_rates(jacobian, i, i + 1, syms["alpha"], syms["beta"])
        jacobian = add_rates(jacobian, i, i + 3, g_cass, syms["omega"])
    for i in (1 + r_idx, 6 + r_idx):
        jacobian = add_rates(jacobian, i, i + 1, syms["f"], syms["g"])
        jacobian = add_rates(jacobian, i, i + 3, g_a_cass, syms["omega_b"])
    for i in (3 + r_idx, 8 + r_idx):
        jacobian = add_rates(jacobian, i, i + 1, syms["alpha_a"], syms["beta_b"])


# RyR transition constants

def k_01(ca, syms):
    return syms["K0_01"] * (ca**2)


def k_10(ca, syms):
    return syms["K0_10"]


def k_12(ca, syms):
    return syms["K0_12"] * (ca**2)


def k_21(ca, syms):
    return syms["K0_21"] * syms["K0_32"] / (syms["K0_23"] * (ca**2) + syms["K0_32"])


def k_13(ca, syms):
    return syms["K0_14"] * (ca**2)


def k_31(ca, syms):
    return syms["K0_41"] * syms["K0_54"] / (syms["K0_45"] * (ca**2) + syms["K0_54"])


def k_23(ca, syms):
    return syms["K0_34"] * syms["K0_23"] * (ca**2) / (syms["K0_32"] + syms["K0_23"] * (ca**2))


def k_32(ca, syms):
    return syms["K0_54"] * syms["K0_43"] * (ca**2) / (syms["K0_54"] + syms["K0_45"] * (ca**2))


for l_idx in range(0, 10):
    ca = [syms["ca_ss"][CASS_INDEX[l_idx + i*10]] for i in range(4)]
    jacobian = add_rates(jacobian, 0 + l_idx, 10 + l_idx, k_01(ca[0], syms), k_10(ca[1], syms))
    jacobian = add_rates(jacobian, 10 + l_idx, 20 + l_idx, k_12(ca[1], syms), k_21(ca[2], syms))
    jacobian = add_rates(jacobian, 20 + l_idx, 30 + l_idx, k_23(ca[2], syms), k_32(ca[3], syms))
    jacobian = add_rates(jacobian, 10 + l_idx, 30 + l_idx, k_13(ca[1], syms), k_31(ca[3], syms))


x_sym = sp.MatrixSymbol('x', 40, 1)
x_states = sp.Matrix(x_sym)
dxdt = jacobian * x_states
for i, x in enumerate(x_states):
    lhs = sp.diff(x(syms["t"]), syms["t"])
    rhs = dxdt[i]
    deps.append((lhs, rhs))
    equations[f"dx{i}dt"] = sp.Eq(lhs, rhs, evaluate=False)


syms["ca_out"] = sp.Symbol("[Ca^{2+}]_{o}", nonnegative=True)
syms["ca_in"] = sp.Symbol("[Ca^{2+}]_{in}", nonnegative=True)
syms["ca_jsr"] = sp.Symbol("[Ca^{2+}]_{JSR}", nonnegative=True)
syms["ca_nsr"] = sp.Symbol("[Ca^{2+}]_{NSR}", nonnegative=True)
syms["r_ryr"] = sp.Symbol("r_{RyR}", nonnegative=True)
syms["r_xfer"] = sp.Symbol("r_{xfer}", nonnegative=True)
syms["p_cicr"] = sp.symbols("p_{0:4}", nonnegative=True)


def get_ical_jrel_jxfer_cass(
    x_cicr, ca_in, ca_jsr, vfrt, evfrtm1, cass_out,
        CA_OUT, V_SS, N_CARU=3e5, P_LCC=9.13e-13, RATE_XFER=220, RATE_RYR=16.6, CM=1, F=96485):

    """
    Get LCC current (iCaL), Ca flux from SR to subspace (jRel),
    Ca flux from subspace to cytosol (jXfer), and subspace calcium vector (caSS[4])

    - Arguments
    x_cicr[40]: CICR states vector
    
    """
    CA_OUT_TERM = 0.341 * CA_OUT
    # Probability distribution of subspace calcium (CaSS) vector elements
    p1 = x_cicr[2] + x_cicr[12] + x_cicr[32]  # LCC open
    p2 = (x_cicr[20] + x_cicr[21] + x_cicr[23] + x_cicr[24] + x_cicr[25] + 
          x_cicr[26] + x_cicr[27] + x_cicr[28] + x_cicr[29])  # RyR open
    p3 = x_cicr[22]  # Both open
    p0 = 1 - p1 - p2 - p3

    # Get subspace calcium concentration of each state from algebraic relations under quasi-steady state assumption
    e2vfrtm1 = evfrtm1 * (evfrtm1 + 1)
    e2vfrt = e2vfrtm1 + 1
    ca_out_weight = 2 * P_LCC / V_SS * vfrt / e2vfrtm1
    ca_out_value = CA_OUT_TERM * ca_out_weight
    ca_out_weight *= e2vfrt
    cain_weight = RATE_XFER
    ca_in_value = ca_in * cain_weight
    ca_jsr_weight = RATE_RYR
    ca_jsr_value = ca_jsr * ca_jsr_weight
    cass_out[0] = ca_in
    cass_out[1] = (ca_out_value + ca_in_value) / (ca_out_weight + cain_weight)
    cass_out[2] = (ca_jsr_value + ca_in_value) / (ca_jsr_weight + cain_weight)
    cass_out[3] = (ca_out_value + ca_in_value + ca_jsr_value) / (ca_out_weight + cain_weight + ca_jsr_weight)

    # LCC calcium current (iCal) (uA/cm^2)
    lcc_ca_diff = p1 * (cass_out[1] * e2vfrt - CA_OUT_TERM) + p3 * (cass_out[3] * e2vfrt - CA_OUT_TERM)
    i_cal = 4 * N_CARU / CM * P_LCC * F * vfrt * lcc_ca_diff / e2vfrtm1

    # Ca flux from SR to subspace (jRel) (mM/ms)
    j_rel = N_CARU * RATE_RYR * (p2 * (ca_jsr - cass_out[2]) + p3 * (ca_jsr - cass_out[3]))

    # Ca flux from subspace to cytosol (jXfer) (mM/ms)
    cass_avg = cass_out[0] * p0 + cass_out[1] * p1 + cass_out[2] * p2 + cass_out[3] * p3
    j_xfer = N_CARU * RATE_XFER * (cass_avg - ca_in)

    return (i_cal, j_rel, j_xfer, cass_out)


def _calc_markov(x, idx_a, idx_b, kf, kb, jac, nl):
    jac[idx_a] -= kf
    jac[idx_b] -= kb
    nl[idx_a] += x[idx_b] * kb
    nl[idx_b] += x[idx_a] * kf
    return jac, nl


def cicr_rates(exp):
    """
        Calculate the gradient (jacobian) and the non-linear part of CICR states

        Args:

            -x[40][N]: CICR states proportion
            -vm[N]: membrane potential (mV)
            -cass[4][N]: subspace calcium vector(mM)
        Returns:
            -jac_out[40][N]: gradient (dF/dy)
            -nonlinear_out[40][N]: nonlinear part (dy/dt - dF/dy)
    """
    def _cicr_rates(x, vm, ca_ss, jac_out, nl_out):
        # Constants
        F, G, A, B = 0.5, 2, 5, 7
        GAMMA0, OMEGA = 7.5, 0.06
        K0_01, K0_10 = 5265, 1500
        K0_12, K0_21 = 2.36E8, 9.6
        K0_23, K0_32 = 1.42E6, 13.65
        K0_14, K0_41 = 2.36E6, 0.0013
        K0_45, K0_54 = 1.89E7, 30
        K0_34, K0_43 = 0.07, 93.39

        # LCC transition rates
        alpha = 0.78113 * exp(0.08 * (vm - 8))
        alpha_a = alpha * A
        beta = 0.34319 * exp(-0.086676 * (vm - 8))
        beta_b = beta / B
        y_inf = 0.25 + 0.75 / (1 + exp((vm + 25) / 5))
        tau_y_inv = 1 / (17 + 500 / ((1 + exp((vm+28) / 10)) * (1 + exp(-(vm + 42) / 14))))
        kb = y_inf * tau_y_inv
        kf = (1 - y_inf) * tau_y_inv
        gamma = GAMMA0 * ca_ss
        gamma_a = gamma * A
        omega_b = OMEGA / B

        # RyR transition rates
        cass_sq = ca_ss ** 2
        k01 = K0_01 * cass_sq
        k10 = K0_10
        k12 = K0_12 * cass_sq
        k21 = K0_21 * K0_32 / (K0_23 * cass_sq + K0_32)
        k13 = K0_14 * cass_sq
        k31 = K0_41 * K0_54 / (K0_45 * cass_sq + K0_54)
        k23 = K0_34 * K0_23 * cass_sq / (K0_32 + K0_23 * cass_sq)
        k32 = K0_54 * K0_43 * cass_sq / (K0_54 + K0_45 * cass_sq)

        for r_idx in range(0, 40, 10):
            cass_idx = CASS_INDEX[r_idx]
            for i in range(5):
                jac_out, nl_out = _calc_markov(x, r_idx + i, r_idx + i + 5, kf, kb, jac_out, nl_out)
            for i in (0, 5):
                jac_out, nl_out = _calc_markov(x, r_idx + i, r_idx + i + 3, gamma[cass_idx], OMEGA, jac_out, nl_out)
                jac_out, nl_out = _calc_markov(x, r_idx + i, r_idx + i + 1, alpha, beta, jac_out, nl_out)
            for i in (1, 6):
                jac_out, nl_out = _calc_markov(x, r_idx + i, r_idx + i + 3, gamma_a[cass_idx], omega_b, jac_out, nl_out)
                jac_out, nl_out = _calc_markov(x, r_idx + i, r_idx + i + 1, F, G, jac_out, nl_out)
            for i in (3, 8):
                jac_out, nl_out = _calc_markov(x, r_idx + i, r_idx + i + 1, alpha_a, beta_b, jac_out, nl_out)

        for l_idx in range(10):
            jac_out, nl_out = _calc_markov(x, 0 + l_idx, 10 + l_idx, k01[CASS_INDEX[l_idx]], k10, jac_out, nl_out)
            jac_out, nl_out = _calc_markov(x, 10 + l_idx, 20 + l_idx,
                                        k12[CASS_INDEX[10 + l_idx]], k21[CASS_INDEX[20 + l_idx]], jac_out, nl_out)
            jac_out, nl_out = _calc_markov(x, 20 + l_idx, 30 + l_idx, k23[CASS_INDEX[20 + l_idx]],
                                        k32[CASS_INDEX[30 + l_idx]], jac_out, nl_out)
            jac_out, nl_out = _calc_markov(x, 10 + l_idx, 30 + l_idx, 
                                        k13[CASS_INDEX[10 + l_idx]], k31[CASS_INDEX[30 + l_idx]], jac_out, nl_out)

        return jac_out, nl_out
    return _cicr_rates
