import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.decomposition import PCA

# -----------------------------
# LOAD DATA
# -----------------------------
df = pd.read_csv("finals.all.csv", sep=";")

# drop rows without pole values
df = df.dropna(subset=["x_pole", "y_pole", "Year", "Month", "Day"])

x = df["x_pole"].astype(float).values
y = df["y_pole"].astype(float).values

# decimal time (safe, monotonic)
years = (
    df["Year"].values +
    (df["Month"].values - 1)/12 +
    (df["Day"].values - 1)/365.25
)

# ensure strictly increasing time
t = years.copy()
dt = np.gradient(t)
dt[dt == 0] = np.nanmedian(dt)

# -----------------------------
# VELOCITY + ACCELERATION
# -----------------------------
vx = np.gradient(x, t)
vy = np.gradient(y, t)

ax = np.gradient(vx, t)
ay = np.gradient(vy, t)

speed = np.sqrt(vx**2 + vy**2)
acc_mag = np.sqrt(ax**2 + ay**2)

# -----------------------------
# CURVATURE
# κ = |vx ay − vy ax| / (speed^3)
# -----------------------------
curvature = np.abs(vx*ay - vy*ax) / (speed**3 + 1e-12)

# -----------------------------
# PCA ANISOTROPY (ROLLING)
# -----------------------------
window = 90   # days
half = window // 2

anisotropy_bearing = []
R_t = []
center_times = []

for i in range(half, len(x)-half):

    xs = x[i-half:i+half]
    ys = y[i-half:i+half]

    pts = np.vstack([xs, ys]).T

    if np.any(~np.isfinite(pts)):
        anisotropy_bearing.append(np.nan)
        R_t.append(np.nan)
        center_times.append(t[i])
        continue

    pca = PCA(n_components=2)
    pca.fit(pts)

    eigvals = pca.explained_variance_
    eigvecs = pca.components_

    # anisotropy strength
    R = (eigvals[0] - eigvals[1]) / (eigvals[0] + eigvals[1] + 1e-12)

    # major axis direction
    v = eigvecs[0]
    ang = np.degrees(np.arctan2(v[1], v[0])) % 360

    anisotropy_bearing.append(ang)
    R_t.append(R)
    center_times.append(t[i])

anisotropy_bearing = np.array(anisotropy_bearing)
R_t = np.array(R_t)
center_times = np.array(center_times)

# trim kinematic fields to match PCA centers
curvature = curvature[half:len(curvature)-half]
acc_mag = acc_mag[half:len(acc_mag)-half]
vx = vx[half:len(vx)-half]
vy = vy[half:len(vy)-half]

# -----------------------------
# DRIFT AXIS MISALIGNMENT
# -----------------------------
drift_dir = np.degrees(np.arctan2(vy, vx)) % 360

def ang_diff(a, b):
    d = np.abs(a-b)
    return np.minimum(d, 360-d)

drift_misalignment = ang_diff(drift_dir, anisotropy_bearing)

# -----------------------------
# PHASE LABELS
# -----------------------------
forcing_axis = 285
inertial_axis = 31
excursion_axis = 170

phase_labels = []

for ang in anisotropy_bearing:
    if not np.isfinite(ang):
        phase_labels.append(np.nan)
        continue

    d = [
        ang_diff(ang, forcing_axis),
        ang_diff(ang, inertial_axis),
        ang_diff(ang, excursion_axis)
    ]
    phase_labels.append(np.argmin(d))

phase_labels = np.array(phase_labels)

# -----------------------------
# CORRELATIONS
# -----------------------------
valid = (
    np.isfinite(curvature) &
    np.isfinite(R_t) &
    np.isfinite(acc_mag)
)

print("Curvature vs R:", np.corrcoef(curvature[valid], R_t[valid])[0,1])
print("Acceleration vs R:", np.corrcoef(acc_mag[valid], R_t[valid])[0,1])
print("Misalignment vs R:", np.corrcoef(drift_misalignment[valid], R_t[valid])[0,1])

# -----------------------------
# PLOTS
# -----------------------------
plt.figure(figsize=(10,6))
plt.scatter(curvature, R_t, s=5)
plt.xlabel("Curvature")
plt.ylabel("R(t)")
plt.title("Phase Relationship: Curvature vs Anisotropy Strength")
plt.show()

plt.figure(figsize=(10,6))
plt.scatter(acc_mag, R_t, s=5)
plt.xlabel("Acceleration magnitude")
plt.ylabel("R(t)")
plt.title("Acceleration vs Anisotropy Strength")
plt.show()

plt.figure(figsize=(10,6))
plt.scatter(drift_misalignment, R_t, s=5)
plt.xlabel("Drift-axis misalignment")
plt.ylabel("R(t)")
plt.title("Drift Alignment vs Anisotropy Strength")
plt.show()
