# Size of variable arrays:
sizeAlgebraic = 29
sizeStates = 9
sizeConstants = 45
from math import *
from numpy import *

def createLegends():
    legend_states = [""] * sizeStates
    legend_rates = [""] * sizeStates
    legend_algebraic = [""] * sizeAlgebraic
    legend_voi = ""
    legend_constants = [""] * sizeConstants
    legend_voi = "time in component environment (millisecond)"
    legend_constants[0] = "RT in component parameters (joule_per_mole)"
    legend_states[0] = "V_m in component membrane (millivolt)"
    legend_constants[1] = "C_m in component membrane (picofarad)"
    legend_algebraic[16] = "I_Na in component I_Na (picoampere)"
    legend_algebraic[18] = "I_T in component I_T (picoampere)"
    legend_algebraic[20] = "I_KCa in component I_KCa (picoampere)"
    legend_algebraic[25] = "I_Ca in component I_Ca (picoampere)"
    legend_algebraic[26] = "I_L in component I_L (picoampere)"
    legend_algebraic[28] = "I_K in component I_K (picoampere)"
    legend_algebraic[13] = "I_stim in component I_stim (picoampere)"
    legend_constants[2] = "I_mag in component I_stim (picoampere)"
    legend_constants[3] = "t_0 in component I_stim (millisecond)"
    legend_constants[4] = "t_1 in component I_stim (millisecond)"
    legend_algebraic[3] = "H_0 in component I_stim (dimensionless)"
    legend_algebraic[9] = "H_1 in component I_stim (dimensionless)"
    legend_constants[5] = "g_Na in component I_Na (nanosiemens)"
    legend_constants[6] = "E_Na in component I_Na (millivolt)"
    legend_states[1] = "m in component m (dimensionless)"
    legend_states[2] = "h in component h (dimensionless)"
    legend_constants[7] = "E_m in component m (millivolt)"
    legend_constants[8] = "V_alpha_m in component m (millivolt)"
    legend_constants[9] = "V_beta_m in component m (millivolt)"
    legend_algebraic[0] = "alpha_m in component m (per_millisecond)"
    legend_algebraic[6] = "beta_m in component m (per_millisecond)"
    legend_constants[10] = "alpha_m_max in component m (per_millivolt_per_millisecond)"
    legend_constants[11] = "beta_m_max in component m (per_millisecond)"
    legend_constants[12] = "E_h in component h (millivolt)"
    legend_constants[13] = "V_alpha_h in component h (millivolt)"
    legend_constants[14] = "V_beta_h in component h (millivolt)"
    legend_algebraic[1] = "alpha_h in component h (per_millisecond)"
    legend_algebraic[7] = "beta_h in component h (per_millisecond)"
    legend_constants[15] = "alpha_h_max in component h (per_millisecond)"
    legend_constants[16] = "beta_h_max in component h (per_millisecond)"
    legend_constants[17] = "R_s in component I_T (per_nanosiemens)"
    legend_constants[18] = "C_t in component I_T (picofarad)"
    legend_states[3] = "V_T in component V_T (millivolt)"
    legend_constants[19] = "g_KCa in component I_KCa (nanosiemens)"
    legend_constants[20] = "E_K in component I_K (millivolt)"
    legend_states[4] = "o in component o (dimensionless)"
    legend_algebraic[19] = "w in component w (dimensionless)"
    legend_states[5] = "c in component calcium_handling (dimensionless)"
    legend_algebraic[15] = "o_oinf in component o (dimensionless)"
    legend_algebraic[2] = "alpha_Vm in component o (per_millisecond)"
    legend_algebraic[8] = "beta_Vm in component o (per_millisecond)"
    legend_algebraic[12] = "tau_Vm in component o (millisecond)"
    legend_constants[21] = "d_1 in component o (dimensionless)"
    legend_constants[22] = "d_2 in component o (dimensionless)"
    legend_constants[23] = "k_1 in component o (dimensionless)"
    legend_constants[24] = "k_2 in component o (dimensionless)"
    legend_constants[25] = "a_bar in component o (per_millisecond)"
    legend_constants[26] = "b_bar in component o (per_millisecond)"
    legend_constants[27] = "kd in component w (dimensionless)"
    legend_states[6] = "cer in component calcium_handling (dimensionless)"
    legend_algebraic[27] = "j_mem in component calcium_handling (picoampere)"
    legend_algebraic[21] = "j_leak in component calcium_handling (picoampere)"
    legend_algebraic[22] = "j_serca in component calcium_handling (picoampere)"
    legend_algebraic[23] = "j_er in component calcium_handling (picoampere)"
    legend_constants[28] = "alpha in component calcium_handling (dimensionless)"
    legend_constants[29] = "k_pmca in component calcium_handling (dimensionless)"
    legend_constants[30] = "k_serca in component calcium_handling (picoampere)"
    legend_constants[31] = "p_leak in component calcium_handling (picoampere)"
    legend_constants[32] = "f_er in component calcium_handling (per_picoampere_per_millisecond)"
    legend_constants[33] = "f_cyt in component calcium_handling (per_picoampere_per_millisecond)"
    legend_constants[34] = "v_cytver in component calcium_handling (dimensionless)"
    legend_states[7] = "d in component d (dimensionless)"
    legend_algebraic[24] = "g_Ca in component I_Ca (nanosiemens)"
    legend_constants[35] = "g_Ca0 in component I_Ca (nanosiemens_per_millivolt)"
    legend_constants[36] = "E_Ca in component I_Ca (millivolt)"
    legend_algebraic[14] = "alpha_d in component d (per_millisecond)"
    legend_algebraic[17] = "beta_d in component d (per_millisecond)"
    legend_algebraic[10] = "tau_d in component d (millisecond)"
    legend_algebraic[4] = "d_infinity in component d (dimensionless)"
    legend_constants[37] = "E_L in component I_L (millivolt)"
    legend_constants[38] = "g_Lmax in component I_L (nanosiemens)"
    legend_states[8] = "n in component n (dimensionless)"
    legend_constants[39] = "g_K in component I_K (nanosiemens)"
    legend_constants[40] = "E_n in component n (millivolt)"
    legend_constants[41] = "V_alpha_n in component n (millivolt)"
    legend_constants[42] = "V_beta_n in component n (millivolt)"
    legend_algebraic[5] = "alpha_n in component n (per_millisecond)"
    legend_algebraic[11] = "beta_n in component n (per_millisecond)"
    legend_constants[43] = "alpha_n_max in component n (per_millivolt_per_millisecond)"
    legend_constants[44] = "beta_n_max in component n (per_millisecond)"
    legend_rates[0] = "d/dt V_m in component membrane (millivolt)"
    legend_rates[1] = "d/dt m in component m (dimensionless)"
    legend_rates[2] = "d/dt h in component h (dimensionless)"
    legend_rates[3] = "d/dt V_T in component V_T (millivolt)"
    legend_rates[4] = "d/dt o in component o (dimensionless)"
    legend_rates[5] = "d/dt c in component calcium_handling (dimensionless)"
    legend_rates[6] = "d/dt cer in component calcium_handling (dimensionless)"
    legend_rates[7] = "d/dt d in component d (dimensionless)"
    legend_rates[8] = "d/dt n in component n (dimensionless)"
    return (legend_states, legend_algebraic, legend_voi, legend_constants)

def initConsts():
    constants = [0.0] * sizeConstants; states = [0.0] * sizeStates;
    constants[0] = 2577
    states[0] = -75
    constants[1] = 0.009
    constants[2] = 2
    constants[3] = 5
    constants[4] = 6
    constants[5] = 0.45
    constants[6] = 50
    states[1] = 0
    states[2] = 1
    constants[7] = -42
    constants[8] = 10
    constants[9] = 18
    constants[10] = 0.208
    constants[11] = 2.081
    constants[12] = -41
    constants[13] = 14.7
    constants[14] = 7.6
    constants[15] = 0.0156
    constants[16] = 3.382
    constants[17] = 15
    constants[18] = 0.04
    states[3] = -70
    constants[19] = 0.5
    constants[20] = -70
    states[4] = 0.15
    states[5] = 0.15
    constants[21] = 0.84
    constants[22] = 1
    constants[23] = 0.18
    constants[24] = 0.011
    constants[25] = 0.48
    constants[26] = 0.28
    constants[27] = 0.18
    states[6] = 200
    constants[28] = 4.5e-6
    constants[29] = 0.2
    constants[30] = 0.4
    constants[31] = 0.0005
    constants[32] = 0.01
    constants[33] = 0.01
    constants[34] = 5
    states[7] = 0
    constants[35] = 0.05
    constants[36] = 50
    constants[37] = -75
    constants[38] = 0.0024
    states[8] = 0
    constants[39] = 0.415
    constants[40] = -40
    constants[41] = 7
    constants[42] = 40
    constants[43] = 0.0229
    constants[44] = 0.09616
    return (states, constants)

def computeRates(voi, states, constants):
    rates = [0.0] * sizeStates; algebraic = [0.0] * sizeAlgebraic
    rates[3] = (states[0]-states[3])/(constants[17]*constants[18])
    algebraic[0] = (constants[10]*(states[0]-constants[7]))/(1.00000-exp((constants[7]-states[0])/constants[8]))
    algebraic[6] = constants[11]*exp((constants[7]-states[0])/constants[9])
    rates[1] = algebraic[0]*(1.00000-states[1])-algebraic[6]*states[1]
    algebraic[1] = constants[15]*exp((states[0]-constants[12])/constants[13])
    algebraic[7] = constants[16]/(1.00000+exp((constants[12]-states[0])/constants[14]))
    rates[2] = algebraic[1]*(1.00000-states[2])-algebraic[7]*states[2]
    algebraic[5] = (constants[43]*(states[0]-constants[40]))/(1.00000-exp((constants[40]-states[0])/constants[41]))
    algebraic[11] = constants[44]*exp((constants[40]-states[0])/constants[42])
    rates[8] = algebraic[5]*(1.00000-states[8])-algebraic[11]*states[8]
    algebraic[2] = constants[25]/(1.00000+constants[23]*exp((-2.00000*constants[21]*96.4850*states[0])/(constants[0]*states[5])))
    algebraic[8] = constants[26]/(1.00000+states[5]/(constants[24]*exp((-2.00000*constants[22]*96.4850*states[0])/constants[0])))
    algebraic[12] = 1.00000/(algebraic[2]+algebraic[8])
    algebraic[15] = algebraic[2]*algebraic[12]
    rates[4] = (algebraic[15]-states[4])/algebraic[12]
    algebraic[10] = (80.0000*1.00000)/cosh(-0.0310000*(states[0]+37.1000))
    algebraic[4] = 1.00000/(1.00000+exp((-24.6000-states[0])/11.3000))
    algebraic[14] = algebraic[4]/algebraic[10]
    algebraic[17] = (1.00000-algebraic[4])/algebraic[10]
    rates[7] = algebraic[14]*(1.00000-states[7])-algebraic[17]*states[7]
    algebraic[21] = constants[31]*(states[6]-states[5])
    algebraic[22] = constants[30]*states[5]
    algebraic[23] = algebraic[21]-algebraic[22]
    rates[6] = -constants[32]*constants[34]*algebraic[23]
    algebraic[24] = (-constants[35]*states[0])/(exp(0.117000*states[0])-1.00000)
    algebraic[25] = algebraic[24]*(power(states[7], 2.00000))*(states[0]-constants[36])
    algebraic[27] = -constants[28]*algebraic[25]*constants[29]*states[5]
    rates[5] = constants[33]*(algebraic[27]+algebraic[23])
    algebraic[16] = constants[5]*(power(states[1], 3.00000))*states[2]*(states[0]-constants[6])
    algebraic[18] = (states[0]-states[3])/constants[17]
    algebraic[19] = (power(states[5], 5.00000))/(power(states[5], 5.00000)+power(constants[27], 5.00000))
    algebraic[20] = constants[19]*states[4]*algebraic[19]*(states[0]-constants[20])
    algebraic[26] = constants[38]*(states[0]-constants[37])
    algebraic[28] = constants[39]*(power(states[8], 4.00000))*(states[0]-constants[20])
    algebraic[3] = custom_piecewise([less(voi , constants[3]), 0.00000 , True, 1.00000])
    algebraic[9] = custom_piecewise([less(voi , constants[4]), 0.00000 , True, 1.00000])
    algebraic[13] = constants[2]*(algebraic[3]-algebraic[9])
    rates[0] = (algebraic[13]-(algebraic[16]+algebraic[25]+algebraic[28]+algebraic[26]+algebraic[18]+algebraic[20]))/constants[1]
    return(rates)

def computeAlgebraic(constants, states, voi):
    algebraic = array([[0.0] * len(voi)] * sizeAlgebraic)
    states = array(states)
    voi = array(voi)
    algebraic[0] = (constants[10]*(states[0]-constants[7]))/(1.00000-exp((constants[7]-states[0])/constants[8]))
    algebraic[6] = constants[11]*exp((constants[7]-states[0])/constants[9])
    algebraic[1] = constants[15]*exp((states[0]-constants[12])/constants[13])
    algebraic[7] = constants[16]/(1.00000+exp((constants[12]-states[0])/constants[14]))
    algebraic[5] = (constants[43]*(states[0]-constants[40]))/(1.00000-exp((constants[40]-states[0])/constants[41]))
    algebraic[11] = constants[44]*exp((constants[40]-states[0])/constants[42])
    algebraic[2] = constants[25]/(1.00000+constants[23]*exp((-2.00000*constants[21]*96.4850*states[0])/(constants[0]*states[5])))
    algebraic[8] = constants[26]/(1.00000+states[5]/(constants[24]*exp((-2.00000*constants[22]*96.4850*states[0])/constants[0])))
    algebraic[12] = 1.00000/(algebraic[2]+algebraic[8])
    algebraic[15] = algebraic[2]*algebraic[12]
    algebraic[10] = (80.0000*1.00000)/cosh(-0.0310000*(states[0]+37.1000))
    algebraic[4] = 1.00000/(1.00000+exp((-24.6000-states[0])/11.3000))
    algebraic[14] = algebraic[4]/algebraic[10]
    algebraic[17] = (1.00000-algebraic[4])/algebraic[10]
    algebraic[21] = constants[31]*(states[6]-states[5])
    algebraic[22] = constants[30]*states[5]
    algebraic[23] = algebraic[21]-algebraic[22]
    algebraic[24] = (-constants[35]*states[0])/(exp(0.117000*states[0])-1.00000)
    algebraic[25] = algebraic[24]*(power(states[7], 2.00000))*(states[0]-constants[36])
    algebraic[27] = -constants[28]*algebraic[25]*constants[29]*states[5]
    algebraic[16] = constants[5]*(power(states[1], 3.00000))*states[2]*(states[0]-constants[6])
    algebraic[18] = (states[0]-states[3])/constants[17]
    algebraic[19] = (power(states[5], 5.00000))/(power(states[5], 5.00000)+power(constants[27], 5.00000))
    algebraic[20] = constants[19]*states[4]*algebraic[19]*(states[0]-constants[20])
    algebraic[26] = constants[38]*(states[0]-constants[37])
    algebraic[28] = constants[39]*(power(states[8], 4.00000))*(states[0]-constants[20])
    algebraic[3] = custom_piecewise([less(voi , constants[3]), 0.00000 , True, 1.00000])
    algebraic[9] = custom_piecewise([less(voi , constants[4]), 0.00000 , True, 1.00000])
    algebraic[13] = constants[2]*(algebraic[3]-algebraic[9])
    return algebraic

def custom_piecewise(cases):
    """Compute result of a piecewise function"""
    return select(cases[0::2],cases[1::2])

def solve_model():
    """Solve model with ODE solver"""
    from scipy.integrate import ode
    # Initialise constants and state variables
    (init_states, constants) = initConsts()

    # Set timespan to solve over
    voi = linspace(0, 10, 500)

    # Construct ODE object to solve
    r = ode(computeRates)
    r.set_integrator('vode', method='bdf', atol=1e-06, rtol=1e-06, max_step=1)
    r.set_initial_value(init_states, voi[0])
    r.set_f_params(constants)

    # Solve model
    states = array([[0.0] * len(voi)] * sizeStates)
    states[:,0] = init_states
    for (i,t) in enumerate(voi[1:]):
        if r.successful():
            r.integrate(t)
            states[:,i+1] = r.y
        else:
            break

    # Compute algebraic variables
    algebraic = computeAlgebraic(constants, states, voi)
    return (voi, states, algebraic)

def plot_model(voi, states, algebraic):
    """Plot variables against variable of integration"""
    import pylab
    (legend_states, legend_algebraic, legend_voi, legend_constants) = createLegends()
    pylab.figure(1)
    pylab.plot(voi,vstack((states,algebraic)).T)
    pylab.xlabel(legend_voi)
    pylab.legend(legend_states + legend_algebraic, loc='best')
    pylab.show()

if __name__ == "__main__":
    (voi, states, algebraic) = solve_model()
    plot_model(voi, states, algebraic)