from pylab import *
from numpy import *

def Step(tstart=0, tend=300, Step_tstart = 20, Step_tend = 270, I_amp=.5, dt=0.01):
    """
    DEFINITION
    Runs the TypeX model for a step current.
    
    INPUT
    tstart: starting time (ms)
    tend: ending time (ms)
    Step_tstart: time at which the step current begins (ms)
    Step_tend: time  at which the step current ends (ms)
    I_amp: magnitude of the current step (uA/cm^2)
    dt: integration timestep
    
    OUTPUT

    [t,v,w,I]
    
    t - arange(tstart,tend,dt)
    v - voltage trace
    w - slow recovery variable
    I - current injected
    
    ALGORITHM
    uses Forward-Euler numerical integration in ForwardEuler
    
    -R.Naud 02.2009.
    """

    # make time bins
    t = arange(tstart,tend,dt)  
    # make current array for each time bin
    I = zeros(t.shape)

    # find array index of I start and stop times
    index_start = searchsorted(t,Step_tstart)
    index_end = searchsorted(t,Step_tend)
    # assign amplitude to that range
    I[index_start:index_end] = I_amp

    # run the integrator
    [v, w] = ForwardEuler(t,I)

    return [t,v,w,I]

def PlotStep(tstart=0, tend=300, Step_tstart = 20, Step_tend = 270, I_amp=.5, dt=0.01):
    """
    
    DEFINITION
    Plots the TypeX model for a step current.
    
    INPUT
    tstart: starting time (ms)
    tend: ending time (ms)
    Step_tstart: time at which the step current begins (ms)
    Step_tend: time  at which the step current ends (ms)
    I_amp: magnitude of the current step (uA/cm^2)
    dt: integration timestep
    
    OUTPUT
    graph with three panels: 1) voltage trace, 2) slow recovery variable w 3) current injected.

    """

    t,v,w,I = Step(tstart, tend, Step_tstart, Step_tend, I_amp, dt)

    # open new figure and plot
    figure()
    # plot voltage time series
    subplot(311)
    plot(t,v,lw=2)
    xlabel('t')
    ylabel('v')
    # plot activation and inactivation variables
    subplot(312)
    plot(t,w,'k',lw=2)
    xlabel('t')
    ylabel('w')
    # plot current
    subplot(313)
    plot(t,I,lw=2)
    axis((tstart, tend, 0, I_amp*1.1))
    xlabel('t (ms)')
    ylabel('I (pA)')

def ForwardEuler(t,I):
    """
    solve the original TypeX model
    using Euler's method

    Inputs:
       t   - equally spaced time bins (ms)
             e.g. arange(0,100,0.1)
             
       I   - current input to LIF,
             same shape as t 
    OUPTUT
    [v,w]
    the timeseries of the voltage (v, in mV), the slow recovery variable (w)        
    
    ALGORITHM
    uses the equations and parameters given in 'Spiking Neuron Models' p70-71.
    These parameters are based on the voltage scale where the resting potential is zero.
    """
    
    # simple error check
    if I.shape!=t.shape:
        raise TypeError, 'require I.shape == t.shape'
    
    # setting Parameter Values
    # in mV, uF/cm2, mS/cm2
    Params = {'a': 1.25, 
          'tau': 15.6}
    
    
    L = len(t)
    dt = t[1]-t[0]

    # allocate array for voltage and gating variables
    v = zeros(I.shape) # this sets initial condition of v to be zero as required
    w = zeros(I.shape)
    
    
    # let initial value be m_inf(initial voltage), etc.
    w[0] = -.5
    v[0] = -1.3
    
    
    for i in range(1,len(t)):
        # forward euler step
        v[i] = v[i-1] + dudt(v[i-1],w[i-1], I[i-1])*dt
        w[i] = w[i-1] + dwdt(v[i-1], w[i-1], Params)*dt
    
    return [v, w]

# define derivatives; factor 1e3 needed for t in ms

def dudt(v,w,I):
    """
    This is the first equation of the TypeX system, the voltage ordinary differential equation

    INPUTS
    
    (state variables)
    u: voltage (float)
    w: recovery variable (float)
    
    I: current (float)
    """
    dudt = v - (v**3)/3 - w + I
    return dudt

def dwdt(v,w,Params):
    """
    This is the second equation of the TypeX system,
    the recovery variable ordinary differential equation

    INPUTS
    (state variables)
    u: voltage (float)
    w: recovery variable (float)
    
    I: current (float)
    
    Params is a dictionary with two keys: 'a' the slope of the w-nullcline 
    and 'tau' the time constant of the recovery variable
    """
    dwdt = (Params['a']*(v+.7) - w)/Params['tau']
    return dwdt