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

# -----------------------------
# Load diagnostics
# -----------------------------

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

v = df["AngularVelocity_deg_kyr"].values
k = df["Curvature"].values
L = df["AngularMomentum"].values
age = df["Age_ka"].values

# Remove NaNs
mask = np.isfinite(v) & np.isfinite(k) & np.isfinite(L)

v = v[mask]
k = k[mask]
L = L[mask]
age = age[mask]

# -----------------------------
# 3D phase space
# -----------------------------

fig = plt.figure(figsize=(8,6))

ax = fig.add_subplot(111, projection='3d')

sc = ax.scatter(
    v,
    k,
    L,
    c=age,
    s=20
)

ax.set_xlabel("Angular Velocity")
ax.set_ylabel("Curvature")
ax.set_zlabel("Angular Momentum")

ax.set_title("Geomagnetic Pole Motion Phase Space")

cbar = plt.colorbar(sc)
cbar.set_label("Age (ka)")

plt.tight_layout()
plt.savefig("geomagnetic_phase_space.png")

plt.show()