%%capture
from matplotlib.widgets import AxesWidget
import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import solve_ivp
import os

for b in [0.8, 1.25, 2.0]:
    for i, alpha in enumerate(np.sin(np.linspace(-np.pi/2, np.pi/2, 121))):
        # Define the parameter values
        a = 0.7
        tau = 12.5
        R = 0.1
        I_ext_0 = (2/3 + (a-1)/b)/R
        I_ext_1 = (-2/3 + (a+1)/b)/R
        I_ext_min = min(I_ext_0, I_ext_1) - 2.0 / R
        I_ext_max = max(I_ext_0, I_ext_1) + 2.0 / R
        I_ext = I_ext_min * (alpha - 1.0)/(-2.0) + I_ext_max * (alpha + 1.0)/(+2.0)

        # 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]

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

        t_span = [0, 50]
        trajectory_resolution = 30
        initial_conditions = [(x, y) for x in np.linspace(vmin, vmax, trajectory_resolution) for y in np.linspace(wmin, wmax, trajectory_resolution)]
        sols = {}
        for ic in initial_conditions:
            sols[ic] = solve_ivp(system, 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)

        # 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([])

        dir_path = f"./{b}"
        if not os.path.exists(dir_path):
            os.makedirs(dir_path)

        fig.savefig(f"{dir_path}/{i}.png")

import imageio.v3 as iio
import os
from natsort import natsorted
import moviepy.editor as mp

for b in [0.8, 1.25, 2.0]:
    dir_path = f"./{b}"
    file_names = natsorted((fn for fn in os.listdir(dir_path) if fn.endswith('.png')))

    # Create a list of image files and set the frame rate
    images = []
    fps = 12

    # Iterate over the file names and append the images to the list
    for file_name in file_names:
        file_path = os.path.join(dir_path, file_name)
        images.append(iio.imread(file_path))

    filename = f"fitzhugh-nagumo_b={b}"
    iio.imwrite(f"{filename}.gif", images + list(reversed(images)), duration=1000/fps, rewind=True)
    clip = mp.ImageSequenceClip(images + list(reversed(images)), fps=fps)
    clip.write_videofile(f"{filename}.mp4")