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

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

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

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

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

dt = np.gradient(age)

alpha = np.gradient(omega) / dt

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

# ------------------------------------------------------
# Rebuild SINDy model
# ------------------------------------------------------

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()

# ------------------------------------------------------
# Extract conservative acceleration component
# ------------------------------------------------------

omega_grid = np.linspace(-3,3,400)

alpha_zero = np.zeros_like(omega_grid)

states = np.vstack([omega_grid, alpha_zero]).T

acc = model.predict(states)[:,1]

# ------------------------------------------------------
# Compute potential V(ω)
# ------------------------------------------------------

V = -cumulative_trapezoid(acc, omega_grid, initial=0)

# normalize
V -= V.min()

# ------------------------------------------------------
# Plot potential curve
# ------------------------------------------------------

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

plt.plot(omega_grid, V, linewidth=3)

plt.xlabel("Angular velocity ω")
plt.ylabel("Potential V(ω)")

plt.title("Reconstructed Geomagnetic Stability Potential")

plt.grid(alpha=0.3)

plt.tight_layout()

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

plt.show()