# Size of variable arrays:
sizeAlgebraic = 29
sizeStates = 9
sizeConstants = 48
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] = "isotonic_mode in component parameters (dimensionless)"
    legend_constants[1] = "alpha_1 in component parameters (per_um)"
    legend_constants[2] = "beta_1 in component parameters (mN)"
    legend_constants[3] = "alpha_2 in component parameters (per_um)"
    legend_constants[4] = "beta_2 in component parameters (mN)"
    legend_constants[5] = "alpha_3 in component parameters (per_um)"
    legend_constants[6] = "beta_3 in component parameters (mN)"
    legend_constants[7] = "lambda in component parameters (mN)"
    legend_constants[8] = "A_half in component parameters (dimensionless)"
    legend_constants[9] = "mu in component parameters (dimensionless)"
    legend_constants[10] = "chi in component parameters (dimensionless)"
    legend_constants[11] = "chi_0 in component parameters (dimensionless)"
    legend_constants[12] = "m_0 in component parameters (dimensionless)"
    legend_constants[13] = "v_max in component parameters (um_per_msec)"
    legend_constants[14] = "a in component parameters (dimensionless)"
    legend_constants[15] = "d_h in component parameters (dimensionless)"
    legend_constants[16] = "alpha_P in component parameters (dimensionless)"
    legend_constants[17] = "alpha_PG in component parameters (dimensionless)"
    legend_constants[18] = "S_0 in component parameters_izakov_et_al_1991 (um)"
    legend_algebraic[3] = "q_v in component parameters_izakov_et_al_1991 (per_millisecond)"
    legend_constants[19] = "q_1 in component parameters_izakov_et_al_1991 (per_millisecond)"
    legend_constants[20] = "q_2 in component parameters_izakov_et_al_1991 (per_millisecond)"
    legend_constants[21] = "q_3 in component parameters_izakov_et_al_1991 (per_millisecond)"
    legend_constants[45] = "v_1 in component parameters_izakov_et_al_1991 (um_per_msec)"
    legend_constants[22] = "alpha_P in component parameters_izakov_et_al_1991 (per_um)"
    legend_constants[23] = "alpha_S in component parameters_izakov_et_al_1991 (per_um)"
    legend_constants[24] = "alpha_G in component parameters_izakov_et_al_1991 (dimensionless)"
    legend_constants[25] = "a_on in component parameters_izakov_et_al_1991 (per_millisecond)"
    legend_constants[26] = "a_off in component parameters_izakov_et_al_1991 (per_millisecond)"
    legend_constants[27] = "k_A in component parameters_izakov_et_al_1991 (dimensionless)"
    legend_states[0] = "v in component CE_velocity (um_per_msec)"
    legend_algebraic[26] = "F_CE in component force (mN)"
    legend_algebraic[5] = "F_muscle in component force (mN)"
    legend_algebraic[0] = "F_XSE in component force (mN)"
    legend_algebraic[1] = "F_SE in component force (mN)"
    legend_algebraic[2] = "F_PE in component force (mN)"
    legend_states[1] = "N in component crossbridge_kinetics (dimensionless)"
    legend_algebraic[13] = "k_P_vis in component CE_velocity (mN)"
    legend_algebraic[15] = "k_S_vis in component PE_velocity (mN)"
    legend_states[2] = "w in component PE_velocity (um_per_msec)"
    legend_states[3] = "l_1 in component length (um)"
    legend_states[4] = "l_2 in component length (um)"
    legend_states[5] = "l_3 in component length (um)"
    legend_algebraic[25] = "p_v in component average_crossbridge_force (dimensionless)"
    legend_algebraic[22] = "K_chi in component crossbridge_kinetics (per_millisecond)"
    legend_algebraic[6] = "M_A in component crossbridge_kinetics (dimensionless)"
    legend_algebraic[7] = "n_1 in component crossbridge_kinetics (dimensionless)"
    legend_algebraic[8] = "L_oz in component crossbridge_kinetics (dimensionless)"
    legend_algebraic[20] = "k_p_v in component crossbridge_kinetics (per_millisecond)"
    legend_algebraic[21] = "k_m_v in component crossbridge_kinetics (per_millisecond)"
    legend_states[6] = "A in component calcium_handling (dimensionless)"
    legend_algebraic[19] = "G_star in component average_crossbridge_force (dimensionless)"
    legend_algebraic[9] = "dl_1_dt in component length (um_per_msec)"
    legend_algebraic[11] = "dl_2_dt in component length (um_per_msec)"
    legend_algebraic[4] = "dl_3_dt in component length (um_per_msec)"
    legend_algebraic[27] = "phi_chi in component CE_velocity (um_per_msec2)"
    legend_algebraic[28] = "p_prime_v in component average_crossbridge_force (per_millisecond)"
    legend_constants[28] = "alpha_P_lengthening in component CE_velocity (per_um)"
    legend_constants[29] = "beta_P_lengthening in component CE_velocity (mN)"
    legend_constants[30] = "alpha_P_shortening in component CE_velocity (per_um)"
    legend_constants[31] = "beta_P_shortening in component CE_velocity (mN)"
    legend_constants[32] = "alpha_S_lengthening in component PE_velocity (per_um)"
    legend_constants[33] = "beta_S_lengthening in component PE_velocity (mN)"
    legend_constants[34] = "alpha_S_shortening in component PE_velocity (per_um)"
    legend_constants[35] = "beta_S_shortening in component PE_velocity (mN)"
    legend_algebraic[18] = "P_star in component average_crossbridge_force (dimensionless)"
    legend_algebraic[17] = "gamma in component average_crossbridge_force (dimensionless)"
    legend_constants[46] = "case_1 in component average_crossbridge_force (per_millisecond)"
    legend_algebraic[23] = "case_2 in component average_crossbridge_force (per_millisecond)"
    legend_constants[47] = "case_3 in component average_crossbridge_force (per_millisecond)"
    legend_algebraic[24] = "case_4 in component average_crossbridge_force (per_millisecond)"
    legend_algebraic[14] = "dA_dt in component calcium_handling (per_millisecond)"
    legend_algebraic[10] = "N_A in component calcium_handling (dimensionless)"
    legend_algebraic[12] = "pi_N_A in component calcium_handling (dimensionless)"
    legend_states[7] = "B in component calcium_handling (dimensionless)"
    legend_algebraic[16] = "dB_dt in component calcium_handling (per_millisecond)"
    legend_states[8] = "Ca_C in component calcium_handling (dimensionless)"
    legend_constants[36] = "A_tot in component calcium_handling (dimensionless)"
    legend_constants[37] = "B_tot in component calcium_handling (dimensionless)"
    legend_constants[38] = "b_on in component calcium_handling (per_millisecond)"
    legend_constants[39] = "b_off in component calcium_handling (per_millisecond)"
    legend_constants[40] = "a_c in component calcium_handling (per_millisecond2)"
    legend_constants[41] = "b_c in component calcium_handling (per_millisecond2)"
    legend_constants[42] = "r_Ca in component calcium_handling (per_millisecond)"
    legend_constants[43] = "q_Ca in component calcium_handling (dimensionless)"
    legend_constants[44] = "t_d in component calcium_handling (millisecond)"
    legend_rates[1] = "d/dt N in component crossbridge_kinetics (dimensionless)"
    legend_rates[3] = "d/dt l_1 in component length (um)"
    legend_rates[4] = "d/dt l_2 in component length (um)"
    legend_rates[5] = "d/dt l_3 in component length (um)"
    legend_rates[0] = "d/dt v in component CE_velocity (um_per_msec)"
    legend_rates[2] = "d/dt w in component PE_velocity (um_per_msec)"
    legend_rates[6] = "d/dt A in component calcium_handling (dimensionless)"
    legend_rates[7] = "d/dt B in component calcium_handling (dimensionless)"
    legend_rates[8] = "d/dt Ca_C in component calcium_handling (dimensionless)"
    return (legend_states, legend_algebraic, legend_voi, legend_constants)

def initConsts():
    constants = [0.0] * sizeConstants; states = [0.0] * sizeStates;
    constants[0] = 1
    constants[1] = 19
    constants[2] = 2.9
    constants[3] = 14.6
    constants[4] = 0.01
    constants[5] = 48
    constants[6] = 0.1
    constants[7] = 960
    constants[8] = 0.6
    constants[9] = 3
    constants[10] = 0.705
    constants[11] = 3
    constants[12] = 0.9
    constants[13] = 0.0056
    constants[14] = 0.25
    constants[15] = 0.5
    constants[16] = 4
    constants[17] = 1
    constants[18] = 0.77
    constants[19] = 0.017
    constants[20] = 0.26
    constants[21] = 0.03
    constants[22] = 4
    constants[23] = 4
    constants[24] = 4
    constants[25] = 2.9e-2
    constants[26] = 0.2
    constants[27] = 2.8
    states[0] = 0
    states[1] = 1
    states[2] = 0
    states[3] = 1
    states[4] = 1
    states[5] = 1
    states[6] = 0
    constants[28] = 16
    constants[29] = 15
    constants[30] = 16
    constants[31] = 15
    constants[32] = 39
    constants[33] = 80
    constants[34] = 46
    constants[35] = 60
    states[7] = 0
    states[8] = 0
    constants[36] = 1
    constants[37] = 0.4
    constants[38] = 2.6
    constants[39] = 0.182
    constants[40] = 0.005
    constants[41] = 0.03
    constants[42] = 0.65
    constants[43] = 50
    constants[44] = 33
    constants[45] = constants[13]/10.0000
    constants[46] = (constants[14]*(0.400000+0.400000*constants[14]))/(constants[13]*(power((constants[14]+1.00000)*0.400000, 2.00000)))
    constants[47] = (0.400000*constants[14]+1.00000)/(constants[14]*constants[13])
    return (states, constants)

def computeRates(voi, states, constants):
    rates = [0.0] * sizeStates; algebraic = [0.0] * sizeAlgebraic
    algebraic[4] = custom_piecewise([equal(constants[0] , 1.00000), 0.00000 , True, -states[2]])
    rates[5] = algebraic[4]
    algebraic[9] = states[0]
    rates[3] = algebraic[9]
    algebraic[11] = states[2]
    rates[4] = algebraic[11]
    algebraic[8] = (states[3]+constants[18])/(0.460000+constants[18])
    algebraic[10] = states[1]/(algebraic[8]*states[6])
    algebraic[12] = custom_piecewise([greater_equal(algebraic[10] , 1.00000), 1.00000 , True, power(0.0200000, algebraic[10])])
    algebraic[14] = constants[25]*(constants[36]-states[6])*states[8]-constants[26]*exp(-constants[27]*states[6])*algebraic[12]*states[6]
    rates[6] = algebraic[14]
    algebraic[16] = constants[38]*(constants[37]-states[7])*states[8]-constants[39]*states[7]
    rates[7] = algebraic[16]
    rates[8] = custom_piecewise([less(voi , constants[44]), constants[41]*voi*(1.00000-exp(-constants[40]*(power(voi, 2.00000))))*exp(-constants[40]*(power(voi, 2.00000))) , True, (-algebraic[14]-algebraic[16])-constants[42]*exp(-constants[43]*states[8])*states[8]])
    algebraic[6] = (power(states[6], constants[9]))/(power(states[6], constants[9])+power(constants[8], constants[9]))
    algebraic[7] = 0.600000*states[3]+0.500000
    algebraic[3] = custom_piecewise([less_equal(states[0] , 0.00000), constants[19]-(constants[20]*states[0])/constants[13] , True, constants[21]])
    algebraic[17] = (constants[14]*constants[15]*(power(constants[45]/constants[13], 2.00000)))/(3.00000*constants[14]*constants[15]-((constants[14]+1.00000)*states[0])/constants[13])
    algebraic[18] = custom_piecewise([less_equal(states[0] , 0.00000), (constants[14]*(1.00000+states[0]/constants[13]))/(constants[14]-states[0]/constants[13]) , True, (1.00000+constants[15])-((power(constants[15], 2.00000))*constants[14])/(((constants[14]*constants[15])/algebraic[17])*(power(states[0]/constants[13], 2.00000))+((constants[14]+1.00000)*states[0])/constants[13]+constants[14]*constants[15])])
    algebraic[19] = custom_piecewise([less_equal(-constants[13] , states[0]) & less_equal(states[0] , 0.00000), 1.00000+(0.600000*states[0])/constants[13] , less(0.00000 , states[0]) & less_equal(states[0] , constants[45]), algebraic[18]/((((0.400000*constants[14]+1.00000)/constants[14])*states[0])/constants[13]+1.00000) , True, (algebraic[18]*exp(-constants[24]*(power((states[0]-constants[45])/constants[13], constants[17]))))/((((0.400000*constants[14]+1.00000)/constants[14])*states[0])/constants[13]+1.00000)])
    algebraic[20] = constants[10]*constants[11]*algebraic[3]*constants[12]*algebraic[19]
    algebraic[21] = constants[11]*algebraic[3]*(1.00000-constants[10]*constants[12]*algebraic[19])
    algebraic[22] = algebraic[20]*algebraic[6]*algebraic[7]*algebraic[8]*(1.00000-states[1])-algebraic[21]*states[1]
    rates[1] = algebraic[22]
    algebraic[13] = custom_piecewise([less_equal(algebraic[9] , 0.00000), constants[29]*exp(constants[28]*states[3]) , True, constants[31]*exp(constants[30]*states[3])])
    algebraic[25] = algebraic[18]/algebraic[19]
    algebraic[23] = (constants[14]*1.00000*(1.00000+0.400000*constants[14]+(1.20000*states[0])/constants[13]+0.600000*(power(states[0]/constants[13], 2.00000))))/(constants[13]*(power((constants[14]-states[0]/constants[13])*(1.00000+(0.600000*states[0])/constants[13]), 2.00000)))
    algebraic[24] = (1.00000/constants[13])*exp(-constants[24]*(power(states[0]/constants[13]-constants[45]/constants[13], constants[17])))*((0.400000*constants[14]+1.00000)/constants[14]+constants[24]*constants[16]*(1.00000+((0.400000*constants[14]+1.00000)*states[0])/(constants[14]*constants[13])))*(power(states[0]/constants[13]-constants[45]/constants[13], constants[17]-1.00000))
    rootfind_0(voi, constants, rates, states, algebraic)
    rates[0] = algebraic[27]
    algebraic[15] = custom_piecewise([less_equal(algebraic[11] , algebraic[9]), constants[33]*exp(constants[32]*(states[4]-states[3])) , True, constants[35]*exp(constants[34]*(states[4]-states[3]))])
    rates[2] = custom_piecewise([equal(constants[0] , 1.00000), ((algebraic[15]*(algebraic[27]-constants[23]*(power(states[2]-states[0], 2.00000)))-constants[1]*constants[2]*exp(constants[1]*(states[4]-states[3]))*(states[2]-states[0])*1.00000)-constants[3]*constants[4]*exp(constants[3]*states[4])*states[2]*1.00000)/algebraic[15] , True, ((algebraic[15]*(algebraic[27]-constants[23]*(power(states[2]-states[0], 2.00000)))-constants[1]*constants[2]*exp(constants[1]*(states[4]-states[3]))*(states[2]-states[0])*1.00000)-(constants[3]*constants[4]*exp(constants[3]*states[4])+constants[5]*constants[6]*exp(constants[5]*states[5]))*states[2]*1.00000)/algebraic[15]])
    return(rates)

def computeAlgebraic(constants, states, voi):
    algebraic = array([[0.0] * len(voi)] * sizeAlgebraic)
    states = array(states)
    voi = array(voi)
    algebraic[4] = custom_piecewise([equal(constants[0] , 1.00000), 0.00000 , True, -states[2]])
    algebraic[9] = states[0]
    algebraic[11] = states[2]
    algebraic[8] = (states[3]+constants[18])/(0.460000+constants[18])
    algebraic[10] = states[1]/(algebraic[8]*states[6])
    algebraic[12] = custom_piecewise([greater_equal(algebraic[10] , 1.00000), 1.00000 , True, power(0.0200000, algebraic[10])])
    algebraic[14] = constants[25]*(constants[36]-states[6])*states[8]-constants[26]*exp(-constants[27]*states[6])*algebraic[12]*states[6]
    algebraic[16] = constants[38]*(constants[37]-states[7])*states[8]-constants[39]*states[7]
    algebraic[6] = (power(states[6], constants[9]))/(power(states[6], constants[9])+power(constants[8], constants[9]))
    algebraic[7] = 0.600000*states[3]+0.500000
    algebraic[3] = custom_piecewise([less_equal(states[0] , 0.00000), constants[19]-(constants[20]*states[0])/constants[13] , True, constants[21]])
    algebraic[17] = (constants[14]*constants[15]*(power(constants[45]/constants[13], 2.00000)))/(3.00000*constants[14]*constants[15]-((constants[14]+1.00000)*states[0])/constants[13])
    algebraic[18] = custom_piecewise([less_equal(states[0] , 0.00000), (constants[14]*(1.00000+states[0]/constants[13]))/(constants[14]-states[0]/constants[13]) , True, (1.00000+constants[15])-((power(constants[15], 2.00000))*constants[14])/(((constants[14]*constants[15])/algebraic[17])*(power(states[0]/constants[13], 2.00000))+((constants[14]+1.00000)*states[0])/constants[13]+constants[14]*constants[15])])
    algebraic[19] = custom_piecewise([less_equal(-constants[13] , states[0]) & less_equal(states[0] , 0.00000), 1.00000+(0.600000*states[0])/constants[13] , less(0.00000 , states[0]) & less_equal(states[0] , constants[45]), algebraic[18]/((((0.400000*constants[14]+1.00000)/constants[14])*states[0])/constants[13]+1.00000) , True, (algebraic[18]*exp(-constants[24]*(power((states[0]-constants[45])/constants[13], constants[17]))))/((((0.400000*constants[14]+1.00000)/constants[14])*states[0])/constants[13]+1.00000)])
    algebraic[20] = constants[10]*constants[11]*algebraic[3]*constants[12]*algebraic[19]
    algebraic[21] = constants[11]*algebraic[3]*(1.00000-constants[10]*constants[12]*algebraic[19])
    algebraic[22] = algebraic[20]*algebraic[6]*algebraic[7]*algebraic[8]*(1.00000-states[1])-algebraic[21]*states[1]
    algebraic[13] = custom_piecewise([less_equal(algebraic[9] , 0.00000), constants[29]*exp(constants[28]*states[3]) , True, constants[31]*exp(constants[30]*states[3])])
    algebraic[25] = algebraic[18]/algebraic[19]
    algebraic[23] = (constants[14]*1.00000*(1.00000+0.400000*constants[14]+(1.20000*states[0])/constants[13]+0.600000*(power(states[0]/constants[13], 2.00000))))/(constants[13]*(power((constants[14]-states[0]/constants[13])*(1.00000+(0.600000*states[0])/constants[13]), 2.00000)))
    algebraic[24] = (1.00000/constants[13])*exp(-constants[24]*(power(states[0]/constants[13]-constants[45]/constants[13], constants[17])))*((0.400000*constants[14]+1.00000)/constants[14]+constants[24]*constants[16]*(1.00000+((0.400000*constants[14]+1.00000)*states[0])/(constants[14]*constants[13])))*(power(states[0]/constants[13]-constants[45]/constants[13], constants[17]-1.00000))
    algebraic[15] = custom_piecewise([less_equal(algebraic[11] , algebraic[9]), constants[33]*exp(constants[32]*(states[4]-states[3])) , True, constants[35]*exp(constants[34]*(states[4]-states[3]))])
    algebraic[0] = constants[6]*(exp(constants[5]*states[5])-1.00000)
    algebraic[1] = constants[2]*(exp(constants[1]*(states[4]-states[3]))-1.00000)
    algebraic[2] = constants[4]*(exp(constants[3]*states[4])-1.00000)
    algebraic[5] = algebraic[0]
    algebraic[26] = constants[7]*algebraic[25]*states[1]
    return algebraic

initialGuess0 = None
def rootfind_0(voi, constants, rates, states, algebraic):
    """Calculate values of algebraic variables for DAE"""
    from scipy.optimize import fsolve
    global initialGuess0
    if initialGuess0 is None: initialGuess0 = ones(2)*0.1
    if not iterable(voi):
        soln = fsolve(residualSN_0, initialGuess0, args=(algebraic, voi, constants, rates, states), xtol=1E-6)
        initialGuess0 = soln
        algebraic[27] = soln[0]
        algebraic[28] = soln[1]
    else:
        for (i,t) in enumerate(voi):
            soln = fsolve(residualSN_0, initialGuess0, args=(algebraic[:,i], voi[i], constants, rates[:i], states[:,i]), xtol=1E-6)
            initialGuess0 = soln
            algebraic[27][i] = soln[0]
            algebraic[28][i] = soln[1]

def residualSN_0(algebraicCandidate, algebraic, voi, constants, rates, states):
    resid = array([0.0] * 2)
    algebraic[27] = algebraicCandidate[0]
    algebraic[28] = algebraicCandidate[1]
    resid[0] = (algebraic[27]-(custom_piecewise([equal(constants[0] , 1.00000), (constants[7]*algebraic[22]*algebraic[25]*1.00000+constants[22]*algebraic[13]*(power(states[0], 2.00000))+constants[3]*constants[4]*exp(constants[3]*states[4])*states[2]*1.00000)/(constants[7]*states[1]*algebraic[28]*1.00000+algebraic[13]) , True, (constants[7]*algebraic[22]*algebraic[25]*1.00000+constants[22]*algebraic[13]*(power(states[0], 2.00000))+(constants[3]*constants[4]*exp(constants[3]*states[4])+constants[5]*constants[6]*exp(constants[5]*states[5]))*states[2]*1.00000)/(constants[7]*states[1]*algebraic[28]*1.00000+algebraic[13])])))
    resid[1] = (algebraic[28]-(custom_piecewise([less_equal(states[0] , -constants[13]), constants[46]*algebraic[27] , less(-constants[13] , states[0]) & less_equal(states[0] , 0.00000), algebraic[23]*algebraic[27] , less(0.00000 , states[0]) & less_equal(states[0] , constants[45]), constants[47]*algebraic[27] , True, algebraic[24]*algebraic[27]])))
    return resid

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)