import numpy as np
from scipy.integrate import solve_ivp
import matplotlib.pyplot as plt
import os
np.set_printoptions(threshold=np.inf)


def alpha_m(V_m): 
    return 0.1 * (- V_m + 25) / (np.exp((-V_m + 25) / 10) - 1)
    
def alpha_h(V_m): 
    return 0.07 * np.exp(-V_m / 20)

def alpha_n(V_m): 
    return 0.01 * (- V_m + 10) / (np.exp((- V_m + 10) / 10) - 1)
    
    
def beta_m(V_m): 
    return 4.0 * np.exp(-V_m / 18)
    
def beta_h(V_m): 
    return 1 / (np.exp((- V_m + 30) / 10) + 1)
    
def beta_n(V_m): 
    return 0.125 * np.exp(-V_m / 80)

V_m_range = np.linspace(-100, 100, 200)
U_rest = -65

m_inf = [alpha_m(V) / (alpha_m(V) + beta_m(V)) for V in V_m_range]
h_inf = [alpha_h(V) / (alpha_h(V) + beta_h(V)) for V in V_m_range]
n_inf = [alpha_n(V) / (alpha_n(V) + beta_n(V)) for V in V_m_range]

index_zero = np.argmin(np.abs(V_m_range - U_rest))
print(f"At V_m = {U_rest:.2f} mV, m_inf = {m_inf[index_zero]:.4f}, h_inf = {h_inf[index_zero]:.4f}, n_inf = {n_inf[index_zero]:.4f}")

plt.figure(figsize=(5.5, 5))
plt.plot(V_m_range, m_inf, label='$m_\infty(V_m)$', c='r')
plt.plot(V_m_range, h_inf, label='$h_\infty(V_m)$', c='g')
plt.plot(V_m_range, n_inf, label='$n_\infty(V_m)$', c='b')
plt.axvline(x=U_rest, color='gray', linestyle='--', label=f'$V_m$={U_rest:.2f} mV')
plt.plot(U_rest, m_inf[index_zero], 'ro')
plt.text(U_rest, m_inf[index_zero], f' {m_inf[index_zero]:.2f}', verticalalignment='bottom',
         color='red')
plt.plot(U_rest, h_inf[index_zero], 'go')
plt.text(U_rest, h_inf[index_zero], f' {h_inf[index_zero]:.2f}', verticalalignment='bottom',
         color='green')
plt.plot(U_rest, n_inf[index_zero], 'bo')
plt.text(U_rest, n_inf[index_zero], f'{n_inf[index_zero]:.2f} ', verticalalignment='bottom', 
         horizontalalignment='right', color='blue')
plt.title('Finding steady-state values of m, h, and n')
plt.xlabel('Membrane potential $V_m$ (mV)')
plt.ylabel('Steady-state value')
plt.legend()
plt.grid(True)
plt.show()

m0 = m_inf[index_zero]
h0 = h_inf[index_zero]
n0 = n_inf[index_zero]
y0 = [U_rest, m0, h0, n0]

# Параметры для моделирования:
I_amp = 25.0           # амплитуда внешнего тока
intervals = [(5, 100)]  # участки времени (в мс), в которые ток включён
t0, t_max = 0.0, 100.0  # время моделирования, мс
dt = 0.01              # шаг интегрирования, мс
num_steps = int((t_max - t0) / dt)

# Подготовим массивы для хранения решения:
t_array = np.arange(t0, t_max, dt)
U_array = np.zeros(num_steps)
m_array = np.zeros(num_steps)
h_array = np.zeros(num_steps)
n_array = np.zeros(num_steps)
I_array = np.zeros(num_steps)

# Найдём установившиеся значения m, h, n при U_rest:
m_inf_0 = alpha_m(U_rest) / (alpha_m(U_rest) + beta_m(U_rest))
h_inf_0 = alpha_h(U_rest) / (alpha_h(U_rest) + beta_h(U_rest))
n_inf_0 = alpha_n(U_rest) / (alpha_n(U_rest) + beta_n(U_rest))

# Начальные условия:
U_array[0] = U_rest
m_array[0] = m_inf_0
h_array[0] = h_inf_0
n_array[0] = n_inf_0


# constants:
C_m = 1.0  # membrane capacitance, in uF/cm^2
g_Na = 120.0  # maximum conductances, in mS/cm^2
g_K = 36.0
g_L = 0.3
E_Na = 120  # reversal potentials, in mV
E_K = -77.0 
E_L = -54.387

def I_ext(t, I_amp=1.0, intervals=[[5, 6], [10, 17]]):     
    """Return I_amp if t is within any of the specified intervals, else return 0."""
    for (start, end) in intervals:         
        if start <= t <= end:             
            return I_amp     
    return 0 
 

for i in range(num_steps - 1):
    # Текущее состояние:
    t_current = t_array[i]
    V_m = U_array[i]
    m   = m_array[i]
    h   = h_array[i]
    n   = n_array[i]
    
    # Вычисляем внешний ток:
    I_current = I_ext(t_current, I_amp, intervals)
    I_array[i] = I_current

    # Считаем правые части (производные):
    dVmdt = (I_current - g_Na * m**3 * h * (V_m - E_Na) - g_K  * n**4 * (V_m - E_K) - g_L  * (V_m - E_L)) / C_m
    
    dmdt = alpha_m(V_m)*(1 - m) - beta_m(V_m)*m
    dhdt = alpha_h(V_m)*(1 - h) - beta_h(V_m)*h
    dndt = alpha_n(V_m)*(1 - n) - beta_n(V_m)*n
    
    # Обновляем значения методом Эйлера:
    U_array[i+1] = V_m + dt*dVmdt
    m_array[i+1] = m   + dt*dmdt
    h_array[i+1] = h   + dt*dhdt
    n_array[i+1] = n   + dt*dndt


# Построение графиков:
plt.figure(figsize=(10, 30))

# 1. Мембранный потенциал:
plt.subplot(6, 1, 1)
plt.plot(t_array, U_array, 'k')
plt.axhline(U_rest, color='gray', linestyle='--', label='$U_{rest}$')
plt.ylabel('$V_m$ (мВ)')
plt.title('Метод Эйлера: решение уравнений Ходжкина–Хаксли')
plt.legend(loc='upper right')
plt.ylim(-100, 120)

# 2. Воротные переменные m, h, n:
plt.subplot(6, 1, 2)
plt.plot(t_array, m_array, 'r', label='$m$')
plt.plot(t_array, h_array, 'g', label='$h$')
plt.plot(t_array, n_array, 'b', label='$n$')
plt.legend(loc='upper right')
plt.ylabel('Воротные переменные')

# 3. Внешний ток:
plt.subplot(6, 1, 3)
plt.plot(t_array, I_array, 'm', label='$I_{ext}(t)$')
plt.ylabel('$I_{ext}$ (мкА/см$^2$)')
plt.legend(loc='upper right')

# 4. Фазовый портрет (пример: проекция на (V_m, m)):
plt.subplot(6, 1, 4)
plt.plot(U_array, m_array, 'r')
plt.xlabel('$V_m$ (мВ)')
plt.ylabel('$m$')
plt.xlim(-125,125)
plt.title('Фазовая плоскость $(V_m, m)$')

plt.subplot(6, 1, 5)
plt.plot(U_array, h_array, 'g')
plt.xlabel('$V_m$ (мВ)')
plt.ylabel('$h$')
plt.xlim(-125,125)
plt.title('Фазовая плоскость $(V_m, h)$')

plt.subplot(6, 1, 6)
plt.plot(U_array, n_array, 'b')
plt.xlabel('$V_m$ (мВ)')
plt.ylabel('$n$')
plt.xlim(-125,125)
plt.title('Фазовая плоскость $(V_m, n)$')

plt.tight_layout()
plt.show()

At V_m = -65.00 mV, m_inf = 0.0000, h_inf = 1.0000, n_inf = 0.0015

Численное моделирование нейрона Ходжкина-Хаксли