import pandas as pd
import numpy as np
import pysindy as ps
import matplotlib.pyplot as plt

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

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

# -----------------------------------
# time step
# -----------------------------------

dt = np.mean(np.diff(age))

# -----------------------------------
# Build 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
)

# -----------------------------------
# Fit model
# -----------------------------------

model.fit(X, t=dt)

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

# -----------------------------------
# Simulate
# -----------------------------------

from scipy.integrate import solve_ivp

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

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

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

sim = sol.y.T

# -----------------------------------
# Plot comparison
# -----------------------------------

plt.figure()

plt.plot(age, omega, label="data")
plt.plot(t_eval, sim[:,0], label="model")

plt.legend()
plt.title("Recovered Geodynamo Dynamics")

plt.tight_layout()
plt.show()