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

# ----------------------------------------
# Load stabilized pole dynamics
# ----------------------------------------

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

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

# ----------------------------------------
# Normalize signal
# ----------------------------------------

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

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

tau = 3          # delay (samples)
m = 3            # embedding dimension

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

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

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

# Corresponding ages
age_embedded = age[:N]
energy_embedded = energy[:N]

# ----------------------------------------
# Plot reconstructed attractor
# ----------------------------------------

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

sc = ax.scatter(
    X[:,0],
    X[:,1],
    X[:,2],
    c=energy_embedded,
    s=20
)

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

ax.set_title("Reconstructed Geomagnetic Attractor")

cbar = plt.colorbar(sc)
cbar.set_label("Excursion Energy")

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

# ----------------------------------------
# Highlight excursion trajectories
# ----------------------------------------

threshold = np.percentile(energy_embedded, 98)

excursions = energy_embedded > threshold

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

ax.scatter(
    X[:,0],
    X[:,1],
    X[:,2],
    color="lightgray",
    s=10
)

ax.scatter(
    X[excursions,0],
    X[excursions,1],
    X[excursions,2],
    color="red",
    s=40
)

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

ax.set_title("Excursion Escape Trajectories")

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