import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import networkx as nx
from scipy.spatial.distance import pdist, squareform

# ----------------------------------------
# Load stabilized 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 signal
omega = (omega - np.mean(omega)) / np.std(omega)

# ----------------------------------------
# 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_emb = age[:N]
energy_emb = energy[:N]

# ----------------------------------------
# Recurrence matrix
# ----------------------------------------

D = squareform(pdist(X))

eps = np.percentile(D, 10)

R = D < eps

# Remove self-links
np.fill_diagonal(R, 0)

# ----------------------------------------
# Build network
# ----------------------------------------

G = nx.from_numpy_array(R)

# ----------------------------------------
# Network metrics
# ----------------------------------------

degree = np.array([d for n, d in G.degree()])
clustering = np.array(list(nx.clustering(G).values()))

# ----------------------------------------
# Plot degree vs time
# ----------------------------------------

plt.figure()

plt.plot(age_emb, degree)

plt.gca().invert_xaxis()

plt.xlabel("Age (ka)")
plt.ylabel("Recurrence network degree")

plt.title("Geomagnetic Dynamical Connectivity")

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

plt.show()

# ----------------------------------------
# Plot clustering coefficient
# ----------------------------------------

plt.figure()

plt.plot(age_emb, clustering)

plt.gca().invert_xaxis()

plt.xlabel("Age (ka)")
plt.ylabel("Clustering coefficient")

plt.title("Recurrence Network Clustering")

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

plt.show()

# ----------------------------------------
# Highlight excursion states
# ----------------------------------------

threshold = np.percentile(energy_emb, 98)

excursions = energy_emb > threshold

plt.figure()

plt.scatter(age_emb, degree, s=10, color="gray")
plt.scatter(age_emb[excursions], degree[excursions], s=50, color="red")

plt.gca().invert_xaxis()

plt.xlabel("Age (ka)")
plt.ylabel("Network degree")

plt.title("Excursion States in Recurrence Network")

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

plt.show()

# ----------------------------------------
# Print strongest transitions
# ----------------------------------------

print("\nTop excursion candidates:\n")

excursion_indices = np.where(excursions)[0]

for i in excursion_indices:
    print(age_emb[i], energy_emb[i])