from matplotlib.widgets import AxesWidget
import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import solve_ivp
from scipy import optimize

for eps in [0.2]:
  # Define the parameter values
  a = 0.7
  b = 2.0 # If b < 1.5, then there are stable loops. Else there are no loops.
  tau = 12.5
  R = 0.1
  I_ext = ((a-1)/b + 2/3)/R + eps

  # Define the system of ODEs
  def system(t, y):
      v, w = y
      dv = v - (v ** 3) / 3 - w + R * I_ext
      dw = (1 / tau) * (v + a - b * w)
      return [dv, dw]
  def system_reversed(t, y):
      v, w = y
      dv = v - (v ** 3) / 3 - w + R * I_ext
      dw = (1 / tau) * (v + a - b * w)
      return [-dv, -dw]

  vmin, vmax, wmin, wmax = -2, 2, -2+R*I_ext, 2+R*I_ext

  t_span = [0, 100]
  trajectory_resolution = 10
  def fun(x):
    v = x[0]
    return v-v**3/3 + R * I_ext - (v+a)/b
  sol = optimize.root(fun, [0], method='hybr')
  x_root = sol.x[0]
  y_root = (x_root+a)/b

  # vmin, vmax, wmin, wmax = -1.5, -0.5, -1.1 +1/3 + R * I_ext, -0.8 +1/3 + R * I_ext
  # initial_conditions = [(-1.0, y) for y in np.linspace(-0.16, -0.03, 30)]
  initial_conditions = [(x, y) for x in np.linspace(vmin, vmax, trajectory_resolution) for y in np.linspace(wmin, wmax, trajectory_resolution)]
  epsilon = 0.005
  initial_conditions += [(x, y) for x in np.linspace(x_root - epsilon, x_root+epsilon, trajectory_resolution)  for y in np.linspace(y_root - epsilon, y_root+epsilon, trajectory_resolution)]
  sols = {}
  for ic in initial_conditions:
      sols[ic] = solve_ivp(system, t_span, ic, dense_output=True, max_step=0.1)
  sols_reversed = {}
  for ic in initial_conditions:
      sols_reversed[ic] = solve_ivp(system_reversed, t_span, ic, dense_output=True, max_step=0.1)

  vs = np.linspace(vmin, vmax, 200)
  v_axis = np.linspace(vmin, vmax, 20)
  w_axis = np.linspace(wmin, wmax, 20)

  v_values, w_values = np.meshgrid(v_axis, w_axis)

  dv = v_values - (v_values ** 3) / 3 - w_values + R * I_ext
  dw = (1 / tau) * (v_values + a - b * w_values)

  fig, ax = plt.subplots(figsize=(16,16))
  # integral curves
  for ic in initial_conditions:
    sol = sols[ic]
    ax.plot(sol.y[0], sol.y[1], color='k', alpha=0.4, linewidth=0.5)
    sol = sols_reversed[ic]
    ax.plot(sol.y[0], sol.y[1], color='k', alpha=0.4, linewidth=0.5)

  # vector fields
  arrow_lengths = np.sqrt(dv**2 + dw**2)
  alpha_values = 1 - (arrow_lengths / np.max(arrow_lengths))**0.4
  ax.quiver(v_values, w_values, dv, dw, color='blue', linewidth=0.5, scale=25, alpha=alpha_values)

  # nullclines
  ax.plot(vs, vs - vs**3/3 + R * I_ext,  color="green", alpha=0.4, label="v nullcline")
  ax.plot(vs, (vs + a) / b, color="red", alpha=0.4, label="w nullcline")

  # ax.set_xlabel('v')
  # ax.set_ylabel('w')
  ax.set_title(f'FitzHugh-Nagumo Model\n$b={b:.2f}$\t\t$I_{{ext\}\} = {I_ext:.2f\}\}\})

  # ax.legend()
  ax.set_xlim(vmin, vmax)
  ax.set_ylim(wmin, wmax)
  # ax.set_xticks([])
  # ax.set_yticks([])
  plt.show()
}}