# Size of variable arrays:
sizeAlgebraic = 23
sizeStates = 5
sizeConstants = 22
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_states[0] = "V in component membrane (millivolt)"
    legend_constants[0] = "RToF in component membrane (millivolt)"
    legend_constants[1] = "K_c in component membrane (millimolar)"
    legend_constants[2] = "K_i in component membrane (millimolar)"
    legend_constants[3] = "Na_c in component membrane (millimolar)"
    legend_constants[4] = "Na_i in component membrane (millimolar)"
    legend_algebraic[5] = "i_ns in component non_specific_current (picoA_per_picoF)"
    legend_algebraic[10] = "i_to in component transient_outward_K_current (picoA_per_picoF)"
    legend_algebraic[20] = "i_K1 in component inward_rectifier (picoA_per_picoF)"
    legend_algebraic[15] = "i_Kur in component ultra_rapid_K_current (picoA_per_picoF)"
    legend_algebraic[21] = "i_b_Na in component background_currents (picoA_per_picoF)"
    legend_algebraic[22] = "i_NaK in component sodium_potassium_pump (picoA_per_picoF)"
    legend_algebraic[0] = "i_Stim in component membrane (picoA_per_picoF)"
    legend_constants[5] = "stim_start in component membrane (millisecond)"
    legend_constants[6] = "stim_end in component membrane (millisecond)"
    legend_constants[7] = "stim_period in component membrane (millisecond)"
    legend_constants[8] = "stim_duration in component membrane (millisecond)"
    legend_constants[9] = "stim_amplitude in component membrane (picoA_per_picoF)"
    legend_constants[10] = "g_ns in component non_specific_current (nanoS_per_picoF)"
    legend_constants[20] = "E_K in component transient_outward_K_current (millivolt)"
    legend_constants[11] = "g_to in component transient_outward_K_current (nanoS_per_picoF)"
    legend_states[1] = "oa in component transient_outward_K_current_oa_gate (dimensionless)"
    legend_states[2] = "oi in component transient_outward_K_current_oi_gate (dimensionless)"
    legend_algebraic[1] = "alpha_oa in component transient_outward_K_current_oa_gate (per_millisecond)"
    legend_algebraic[6] = "beta_oa in component transient_outward_K_current_oa_gate (per_millisecond)"
    legend_algebraic[16] = "tau_oa in component transient_outward_K_current_oa_gate (millisecond)"
    legend_algebraic[11] = "oa_inf in component transient_outward_K_current_oa_gate (dimensionless)"
    legend_algebraic[2] = "alpha_oi in component transient_outward_K_current_oi_gate (per_millisecond)"
    legend_algebraic[7] = "beta_oi in component transient_outward_K_current_oi_gate (per_millisecond)"
    legend_algebraic[17] = "tau_oi in component transient_outward_K_current_oi_gate (millisecond)"
    legend_algebraic[12] = "oi_inf in component transient_outward_K_current_oi_gate (dimensionless)"
    legend_constants[12] = "g_Kur in component ultra_rapid_K_current (nanoS_per_picoF)"
    legend_states[3] = "ua in component ultra_rapid_K_current_ua_gate (dimensionless)"
    legend_states[4] = "ui in component ultra_rapid_K_current_ui_gate (dimensionless)"
    legend_algebraic[3] = "alpha_ua in component ultra_rapid_K_current_ua_gate (per_millisecond)"
    legend_algebraic[8] = "beta_ua in component ultra_rapid_K_current_ua_gate (per_millisecond)"
    legend_algebraic[18] = "tau_ua in component ultra_rapid_K_current_ua_gate (millisecond)"
    legend_algebraic[13] = "ua_inf in component ultra_rapid_K_current_ua_gate (dimensionless)"
    legend_algebraic[4] = "alpha_ui in component ultra_rapid_K_current_ui_gate (per_millisecond)"
    legend_algebraic[9] = "beta_ui in component ultra_rapid_K_current_ui_gate (per_millisecond)"
    legend_algebraic[19] = "tau_ui in component ultra_rapid_K_current_ui_gate (millisecond)"
    legend_algebraic[14] = "ui_inf in component ultra_rapid_K_current_ui_gate (dimensionless)"
    legend_constants[13] = "g_K1 in component inward_rectifier (nanoS_per_picoF)"
    legend_constants[14] = "g_b_Na in component background_currents (nanoS_per_picoF)"
    legend_constants[21] = "E_Na in component background_currents (millivolt)"
    legend_constants[15] = "k_NaK_K in component sodium_potassium_pump (millimolar)"
    legend_constants[16] = "k_NaK_Na in component sodium_potassium_pump (millimolar)"
    legend_constants[17] = "i_NaK_max in component sodium_potassium_pump (picoA_per_picoF)"
    legend_constants[18] = "V_rev in component sodium_potassium_pump (millivolt)"
    legend_constants[19] = "B in component sodium_potassium_pump (millivolt)"
    legend_rates[0] = "d/dt V in component membrane (millivolt)"
    legend_rates[1] = "d/dt oa in component transient_outward_K_current_oa_gate (dimensionless)"
    legend_rates[2] = "d/dt oi in component transient_outward_K_current_oi_gate (dimensionless)"
    legend_rates[3] = "d/dt ua in component ultra_rapid_K_current_ua_gate (dimensionless)"
    legend_rates[4] = "d/dt ui in component ultra_rapid_K_current_ui_gate (dimensionless)"
    return (legend_states, legend_algebraic, legend_voi, legend_constants)

def initConsts():
    constants = [0.0] * sizeConstants; states = [0.0] * sizeStates;
    states[0] = -80
    constants[0] = 26.54
    constants[1] = 5.4
    constants[2] = 139
    constants[3] = 140
    constants[4] = 11.2
    constants[5] = 100
    constants[6] = 10000
    constants[7] = 5000
    constants[8] = 2
    constants[9] = -100
    constants[10] = 0.018
    constants[11] = 0.01652
    states[1] = 0.2488
    states[2] = 0.1866
    constants[12] = 0.6
    states[3] = 0.9933
    states[4] = 0.0999
    constants[13] = 0.03
    constants[14] = 0.00607
    constants[15] = 1
    constants[16] = 11
    constants[17] = 2.002
    constants[18] = -150
    constants[19] = -200
    constants[20] = constants[0]*log(constants[1]/constants[2])
    constants[21] = constants[0]*log(constants[3]/constants[4])
    return (states, constants)

def computeRates(voi, states, constants):
    rates = [0.0] * sizeStates; algebraic = [0.0] * sizeAlgebraic
    algebraic[1] = 0.650000/(exp((states[0]+10.0000)/-8.50000)+exp((states[0]-30.0000)/-59.0000))
    algebraic[6] = 0.650000/(2.50000+exp((states[0]+82.0000)/17.0000))
    algebraic[16] = 15.0000/(algebraic[1]+algebraic[6])
    algebraic[11] = 1.00000/(1.00000+exp((states[0]+20.4700)/-17.5400))
    rates[1] = (algebraic[11]-states[1])/algebraic[16]
    algebraic[2] = 1.00000/(18.5300+exp((states[0]+113.700)/10.9500))
    algebraic[7] = 1.00000/(35.5600+exp((states[0]+1.26000)/-7.44000))
    algebraic[17] = 15.0000/(algebraic[2]+algebraic[7])
    algebraic[12] = 1.00000/(1.00000+exp((states[0]+43.1000)/5.30000))
    rates[2] = (algebraic[12]-states[2])/algebraic[17]
    algebraic[3] = 0.650000/(exp((states[0]+10.0000)/-8.50000)+exp((states[0]-30.0000)/-59.0000))
    algebraic[8] = 0.650000/(2.50000+exp((states[0]+82.0000)/17.0000))
    algebraic[18] = 1.00000/(algebraic[3]+algebraic[8])
    algebraic[13] = 1.00000/(1.00000+exp((states[0]+33.3000)/-9.60000))
    rates[3] = (algebraic[13]-states[3])/algebraic[18]
    algebraic[4] = 1.00000/(21.0000+exp((states[0]-185.000)/28.0000))
    algebraic[9] = 1.00000/exp((states[0]-158.000)/-16.0000)
    algebraic[19] = 5.00000/(algebraic[4]+algebraic[9])
    algebraic[14] = 1.00000/(1.00000+exp((states[0]-99.4500)/27.4800))
    rates[4] = (algebraic[14]-states[4])/algebraic[19]
    algebraic[5] = constants[10]*states[0]
    algebraic[10] = constants[11]*(power(states[1], 3.00000))*states[2]*(states[0]-constants[20])
    algebraic[20] = (constants[13]*(states[0]+86.7500))/(1.00000+exp((states[0]+20.0000)/20.0000))
    algebraic[15] = (constants[12]*0.00500000+0.0500000/(1.00000+exp(-(states[0]-15.0000)/13.0000)))*(power(states[3], 3.00000))*states[4]*(states[0]-constants[20])
    algebraic[21] = constants[14]*(states[0]-constants[21])
    algebraic[22] = (((((constants[17]*constants[1])/(constants[1]+constants[15]))*(power(constants[4], 1.50000)))/(power(constants[4], 1.50000)+power(constants[16], 1.50000)))*(states[0]-constants[18]))/(states[0]-constants[19])
    algebraic[0] = custom_piecewise([greater_equal(voi , constants[5]) & less_equal(voi , constants[6]) & less_equal((voi-constants[5])-floor((voi-constants[5])/constants[7])*constants[7] , constants[8]), constants[9] , True, 0.00000])
    rates[0] = -(algebraic[5]+algebraic[15]+algebraic[10]+algebraic[20]+algebraic[21]+algebraic[22]+algebraic[0])
    return(rates)

def computeAlgebraic(constants, states, voi):
    algebraic = array([[0.0] * len(voi)] * sizeAlgebraic)
    states = array(states)
    voi = array(voi)
    algebraic[1] = 0.650000/(exp((states[0]+10.0000)/-8.50000)+exp((states[0]-30.0000)/-59.0000))
    algebraic[6] = 0.650000/(2.50000+exp((states[0]+82.0000)/17.0000))
    algebraic[16] = 15.0000/(algebraic[1]+algebraic[6])
    algebraic[11] = 1.00000/(1.00000+exp((states[0]+20.4700)/-17.5400))
    algebraic[2] = 1.00000/(18.5300+exp((states[0]+113.700)/10.9500))
    algebraic[7] = 1.00000/(35.5600+exp((states[0]+1.26000)/-7.44000))
    algebraic[17] = 15.0000/(algebraic[2]+algebraic[7])
    algebraic[12] = 1.00000/(1.00000+exp((states[0]+43.1000)/5.30000))
    algebraic[3] = 0.650000/(exp((states[0]+10.0000)/-8.50000)+exp((states[0]-30.0000)/-59.0000))
    algebraic[8] = 0.650000/(2.50000+exp((states[0]+82.0000)/17.0000))
    algebraic[18] = 1.00000/(algebraic[3]+algebraic[8])
    algebraic[13] = 1.00000/(1.00000+exp((states[0]+33.3000)/-9.60000))
    algebraic[4] = 1.00000/(21.0000+exp((states[0]-185.000)/28.0000))
    algebraic[9] = 1.00000/exp((states[0]-158.000)/-16.0000)
    algebraic[19] = 5.00000/(algebraic[4]+algebraic[9])
    algebraic[14] = 1.00000/(1.00000+exp((states[0]-99.4500)/27.4800))
    algebraic[5] = constants[10]*states[0]
    algebraic[10] = constants[11]*(power(states[1], 3.00000))*states[2]*(states[0]-constants[20])
    algebraic[20] = (constants[13]*(states[0]+86.7500))/(1.00000+exp((states[0]+20.0000)/20.0000))
    algebraic[15] = (constants[12]*0.00500000+0.0500000/(1.00000+exp(-(states[0]-15.0000)/13.0000)))*(power(states[3], 3.00000))*states[4]*(states[0]-constants[20])
    algebraic[21] = constants[14]*(states[0]-constants[21])
    algebraic[22] = (((((constants[17]*constants[1])/(constants[1]+constants[15]))*(power(constants[4], 1.50000)))/(power(constants[4], 1.50000)+power(constants[16], 1.50000)))*(states[0]-constants[18]))/(states[0]-constants[19])
    algebraic[0] = custom_piecewise([greater_equal(voi , constants[5]) & less_equal(voi , constants[6]) & less_equal((voi-constants[5])-floor((voi-constants[5])/constants[7])*constants[7] , constants[8]), constants[9] , True, 0.00000])
    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)