from tqdm import tqdm
import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import solve_ivp
import os

for i, mu in enumerate(tqdm(np.linspace(-1, 1, 240))):
  def system(t, y):
      v, w = y
      dv = mu * v + w - v**2
      dw = -v + mu * w + 2 * v**2
      return [dv, dw]
  def system_reversed(t, y):
      v, w = y
      dv = mu * v + w - v**2
      dw = -v + mu * w + 2 * v**2
      return [-dv, -dw]

  x_root = (mu**2+1)/(2+mu)
  y_root = -mu * x_root + x_root ** 2
  vmin, vmax, wmin, wmax = -1,1,-1,1
  # vmin,vmax,wmin,wmax= x_root-0.0005,x_root+0.0005, y_root-0.0005, y_root+0.0005

  t_span = [0, 10]
  trajectory_resolution = 10

  epsilon = 0.01
  
  initial_conditions = [(x, y)  for x in np.linspace(vmin, vmax, trajectory_resolution) for y in np.linspace(wmin, wmax, trajectory_resolution)]
  initial_conditions_2 = [(x_root + dx, y_root + dy) for dx in np.linspace(-epsilon, epsilon, 10) for dy in np.linspace(-epsilon, epsilon, 10)]
  sols = {}
  sols_2 = {}
  sols_reversed = {}
  sols_reversed_2 = {}
  for ic in initial_conditions:
      sols[ic] = solve_ivp(system, t_span, ic, dense_output=True, max_step=0.05)
      sols_reversed[ic] = solve_ivp(system_reversed, t_span, ic, dense_output=True, max_step=0.05)
  for ic in initial_conditions_2:
      sols_2[ic] = solve_ivp(system, t_span, ic, dense_output=True, max_step=0.05)
      sols_reversed_2[ic] = solve_ivp(system_reversed, t_span, ic, dense_output=True, max_step=0.05)

  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, dw = system(0, [v_values, w_values])

  fig, ax = plt.subplots(figsize=(16,16))
  # ax.scatter(x_root, y_root)
  # integral curves
  for ic in initial_conditions:
    sol = sols[ic]
    ax.plot(sol.y[0], sol.y[1],alpha=0.4, linewidth=0.5, color='k')
    sol = sols_reversed[ic]
    ax.plot(sol.y[0], sol.y[1], alpha=0.4, linewidth=0.5, color='k')
  for ic in initial_conditions_2:
    sol = sols_2[ic]
    ax.plot(sol.y[0], sol.y[1],alpha=0.8, linewidth=0.5, color='r')
    sol = sols_reversed_2[ic]
    ax.plot(sol.y[0], sol.y[1], alpha=0.8, linewidth=0.5, color='b')

  # 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**2-mu*vs,  color="green", alpha=0.2, label="x nullcline")
  if np.abs(mu) < 0.001:
    ax.axvline(0, wmin, wmax, color="red", alpha=0.2, label="y nullcline")
    ax.axvline(1/2, wmin, wmax, color="red", alpha=0.2, label="y nullcline")
  else:  
    ax.plot(vs, (vs-2*vs**2)/mu, color="red", alpha=0.2, label="y nullcline")

  ax.set_title(f'Hopf Bifurcation Model
$\mu={mu:.3f}}})

  # ax.legend()
  ax.set_xlim(vmin, vmax)
  ax.set_ylim(wmin, wmax)
  ax.set_xticks([])
  ax.set_yticks([])
  dir_path = f"./hopf"
  if not os.path.exists(dir_path):
    os.makedirs(dir_path)

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

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

for dir_path in ["./hopf"]:
    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 = 24

    # 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 = dir_path[2:]
    iio.imwrite(f"{filename}.gif", images, duration=1000/fps, rewind=True)
    clip = mp.ImageSequenceClip(images, fps=fps)
    clip.write_videofile(f"{filename}.mp4")