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 TypeY 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 TypeY 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 ForwardEuler3(t,I,w0 = 0,v0=-70,Params={'a':0,'EL':-70.0,'tauw':10,'b':0,'Ereset':-73,'gL':4.0,'vT':-50,'DT':2,'C':40,'a2':0,'a3':0,'b2':0,'b3':0, 'tauw2':10,'tauw3':10}):
    """
    solve the original TypeY 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), and 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'
    
    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)
    w2 = zeros(I.shape)
    w3 = zeros(I.shape)
    I = array(I).astype(float)
    
    # let initial value be m_inf(initial voltage), etc.
    v[0] = v0
    w[0] = w0
    w2[0] = w0
    w3[0] = w0
    a=Params.get('a')
    a2 = Params.get('a2')
    a3 = Params.get('a3')
    b2 = Params.get('b2')
    b3 = Params.get('b3')
    EL= Params.get('EL')
    tauw=Params.get('tauw')
    tauw2 = Params.get('tauw2')
    tauw3 = Params.get('tauw3')
    b=Params.get('b')
    Ereset = Params.get('Ereset')
    
    i=0
    spks = array([])
    while i <len(I)-1:
        # forward euler step
        i=i+1
        v[i] = v[i-1] + dudt3(v[i-1],w[i-1],w2[i-1],w3[i-1], I[i-1],Params)*dt
        w[i] = w[i-1] + ((a*(v[i-1]-EL)-w[i-1])/tauw)*dt
        w2[i] = w2[i-1] + ((a2*(v[i-1]-EL)-w2[i-1])/tauw2)*dt
        w3[i] = w3[i-1] + ((a3*(v[i-1]-EL)-w3[i-1])/tauw3)*dt
        if v[i]>0:
            v[i:i+4]=0
            w[i:i+4]=w[i]
            w2[i:i+4] = w2[i]
            w3[i:i+4] = w3[i]
            v[i+5]=Ereset
            w[i+5]=w[i]+b
            w2[i+5] = w2[i] + b2
            w3[i+5] = w3[i] + b3
            spks = append(spks,i)
            i=i+5
    
    return [v, w,w2,w3,spks]



def ForwardEuler(t,I,w0 = 0,v0=-70,Params={'a':0,'EL':-70.0,'tauw':10,'b':0,'Ereset':-73,'gL':4.0,'vT':-50,'DT':2,'C':40}):
    """
    solve the original TypeY 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,spks]
    the timeseries of the voltage (v, in mV), and the slow recovery variable (w) and spike times (spks)
    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'
    
    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)
    I = array(I).astype(float)
    
    # let initial value be m_inf(initial voltage), etc.
    v[0] = v0
    w[0] = w0
    a=Params.get('a')
    EL= Params.get('EL')
    tauw=Params.get('tauw')
    b=Params.get('b')
    Ereset = Params.get('Ereset')
    
    i=0
    spks = array([])
    while i <len(I)-1:
        # forward euler step
        i=i+1
        v[i] = v[i-1] + dudt(v[i-1],w[i-1], I[i-1],Params)*dt
        w[i] = w[i-1] + ((a*(v[i-1]-EL)-w[i-1])/tauw)*dt
        if v[i]>0:
            v[i:i+4]=0
            w[i:i+4]=w[i]
            v[i+5]=Ereset
            w[i+5]=w[i]+b
            spks = append(spks,i)
            i=i+5
    
    return [v, w,spks]



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

def dudt(v,w,I,Params={'gL':4.0,'EL':-70.0,'vT':-50,'DT':2,'C':40}):
    """
    This is the first equation of the TypeY system, the voltage ordinary differential equation

    INPUTS
    (state variables)
    u: voltage (float)
    w: recovery variable (float)
    
    I: current (float)
    Params: dictionary with gL, EL, EL vT DT C

    """
    gL=Params.get('gL')
    EL=Params.get('EL')
    vT=Params.get('vT')
    DT=Params.get('DT')
    C=Params.get('C')
    
    dudt = -gL*(v-EL)/C +gL*exp((v-vT)/DT)/C+ I/C -w/C
    return dudt

def dudt3(v,w,w2,w3,I,Params={'gL':4.0,'EL':-70.0,'vT':-50,'DT':2,'C':40}):
    """
    This is the first equation of the TypeY system, the voltage ordinary differential equation

    INPUTS
    (state variables)
    u: voltage (float)
    w: recovery variable (float)
    
    I: current (float)
    Params: dictionary with gL, EL, EL vT DT C

    """
    gL=Params.get('gL')
    EL=Params.get('EL')
    vT=Params.get('vT')
    DT=Params.get('DT')
    C=Params.get('C')
    
    dudt = -gL*(v-EL)/C +gL*exp((v-vT)/DT)/C+ I/C -w/C - w2/C - w3/C
    return dudt