# -*- coding: utf-8 -*-
from scipy.integrate import odeint
from pylab import figure, plot, show, xlabel, ylabel ,legend, subplot,title
import numpy as np
    
def HH(y,t,I):
    
    V = y[0]
    n = y[1]
    m = y[2]
    h = y[3]
    
    # stałe:
    C =1.
    g_Na = 120.
    V_Na = 50.
    g_K = 36.
    V_K = -77.
    g_L = 0.3
    V_L = -54.
    
    # prawa strona
    dV = 1./C * (I - g_Na*m**3*h*(V-V_Na) -g_K*n**4*(V-V_K) -g_L*(V-V_L)  )
    dn = a_n(V)*(1-n) -b_n(V)*n    
    dm = a_m(V)*(1-m) -b_m(V)*m 
    dh = a_h(V)*(1-h) -b_h(V)*h
  
    return np.array([dV, dn, dm, dh]) 
                
# funkcje pomocnicze do HH
def a_n(V):
    return 0.01*(V+55.)/(1-np.exp(-(V+55.)/10.))
def b_n(V):
    return 0.125*np.exp(-(V+65.)/80.)
def a_m(V):
    return 0.1*(V+40.)/(1-np.exp(-(V+40.)/10.))
def b_m(V):
    return 4*np.exp(-(V+65.)/18.)
def a_h(V):
    return 0.07*np.exp(-(V+65.)/20.)
def b_h(V):
    return 1./(1+np.exp(-(V+35.)/10))
    
def n_inf(V):
    return a_n(V)/(a_n(V)+b_n(V))
def tau_n(V):
    return 1./(a_n(V)+b_n(V))
def m_inf(V):
    return a_m(V)/(a_m(V)+b_m(V))
def tau_m(V):
    return 1./(a_m(V)+b_m(V))
def h_inf(V):
    return a_h(V)/(a_h(V)+b_h(V))
def tau_h(V):
    return 1./(a_h(V)+b_h(V))

def rysuj_stale_czasowe():    
    V = np.linspace(-100,100,201)
    subplot(3,2,1)
    plot(V,n_inf(V))
    title('n_inf')
    
    subplot(3,2,2)
    plot(V,tau_n(V))
    title('tau_n')
    
    subplot(3,2,3)
    plot(V,m_inf(V))
    title('m_inf')
    
    subplot(3,2,4)
    plot(V,tau_m(V))
    title('tau_m')
    
    subplot(3,2,5)
    plot(V,h_inf(V))
    title('h_inf')
    
    subplot(3,2,6)
    plot(V,tau_h(V))
    title('tau_h')
    
    
    show()

#rysuj_stale_czasowe()
## całkowanie
time = np.linspace(0.0,200.0,1000) 
# warunki początkowe
yinit = np.array([0.,0.,0.,0.]) 

# prametry
I = 0
# przebieg w czasie

y1 = odeint(HH, yinit, time, (I,),atol=1e-13,rtol=1e-13)
yinit=y1[-1,:]
time2 = time[-1]+np.linspace(0.0,200.0,1000)
I =  -10 #6.125  #2.222
y2 = odeint(HH, yinit, time2, (I,),atol=1e-13,rtol=1e-13)
yinit=y2[-1,:]
time3 = time2[-1]+np.linspace(0.0,200.0,1000)
I =  0 #6.125  #2.222
y3 = odeint(HH, yinit, time3, (I,),atol=1e-13,rtol=1e-13)

figure(1) 
plot(np.concatenate((time,time2,time3)),np.concatenate((y1[:,0],y2[:,0],y3[:,0]))) # y[:,0] pierwsza kolumna y - zawiera zmienną u 
#plot(time,y[:,1]) # y[:,1] druga kolumna y - zawiera zmienną v 
xlabel('t') 
ylabel('y') 
legend(('v',))
show()