import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import solve_ivp
import pysindy as ps
import pandas as pd

# ------------------------------------------------------------
# Load the VGP dynamical dataset used to train the SINDy model
# ------------------------------------------------------------

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

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

# Normalize
omega = (omega - omega.mean()) / omega.std()

# Angular acceleration
dt = np.gradient(age)
alpha = np.gradient(omega) / dt

X = np.vstack([omega,alpha]).T

# ------------------------------------------------------------
# Rebuild the SINDy model (same settings as before)
# ------------------------------------------------------------

library = ps.PolynomialLibrary(degree=2, include_bias=False)

optimizer = ps.STLSQ(
    threshold=0.05,
    alpha=0.01
)

model = ps.SINDy(
    optimizer=optimizer,
    feature_library=library
)

model.fit(X, t=np.mean(np.diff(age)))

print("\nRecovered dynamical system:\n")
model.print()

# ------------------------------------------------------------
# Define RHS for integration
# ------------------------------------------------------------

def sindy_rhs(t, y):
    return model.predict(y.reshape(1,-1))[0]

# ------------------------------------------------------------
# Phase portrait simulation
# ------------------------------------------------------------

t_span = (0,50)
t_eval = np.linspace(0,50,2000)

sol = solve_ivp(
    sindy_rhs,
    t_span,
    X[0],
    t_eval=t_eval,
    method="RK45"
)

sim = sol.y.T

# ------------------------------------------------------------
# PHASE PORTRAIT
# ------------------------------------------------------------

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

plt.plot(sim[:,0], sim[:,1], color="darkorange")
plt.scatter(X[:,0], X[:,1], s=8, alpha=0.4)

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

plt.title("Geomagnetic Pole Phase Portrait")

plt.tight_layout()
plt.savefig("geomagnetic_phase_portrait.png", dpi=300)
plt.show()

# ------------------------------------------------------------
# POTENTIAL LANDSCAPE
# ------------------------------------------------------------

grid = 80

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

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

V = np.zeros_like(W)

# evaluate dynamical magnitude
for i in range(grid):
    for j in range(grid):

        state = np.array([W[i,j], A[i,j]])

        flow = model.predict(state.reshape(1,-1))[0]

        V[i,j] = np.linalg.norm(flow)

# ------------------------------------------------------------
# Plot potential landscape
# ------------------------------------------------------------

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

plt.contourf(W, A, V, levels=40, cmap="inferno")
plt.colorbar(label="Dynamical magnitude")

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

plt.title("Geomagnetic Dynamical Landscape")

plt.tight_layout()
plt.savefig("geomagnetic_potential_landscape.png", dpi=300)
plt.show()