import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.neighbors import NearestNeighbors

# ---------------------------------------------------
# Load pole dynamics dataset
# ---------------------------------------------------

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

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

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

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

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

# phase-space trajectory
X = np.vstack([omega, alpha]).T

# ---------------------------------------------------
# Find nearest neighbors in phase space
# ---------------------------------------------------

nbrs = NearestNeighbors(n_neighbors=2).fit(X)

distances, indices = nbrs.kneighbors(X)

neighbor_index = indices[:,1]

# ---------------------------------------------------
# Track divergence
# ---------------------------------------------------

max_steps = 50
divergence = np.zeros(max_steps)

count = np.zeros(max_steps)

for i in range(len(X)-max_steps):

    j = neighbor_index[i]

    for k in range(max_steps):

        if i+k < len(X) and j+k < len(X):

            d = np.linalg.norm(X[i+k] - X[j+k])

            if d > 0:

                divergence[k] += np.log(d)

                count[k] += 1

divergence /= count

time = np.arange(max_steps) * dt

# ---------------------------------------------------
# Linear fit to estimate Lyapunov exponent
# ---------------------------------------------------

fit_range = slice(1,15)

coeff = np.polyfit(time[fit_range], divergence[fit_range], 1)

lyapunov = coeff[0]

print("\nEstimated largest Lyapunov exponent:", lyapunov)

# ---------------------------------------------------
# Plot divergence
# ---------------------------------------------------

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

plt.plot(time, divergence, label="trajectory divergence")

plt.plot(time, np.polyval(coeff,time), "--", label="linear fit")

plt.xlabel("Time (kyr)")
plt.ylabel("log divergence")

plt.title("Largest Lyapunov Exponent Estimate")

plt.legend()

plt.grid(alpha=0.3)

plt.tight_layout()

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

plt.show()