import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

# -------------------------------
# Load dynamics
# -------------------------------

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

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

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

# -------------------------------
# Takens embedding
# -------------------------------

tau = 3
m = 3

N = len(omega) - (m - 1) * tau

X = np.zeros((N,m))

for i in range(m):
    X[:,i] = omega[i*tau : i*tau + N]

age = age[:N]
energy = energy[:N]

# -------------------------------
# plot ribbon
# -------------------------------

fig = plt.figure(figsize=(8,6))
ax = fig.add_subplot(111, projection='3d')

for i in range(N-1):

    color = plt.cm.viridis(energy[i]/energy.max())

    ax.plot(
        X[i:i+2,0],
        X[i:i+2,1],
        X[i:i+2,2],
        color=color,
        linewidth=2
    )

ax.set_xlabel("ω(t)")
ax.set_ylabel("ω(t+τ)")
ax.set_zlabel("ω(t+2τ)")

ax.set_title("Geomagnetic Pole Trajectory Ribbon")

plt.tight_layout()
plt.savefig("geomagnetic_ribbon.png")
plt.show()