import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.ndimage import gaussian_filter
from skimage.graph import route_through_array

# ------------------------------------------------
# Load dataset
# ------------------------------------------------

df = pd.read_csv("vgp_spherical_dynamics_stable.csv")

age = df["Age_ka"].values
omega = df["AngularVelocity_deg_kyr"].values

omega = (omega - np.mean(omega)) / np.std(omega)

dt = np.gradient(age)
alpha = np.gradient(omega) / dt

# ------------------------------------------------
# Potential
# ------------------------------------------------

def potential(w):
    return 2.0495*w**2 - 0.3523*w**3

# ------------------------------------------------
# Phase space grid
# ------------------------------------------------

N = 300

w = np.linspace(-3,3,N)
a = np.linspace(-3,3,N)

W,A = np.meshgrid(w,a)

V = potential(W)
V -= np.min(V)

E = 0.5*A**2 + V

# smooth surface slightly
E = gaussian_filter(E,1)

# ------------------------------------------------
# Locate basin minimum
# ------------------------------------------------

basin = np.unravel_index(np.argmin(E),E.shape)

# ------------------------------------------------
# Locate saddle point
# (maximum along potential ridge)
# ------------------------------------------------

ridge = np.argmax(E[:,N//2])
saddle = (ridge,N//2)

print("Basin index:",basin)
print("Saddle index:",saddle)

# ------------------------------------------------
# Minimum energy path
# ------------------------------------------------

path,cost = route_through_array(
    E,
    basin,
    saddle,
    fully_connected=True
)

path = np.array(path)

path_w = w[path[:,1]]
path_a = a[path[:,0]]

# ------------------------------------------------
# Plot
# ------------------------------------------------

plt.figure(figsize=(7,6))

levels = np.linspace(0,np.percentile(E,95),50)

plt.contourf(W,A,E,levels=levels,cmap="inferno")

plt.scatter(omega,alpha,s=10,c="white",alpha=0.4,label="Observed pole motion")

plt.plot(path_w,path_a,color="cyan",linewidth=3,label="Minimum escape path")

plt.xlabel("Angular velocity ω")
plt.ylabel("Angular acceleration α")

plt.title("Geomagnetic Minimum-Energy Escape Path")

plt.legend()

plt.grid(alpha=0.2)

plt.tight_layout()

plt.savefig("geomagnetic_escape_path.png",dpi=300)

plt.show()