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

# --------------------------------
# Parameters
# --------------------------------

MIN_DT = 0.05     # kyr minimum timestep
MAX_VEL = 500     # deg/kyr clip unrealistic velocities

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

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

# Normalize longitude
df["VGP_lon"] = ((df["VGP_lon"] + 180) % 360) - 180

# Sort
df = df.sort_values("Age_ka")

# Remove duplicate ages
df = df.groupby("Age_ka", as_index=False).mean(numeric_only=True)

# Enforce minimum timestep
df = df[df["Age_ka"].diff().fillna(1) > MIN_DT]

df = df.reset_index(drop=True)

# --------------------------------
# Extract arrays
# --------------------------------

lat = np.radians(df["VGP_lat"].values)
lon = np.radians(df["VGP_lon"].values)
t = df["Age_ka"].values * 1000  # years

# --------------------------------
# Convert to unit sphere vectors
# --------------------------------

x = np.cos(lat) * np.cos(lon)
y = np.cos(lat) * np.sin(lon)
z = np.sin(lat)

r = np.vstack((x, y, z)).T

# --------------------------------
# Time derivatives
# --------------------------------

dt = np.gradient(t)

# avoid zero dt
dt[dt == 0] = np.nan

dr = np.gradient(r, axis=0) / dt[:, None]

# angular velocity magnitude
omega = np.linalg.norm(dr, axis=1)

# --------------------------------
# Angular momentum proxy
# --------------------------------

L = np.cross(r, dr)
Lmag = np.linalg.norm(L, axis=1)

# --------------------------------
# Acceleration
# --------------------------------

d2r = np.gradient(dr, axis=0) / dt[:, None]

# --------------------------------
# Curvature
# κ = |v × a| / |v|^3
# --------------------------------

cross_va = np.cross(dr, d2r)

curvature = np.linalg.norm(cross_va, axis=1) / (
    np.linalg.norm(dr, axis=1)**3 + 1e-12
)

# --------------------------------
# Convert velocity to deg/kyr
# --------------------------------

omega_deg = omega * (180 / np.pi) * 1000

# clip unrealistic spikes
omega_deg = np.clip(omega_deg, 0, MAX_VEL)

# --------------------------------
# Clean NaNs
# --------------------------------

curvature = np.nan_to_num(curvature)
Lmag = np.nan_to_num(Lmag)
omega_deg = np.nan_to_num(omega_deg)

# --------------------------------
# Save diagnostics
# --------------------------------

out = pd.DataFrame({
    "Age_ka": df["Age_ka"],
    "AngularVelocity_deg_kyr": omega_deg,
    "AngularMomentum": Lmag,
    "Curvature": curvature
})

out.to_csv("vgp_dynamics.csv", index=False)

# --------------------------------
# Plot Angular Velocity
# --------------------------------

plt.figure()
plt.plot(df["Age_ka"], omega_deg)
plt.xlabel("Age (ka)")
plt.ylabel("Angular velocity (deg/kyr)")
plt.title("Geomagnetic Pole Angular Velocity")
plt.gca().invert_xaxis()
plt.tight_layout()
plt.savefig("vgp_angular_velocity.png")
plt.show()

# --------------------------------
# Plot Angular Momentum
# --------------------------------

plt.figure()
plt.plot(df["Age_ka"], Lmag)
plt.xlabel("Age (ka)")
plt.ylabel("Angular momentum proxy")
plt.title("VGP Angular Momentum")
plt.gca().invert_xaxis()
plt.tight_layout()
plt.savefig("vgp_angular_momentum.png")
plt.show()

# --------------------------------
# Plot Curvature
# --------------------------------

plt.figure()
plt.plot(df["Age_ka"], curvature)
plt.xlabel("Age (ka)")
plt.ylabel("Trajectory curvature")
plt.title("Pole Path Curvature")
plt.gca().invert_xaxis()
plt.tight_layout()
plt.savefig("vgp_path_curvature.png")
plt.show()

print("Diagnostics written to vgp_dynamics.csv")