# Size of variable arrays:
sizeAlgebraic = 35
sizeStates = 9
sizeConstants = 49
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 (second)"
    legend_constants[44] = "E_Na in component reversal_potentials (millivolt)"
    legend_algebraic[24] = "E_Ca in component reversal_potentials (millivolt)"
    legend_constants[45] = "E_K in component reversal_potentials (millivolt)"
    legend_constants[46] = "E_Cl in component reversal_potentials (millivolt)"
    legend_constants[0] = "R in component model_parameters (millijoule_per_mole_kelvin)"
    legend_constants[1] = "T in component model_parameters (kelvin)"
    legend_constants[2] = "F in component model_parameters (coulomb_per_mole)"
    legend_constants[3] = "Nai in component model_parameters (millimolar)"
    legend_constants[4] = "Nao in component model_parameters (millimolar)"
    legend_algebraic[22] = "Cai in component model_parameters (millimolar)"
    legend_constants[5] = "Cao in component model_parameters (millimolar)"
    legend_constants[6] = "Ki in component model_parameters (millimolar)"
    legend_constants[7] = "Ko in component model_parameters (millimolar)"
    legend_constants[8] = "Cli in component model_parameters (millimolar)"
    legend_constants[9] = "Clo in component model_parameters (millimolar)"
    legend_algebraic[0] = "i_Kdr in component delayed_rectifier_K_current (nanoA)"
    legend_constants[10] = "g_Kdr in component delayed_rectifier_K_current (microS)"
    legend_constants[11] = "gamma_KSS in component delayed_rectifier_K_current (dimensionless)"
    legend_states[0] = "V in component membrane (millivolt)"
    legend_states[1] = "xa in component delayed_rectifier_K_current_xa_gate (dimensionless)"
    legend_states[2] = "xi1 in component delayed_rectifier_K_current_xi1_gate (dimensionless)"
    legend_states[3] = "xi2 in component delayed_rectifier_K_current_xi2_gate (dimensionless)"
    legend_algebraic[1] = "xa_infinity in component delayed_rectifier_K_current_xa_gate (dimensionless)"
    legend_algebraic[8] = "tau_xa in component delayed_rectifier_K_current_xa_gate (second)"
    legend_algebraic[2] = "xi1_infinity in component delayed_rectifier_K_current_xi1_gate (dimensionless)"
    legend_constants[12] = "tau_xi1 in component delayed_rectifier_K_current_xi1_gate (second)"
    legend_algebraic[3] = "xi2_infinity in component delayed_rectifier_K_current_xi2_gate (dimensionless)"
    legend_constants[13] = "tau_xi2 in component delayed_rectifier_K_current_xi2_gate (second)"
    legend_algebraic[26] = "i_CaL in component L_type_Ca_current (nanoA)"
    legend_constants[14] = "P_CaL in component L_type_Ca_current (nanoA_per_millimolar)"
    legend_constants[15] = "V_surf in component L_type_Ca_current (millivolt)"
    legend_algebraic[7] = "ICaL_perm1 in component L_type_Ca_current (dimensionless)"
    legend_algebraic[25] = "ICaL_perm2 in component L_type_Ca_current (millimolar)"
    legend_states[4] = "d in component L_type_Ca_current_d_gate (dimensionless)"
    legend_states[5] = "f in component L_type_Ca_current_f_gate (dimensionless)"
    legend_algebraic[4] = "alpha_d in component L_type_Ca_current_d_gate (per_second)"
    legend_algebraic[9] = "beta_d in component L_type_Ca_current_d_gate (per_second)"
    legend_constants[16] = "tau_f in component L_type_Ca_current_f_gate (second)"
    legend_algebraic[5] = "f_infinity in component L_type_Ca_current_f_gate (dimensionless)"
    legend_algebraic[11] = "i_KCa in component Ca_activated_K_current (nanoA)"
    legend_constants[17] = "g_KCa in component Ca_activated_K_current (microS)"
    legend_states[6] = "xCa1 in component Ca_activated_K_current_xCa1_gate (dimensionless)"
    legend_states[7] = "B in component Ca_activated_K_current_B_gate (dimensionless)"
    legend_algebraic[23] = "xCa1_infinity in component Ca_activated_K_current_xCa1_gate (dimensionless)"
    legend_constants[18] = "tau_xCa1 in component Ca_activated_K_current_xCa1_gate (second)"
    legend_algebraic[13] = "K2 in component Ca_activated_K_current_xCa1_gate (millimolar)"
    legend_algebraic[12] = "K4 in component Ca_activated_K_current_xCa1_gate (millimolar)"
    legend_constants[19] = "alpha in component Ca_activated_K_current_xCa1_gate (per_second)"
    legend_constants[20] = "beta in component Ca_activated_K_current_xCa1_gate (per_second)"
    legend_algebraic[14] = "K1 in component Ca_activated_K_current_B_gate (per_millimolar_per_second)"
    legend_algebraic[15] = "K_1 in component Ca_activated_K_current_B_gate (per_second)"
    legend_algebraic[27] = "i_ClCa in component Ca_activated_Cl_current (nanoA)"
    legend_constants[21] = "g_Cl in component Ca_activated_Cl_current (microS)"
    legend_constants[22] = "CaCT in component Ca_activated_Cl_current (millimolar)"
    legend_constants[23] = "h in component Ca_activated_Cl_current (dimensionless)"
    legend_algebraic[32] = "i_cationic in component nonspecific_cationic_current (nanoA)"
    legend_algebraic[29] = "i_nsNa in component nonspecific_cationic_current (nanoA)"
    legend_algebraic[30] = "i_nsK in component nonspecific_cationic_current (nanoA)"
    legend_algebraic[31] = "i_nsCa in component nonspecific_cationic_current (nanoA)"
    legend_constants[24] = "P_nsCa in component nonspecific_cationic_current (nanoA_per_millimolar_mole_per_C)"
    legend_algebraic[28] = "alpha_Ca in component nonspecific_cationic_current (dimensionless)"
    legend_constants[25] = "Km_nsCa in component nonspecific_cationic_current (millimolar)"
    legend_states[8] = "ns in component nonspecific_cationic_current_ns_gate (dimensionless)"
    legend_algebraic[10] = "tau_ns in component nonspecific_cationic_current_ns_gate (second)"
    legend_algebraic[6] = "ns_infinity in component nonspecific_cationic_current_ns_gate (dimensionless)"
    legend_algebraic[33] = "i_pCa in component Ca_pump_current (nanoA)"
    legend_constants[26] = "i_pCamax in component Ca_pump_current (nanoA)"
    legend_constants[27] = "Kmp_Ca in component Ca_pump_current (millimolar)"
    legend_constants[47] = "i_NaK in component Na_K_pump_current (nanoA)"
    legend_constants[28] = "i_NaKmax in component Na_K_pump_current (nanoA)"
    legend_constants[29] = "Km_Na in component Na_K_pump_current (millimolar)"
    legend_constants[30] = "Km_K in component Na_K_pump_current (millimolar)"
    legend_algebraic[34] = "i_bCa in component Ca_background_current (nanoA)"
    legend_constants[31] = "g_bCa in component Ca_background_current (microS)"
    legend_algebraic[16] = "i_bK in component K_background_current (nanoA)"
    legend_constants[32] = "g_bK in component K_background_current (microS)"
    legend_algebraic[17] = "i_bNa in component Na_background_current (nanoA)"
    legend_constants[33] = "g_bNa in component Na_background_current (microS)"
    legend_constants[34] = "cap in component membrane (microF)"
    legend_algebraic[19] = "upeak in component model_parameters (millimolar)"
    legend_algebraic[18] = "up in component model_parameters (millimolar)"
    legend_algebraic[20] = "peak in component model_parameters (millimolar)"
    legend_algebraic[21] = "plat in component model_parameters (millimolar)"
    legend_constants[35] = "A in component model_parameters (millimolar)"
    legend_constants[36] = "B in component model_parameters (millimolar)"
    legend_constants[37] = "C in component model_parameters (millimolar)"
    legend_constants[48] = "x0 in component model_parameters (second)"
    legend_constants[38] = "x1 in component model_parameters (second)"
    legend_constants[39] = "x2 in component model_parameters (second)"
    legend_constants[40] = "ds in component model_parameters (second)"
    legend_constants[41] = "p in component model_parameters (second)"
    legend_constants[42] = "n in component model_parameters (second)"
    legend_constants[43] = "t36 in component model_parameters (second)"
    legend_rates[1] = "d/dt xa in component delayed_rectifier_K_current_xa_gate (dimensionless)"
    legend_rates[2] = "d/dt xi1 in component delayed_rectifier_K_current_xi1_gate (dimensionless)"
    legend_rates[3] = "d/dt xi2 in component delayed_rectifier_K_current_xi2_gate (dimensionless)"
    legend_rates[4] = "d/dt d in component L_type_Ca_current_d_gate (dimensionless)"
    legend_rates[5] = "d/dt f in component L_type_Ca_current_f_gate (dimensionless)"
    legend_rates[6] = "d/dt xCa1 in component Ca_activated_K_current_xCa1_gate (dimensionless)"
    legend_rates[7] = "d/dt B in component Ca_activated_K_current_B_gate (dimensionless)"
    legend_rates[8] = "d/dt ns in component nonspecific_cationic_current_ns_gate (dimensionless)"
    legend_rates[0] = "d/dt V in component membrane (millivolt)"
    return (legend_states, legend_algebraic, legend_voi, legend_constants)

def initConsts():
    constants = [0.0] * sizeConstants; states = [0.0] * sizeStates;
    constants[0] = 8314.472
    constants[1] = 310
    constants[2] = 96485.3415
    constants[3] = 12
    constants[4] = 130
    constants[5] = 2
    constants[6] = 150
    constants[7] = 5
    constants[8] = 55
    constants[9] = 140
    constants[10] = 0.035
    constants[11] = 0.15
    states[0] = -60
    states[1] = 0.001
    states[2] = 0
    states[3] = 0
    constants[12] = 0.25
    constants[13] = 1.98
    constants[14] = 0.003
    constants[15] = 150
    states[4] = 0
    states[5] = 1
    constants[16] = 0.0173
    constants[17] = 2.45
    states[6] = 1
    states[7] = 0.001
    constants[18] = 1
    constants[19] = 280
    constants[20] = 480
    constants[21] = 0.01
    constants[22] = 0.0005
    constants[23] = 3
    constants[24] = 0.000000175
    constants[25] = 0.0012
    states[8] = 1
    constants[26] = 1.15
    constants[27] = 0.05
    constants[28] = 0.7
    constants[29] = 40
    constants[30] = 1
    constants[31] = 1.194e-5
    constants[32] = 0.008729
    constants[33] = 0.003263
    constants[34] = 0.00002
    constants[35] = 0.0008
    constants[36] = 0.00008
    constants[37] = 0.00012
    constants[38] = 15
    constants[39] = 0.8
    constants[40] = 30
    constants[41] = 0.4
    constants[42] = 0.2
    constants[43] = 1.5
    constants[44] = ((constants[0]*constants[1])/constants[2])*log(constants[4]/constants[3])
    constants[45] = ((constants[0]*constants[1])/constants[2])*log(constants[7]/constants[6])
    constants[46] = ((constants[0]*constants[1])/constants[2])*log(constants[8]/constants[9])
    constants[47] = (((constants[28]*constants[7])/(constants[7]+constants[30]))*constants[3])/(constants[3]+constants[29])
    constants[48] = constants[38]+constants[39]
    return (states, constants)

def computeRates(voi, states, constants):
    rates = [0.0] * sizeStates; algebraic = [0.0] * sizeAlgebraic
    algebraic[2] = 1.00000/(1.00000+exp((states[0]+4.30000)/7.50000))
    rates[2] = (algebraic[2]-states[2])/constants[12]
    algebraic[3] = 1.00000/(1.00000+exp((states[0]+4.30000)/7.50000))
    rates[3] = (algebraic[3]-states[3])/constants[13]
    algebraic[5] = 1.00000/(1.00000+exp((23.0000+states[0])/6.60000))
    rates[5] = (algebraic[5]-states[5])/constants[16]
    algebraic[1] = 1.00000/(1.00000+exp((5.50000-states[0])/6.00000))
    algebraic[8] = 0.00326200-3.55200e-05*states[0]
    rates[1] = (algebraic[1]-states[1])/algebraic[8]
    algebraic[4] = (30.0000*(states[0]+18.0000))/(1.00000-exp(-0.250000*(states[0]+18.0000)))
    algebraic[9] = (12.0000*(states[0]+18.0000))/(exp(0.100000*(states[0]+18.0000))-1.00000)
    rates[4] = algebraic[4]-(algebraic[4]+algebraic[9])*states[4]
    algebraic[10] = 0.200250-0.000898750*states[0]
    algebraic[6] = 1.00000/(1.00000+exp((states[0]-69.8000)/-11.9000))
    rates[8] = (algebraic[6]-states[8])/algebraic[10]
    algebraic[18] = custom_piecewise([less(voi , constants[38]), 0.00000 , True, ((voi-constants[38])*constants[35])/constants[39]])
    algebraic[19] = custom_piecewise([less(voi , constants[38]+constants[39]), algebraic[18] , True, 0.00000])
    algebraic[20] = custom_piecewise([less(voi , constants[48]), 0.00000 , True, constants[35]*exp(-(voi-constants[48])/constants[43])])
    algebraic[21] = constants[36]/(1.00000+exp((voi-(constants[38]+constants[41]+constants[40]))/constants[42]))-constants[36]/(1.00000+exp((voi-(constants[38]+constants[41]))/constants[42]))
    algebraic[22] = constants[37]+algebraic[19]+algebraic[20]+algebraic[21]
    algebraic[14] = 0.850000*exp(0.0400000*states[0])
    algebraic[15] = 0.240000*exp(-0.0120000*states[0])
    rates[7] = custom_piecewise([greater_equal(algebraic[22] , 0.100000), algebraic[14]*algebraic[22]*states[6]-algebraic[15]*states[7] , True, 1.00000])
    algebraic[13] = 0.000275000*exp((-1.51000*states[0]*constants[2])/(constants[0]*constants[1]))
    algebraic[12] = 1.25000e-05*exp((-1.99000*states[0]*constants[2])/(constants[0]*constants[1]))
    algebraic[23] = (power(algebraic[22], 2.00000)+algebraic[12]*algebraic[22])/(power(algebraic[22], 2.00000)+algebraic[12]*algebraic[22]*(1.00000+constants[19]/constants[20])+(algebraic[12]*algebraic[13]*constants[19])/constants[20])
    rates[6] = (algebraic[23]-states[6])/constants[18]
    algebraic[0] = constants[10]*(constants[11]+(states[2]+states[3])*(1.00000-constants[11]))*(power(states[1], 2.00000))*(states[0]-constants[45])
    algebraic[7] = ((states[0]-constants[15])*constants[2])/(constants[0]*constants[1]*(1.00000-exp((-(states[0]-constants[15])*2.00000*constants[2])/(constants[0]*constants[1]))))
    algebraic[25] = algebraic[7]*(algebraic[22]*exp(constants[15]/((0.500000*constants[0]*constants[1])/constants[2]))-constants[5]*exp(-(states[0]-constants[15])/((0.500000*constants[0]*constants[1])/constants[2])))
    algebraic[26] = 4.00000*constants[14]*states[4]*states[5]*algebraic[25]
    algebraic[11] = constants[17]*states[6]*states[7]*(states[0]-constants[45])
    algebraic[27] = (constants[21]*(states[0]-constants[46])*1.00000)/(1.00000+power(constants[22]/algebraic[22], constants[23]))
    algebraic[28] = 1.00000/(1.00000+power(constants[25]/algebraic[22], 3.00000))
    algebraic[29] = (((states[8]*algebraic[28]*constants[24]*states[0]*(power(constants[2], 2.00000)))/(constants[0]*constants[1]))*(0.750000*constants[3]*exp((states[0]*constants[2])/(constants[0]*constants[1]))-0.750000*constants[4]))/(exp((states[0]*constants[2])/(constants[0]*constants[1]))-1.00000)
    algebraic[30] = (((states[8]*algebraic[28]*constants[24]*states[0]*(power(constants[2], 2.00000)))/(constants[0]*constants[1]))*(0.750000*constants[6]*exp((states[0]*constants[2])/(constants[0]*constants[1]))-0.750000*constants[7]))/(exp((states[0]*constants[2])/(constants[0]*constants[1]))-1.00000)
    algebraic[31] = (((states[8]*algebraic[28]*constants[24]*4.00000*states[0]*(power(constants[2], 2.00000)))/(constants[0]*constants[1]))*(1.50000*algebraic[22]*exp((states[0]*constants[2])/(constants[0]*constants[1]))-1.50000*constants[5]))/(exp((states[0]*constants[2])/(constants[0]*constants[1]))-1.00000)
    algebraic[32] = algebraic[29]+algebraic[30]+algebraic[31]
    algebraic[33] = (constants[26]*algebraic[22])/(constants[27]+algebraic[22])
    algebraic[24] = ((constants[0]*constants[1])/(2.00000*constants[2]))*log(constants[5]/algebraic[22])
    algebraic[34] = constants[31]*(states[0]-algebraic[24])
    algebraic[16] = constants[32]*(states[0]-constants[45])
    algebraic[17] = constants[33]*(states[0]-constants[44])
    rates[0] = (-1.00000/constants[34])*(algebraic[26]+algebraic[0]+algebraic[11]+algebraic[27]+algebraic[32]+constants[47]+algebraic[33]+algebraic[34]+algebraic[16]+algebraic[17])
    return(rates)

def computeAlgebraic(constants, states, voi):
    algebraic = array([[0.0] * len(voi)] * sizeAlgebraic)
    states = array(states)
    voi = array(voi)
    algebraic[2] = 1.00000/(1.00000+exp((states[0]+4.30000)/7.50000))
    algebraic[3] = 1.00000/(1.00000+exp((states[0]+4.30000)/7.50000))
    algebraic[5] = 1.00000/(1.00000+exp((23.0000+states[0])/6.60000))
    algebraic[1] = 1.00000/(1.00000+exp((5.50000-states[0])/6.00000))
    algebraic[8] = 0.00326200-3.55200e-05*states[0]
    algebraic[4] = (30.0000*(states[0]+18.0000))/(1.00000-exp(-0.250000*(states[0]+18.0000)))
    algebraic[9] = (12.0000*(states[0]+18.0000))/(exp(0.100000*(states[0]+18.0000))-1.00000)
    algebraic[10] = 0.200250-0.000898750*states[0]
    algebraic[6] = 1.00000/(1.00000+exp((states[0]-69.8000)/-11.9000))
    algebraic[18] = custom_piecewise([less(voi , constants[38]), 0.00000 , True, ((voi-constants[38])*constants[35])/constants[39]])
    algebraic[19] = custom_piecewise([less(voi , constants[38]+constants[39]), algebraic[18] , True, 0.00000])
    algebraic[20] = custom_piecewise([less(voi , constants[48]), 0.00000 , True, constants[35]*exp(-(voi-constants[48])/constants[43])])
    algebraic[21] = constants[36]/(1.00000+exp((voi-(constants[38]+constants[41]+constants[40]))/constants[42]))-constants[36]/(1.00000+exp((voi-(constants[38]+constants[41]))/constants[42]))
    algebraic[22] = constants[37]+algebraic[19]+algebraic[20]+algebraic[21]
    algebraic[14] = 0.850000*exp(0.0400000*states[0])
    algebraic[15] = 0.240000*exp(-0.0120000*states[0])
    algebraic[13] = 0.000275000*exp((-1.51000*states[0]*constants[2])/(constants[0]*constants[1]))
    algebraic[12] = 1.25000e-05*exp((-1.99000*states[0]*constants[2])/(constants[0]*constants[1]))
    algebraic[23] = (power(algebraic[22], 2.00000)+algebraic[12]*algebraic[22])/(power(algebraic[22], 2.00000)+algebraic[12]*algebraic[22]*(1.00000+constants[19]/constants[20])+(algebraic[12]*algebraic[13]*constants[19])/constants[20])
    algebraic[0] = constants[10]*(constants[11]+(states[2]+states[3])*(1.00000-constants[11]))*(power(states[1], 2.00000))*(states[0]-constants[45])
    algebraic[7] = ((states[0]-constants[15])*constants[2])/(constants[0]*constants[1]*(1.00000-exp((-(states[0]-constants[15])*2.00000*constants[2])/(constants[0]*constants[1]))))
    algebraic[25] = algebraic[7]*(algebraic[22]*exp(constants[15]/((0.500000*constants[0]*constants[1])/constants[2]))-constants[5]*exp(-(states[0]-constants[15])/((0.500000*constants[0]*constants[1])/constants[2])))
    algebraic[26] = 4.00000*constants[14]*states[4]*states[5]*algebraic[25]
    algebraic[11] = constants[17]*states[6]*states[7]*(states[0]-constants[45])
    algebraic[27] = (constants[21]*(states[0]-constants[46])*1.00000)/(1.00000+power(constants[22]/algebraic[22], constants[23]))
    algebraic[28] = 1.00000/(1.00000+power(constants[25]/algebraic[22], 3.00000))
    algebraic[29] = (((states[8]*algebraic[28]*constants[24]*states[0]*(power(constants[2], 2.00000)))/(constants[0]*constants[1]))*(0.750000*constants[3]*exp((states[0]*constants[2])/(constants[0]*constants[1]))-0.750000*constants[4]))/(exp((states[0]*constants[2])/(constants[0]*constants[1]))-1.00000)
    algebraic[30] = (((states[8]*algebraic[28]*constants[24]*states[0]*(power(constants[2], 2.00000)))/(constants[0]*constants[1]))*(0.750000*constants[6]*exp((states[0]*constants[2])/(constants[0]*constants[1]))-0.750000*constants[7]))/(exp((states[0]*constants[2])/(constants[0]*constants[1]))-1.00000)
    algebraic[31] = (((states[8]*algebraic[28]*constants[24]*4.00000*states[0]*(power(constants[2], 2.00000)))/(constants[0]*constants[1]))*(1.50000*algebraic[22]*exp((states[0]*constants[2])/(constants[0]*constants[1]))-1.50000*constants[5]))/(exp((states[0]*constants[2])/(constants[0]*constants[1]))-1.00000)
    algebraic[32] = algebraic[29]+algebraic[30]+algebraic[31]
    algebraic[33] = (constants[26]*algebraic[22])/(constants[27]+algebraic[22])
    algebraic[24] = ((constants[0]*constants[1])/(2.00000*constants[2]))*log(constants[5]/algebraic[22])
    algebraic[34] = constants[31]*(states[0]-algebraic[24])
    algebraic[16] = constants[32]*(states[0]-constants[45])
    algebraic[17] = constants[33]*(states[0]-constants[44])
    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)