import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.metrics import pairwise_distances


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

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

omega = df["AngularVelocity_deg_kyr"].values

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

N = len(omega)

# --------------------------------------------------
# parameters
# --------------------------------------------------

delay = 2
max_dim = 8

r_vals = np.logspace(-3,1,60)

dimensions = []
D2_values = []

# --------------------------------------------------
# correlation dimension loop
# --------------------------------------------------

for m in range(2, max_dim+1):

    M = N - (m-1)*delay

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

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

    # pairwise distances
    D = pairwise_distances(X)

    D[D==0] = np.nan

    C = []

    for r in r_vals:
        count = np.nansum(D < r)
        norm = M*(M-1)
        C.append(count/norm)

    C = np.array(C)

    # remove invalid values
    mask = (C > 0)

    log_r = np.log(r_vals[mask])
    log_C = np.log(C[mask])

    # find linear scaling region automatically
    slopes = []

    window = 10

    for i in range(len(log_r)-window):

        fit = np.polyfit(
            log_r[i:i+window],
            log_C[i:i+window],
            1
        )

        slopes.append(fit[0])

    slopes = np.array(slopes)

    # choose stable slope region
    D2 = np.median(slopes)

    dimensions.append(m)
    D2_values.append(D2)

    print(f"Embedding dim {m} → D2 ≈ {D2:.3f}")


# --------------------------------------------------
# plot saturation curve
# --------------------------------------------------

plt.figure(figsize=(6,4))

plt.plot(dimensions, D2_values, "-o")

plt.xlabel("Embedding dimension")
plt.ylabel("Correlation dimension D2")

plt.title("Correlation Dimension Saturation")

plt.grid(alpha=0.3)

plt.tight_layout()

plt.savefig("geomagnetic_dimension_saturation.png", dpi=300)

plt.show()