"""A demonstration of a model for a single Hodgkin-Huxley membrane
potential for an oscillatory cortical neuron. 

This is the stripped-down version from Robert Clewley (June 2005). The full version,
with many options and switches, comes with PyDSTool

Thomas Haslwanter
Ver 1.2
26-Feb-2017
"""

#from PyDSTool import *
import PyDSTool as pdt
import matplotlib.pyplot as plt

# ------------------------------------------------------------

def makeHHneuron(name, par_args, ic_args, extra_terms=''):
    '''Define the Hodgkins Huxley equations'''
    vfn_str = '(Iapp'+extra_terms+'-ionic(v,m,h,n))/C'
    mfn_str = 'ma(v)*(1-m)-mb(v)*m'
    nfn_str = 'na(v)*(1-n)-nb(v)*n'
    hfn_str = 'ha(v)*(1-h)-hb(v)*h'
    aux_str = 'm*m*m*h'

    auxdict = {'ionic': (['vv', 'mm', 'hh', 'nn'],
            'gna*mm*mm*mm*hh*(vv-vna) + gk*nn*nn*nn*nn*(vv-vk) + gl*(vv-vl)'),
               'ma': (['v'], '0.32*(v+54)/(1-exp(-(v+54)/4))'),
               'mb': (['v'], '0.28*(v+27)/(exp((v+27)/5)-1)'),
               'ha': (['v'], '.128*exp(-(50+v)/18)'),
               'hb': (['v'], '4/(1+exp(-(v+27)/5))'),
               'na': (['v'], '.032*(v+52)/(1-exp(-(v+52)/5))'),
               'nb': (['v'], '.5*exp(-(57+v)/40)'),
               'ptest': (['p'], '1+p+ma(-50)+C'),
               'atest': (['q'], 'q+mb(-40)')}  # tests fn cross-calling

    DSargs = pdt.args()
    DSargs.varspecs = {'v': vfn_str, 'm': mfn_str,
                       'h': hfn_str, 'n': nfn_str,
                       'v_bd0': 'getbound("v",0)',  # demo of referencing bounds
                       'v_bd1': 'getbound("v",1)'}
    DSargs.pars = par_args
    DSargs.auxvars = ['v_bd0','v_bd1']
    DSargs.fnspecs = auxdict
    DSargs.xdomain = {'v': [-130, 70], 'm': [0,1], 'h': [0,1], 'n': [0,1]}
    DSargs.algparams = {'init_step':0.03}
    DSargs.checklevel = 0
    DSargs.ics = ic_args
    DSargs.name = name
    return pdt.Generator.Vode_ODEsystem(DSargs)

# ------------------------------------------------------------

if __name__=='__main__':
    # need the __main__ to use above functions as imports without running this part
    print('-------- Test: Hodgkin-Huxley system')
    
    # Define the parameters
    par_args = {'gna': 100, 'gk': 80, 'gl': 0.1,
                'vna': 50, 'vk': -100, 'vl': -67,
                'Iapp': 1.75, 'C': 1.0}
    ic_args = {'v':-70.0, 'm': 0, 'h': 1, 'n': 0}

    # Generate the system
    HH = makeHHneuron('HHtest', par_args, ic_args)
    
    # Calculate the trajectory
    HH.set(tdata=[0, 40])
    HHtraj = HH.compute('test_long')
    plotData = HHtraj.sample()
    
    # Plot the data
    plt.ylabel('v')
    plt.xlabel('t')    
    plt.plot(plotData['t'], plotData['v'])
    plt.show()