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=";")
df = df[df["Type"] == "final"].copy()

df["date"] = pd.to_datetime(dict(year=df.Year, month=df.Month, day=df.Day))

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

mask = np.isfinite(x) & np.isfinite(y)
x = x[mask]
y = y[mask]
t = t[mask]

# ============================================================
# SMOOTH + DETREND
# ============================================================

smooth_days = 365 * 2

x_s = pd.Series(x).rolling(smooth_days, center=True).mean().values
y_s = pd.Series(y).rolling(smooth_days, center=True).mean().values

mask2 = np.isfinite(x_s) & np.isfinite(y_s)
x_s = x_s[mask2]
y_s = y_s[mask2]
t_s = t[mask2]

idx = np.arange(len(x_s))
x_s -= np.polyval(np.polyfit(idx, x_s, 1), idx)
y_s -= np.polyval(np.polyfit(idx, y_s, 1), idx)

# ============================================================
# VELOCITY + ACCELERATION
# ============================================================

dx = np.gradient(x_s)
dy = np.gradient(y_s)

ddx = np.gradient(dx)
ddy = np.gradient(dy)

acc_mag = np.sqrt(ddx**2 + ddy**2)
acc_angle = np.degrees(np.arctan2(ddy, ddx)) % 180

# ============================================================
# ROLLING PCA → ANISOTROPY AXIS
# ============================================================

window = 365 * 3
step   = 30

angles = []
times  = []

for i in range(0, len(x_s) - window, step):
    Xw = np.vstack([x_s[i:i+window], y_s[i:i+window]]).T
    
    pca = PCA(n_components=2)
    pca.fit(Xw)
    
    vec = pca.components_[0]
    angle = np.degrees(np.arctan2(vec[1], vec[0])) % 180
    
    angles.append(angle)
    times.append(t_s[i + window//2])

angles = np.array(angles)
times  = np.array(times)

# ============================================================
# COMPUTE R(t)
# ============================================================

def rolling_R(angles, times, window=60):
    Rvals = []
    tvals = []
    
    for i in range(0, len(angles) - window):
        a = angles[i:i+window]
        th = np.radians(a * 2)
        
        C = np.mean(np.cos(th))
        S = np.mean(np.sin(th))
        R = np.sqrt(C**2 + S**2)
        
        Rvals.append(R)
        tvals.append(times[i + window//2])
        
    return np.array(tvals), np.array(Rvals)

tR, Rseries = rolling_R(angles, times, window=60)

# ============================================================
# INTERPOLATE ACCELERATION ONTO PCA TIMES
# ============================================================

acc_mag_interp = np.interp(
    times.astype("int64"),
    t_s.astype("int64"),
    acc_mag
)

acc_ang_interp = np.interp(
    times.astype("int64"),
    t_s.astype("int64"),
    acc_angle
)

# ============================================================
# MISALIGNMENT BETWEEN ACCELERATION DIRECTION & AXIS
# ============================================================

dtheta_acc = np.abs(acc_ang_interp - angles)
dtheta_acc = np.minimum(dtheta_acc, 180 - dtheta_acc)

# ============================================================
# ALIGN BOTH WITH R(t)
# ============================================================

acc_mag_R = np.interp(
    tR.astype("int64"),
    times.astype("int64"),
    acc_mag_interp
)

dtheta_R = np.interp(
    tR.astype("int64"),
    times.astype("int64"),
    dtheta_acc
)

# normalize
acc_norm   = (acc_mag_R - np.mean(acc_mag_R)) / np.std(acc_mag_R)
mis_norm   = (dtheta_R  - np.mean(dtheta_R))  / np.std(dtheta_R)
R_norm     = (Rseries   - np.mean(Rseries))   / np.std(Rseries)

# ============================================================
# GLOBAL CORRELATIONS
# ============================================================

corr_acc = np.corrcoef(acc_norm, R_norm)[0,1]
corr_mis = np.corrcoef(mis_norm, R_norm)[0,1]

print("\nAcceleration magnitude vs R correlation:", round(corr_acc,3))
print("Acceleration direction misalignment vs R correlation:", round(corr_mis,3))

# ============================================================
# PLOT 1 — ACCELERATION MAGNITUDE vs R
# ============================================================

plt.figure(figsize=(12,5))
plt.plot(tR, acc_norm, label="Acceleration magnitude")
plt.plot(tR, R_norm,   label="Anisotropy strength R(t)")
plt.legend()
plt.title("Acceleration Magnitude vs R(t)")
plt.tight_layout()
plt.show()

# ============================================================
# PLOT 2 — ACCELERATION MISALIGNMENT vs R
# ============================================================

plt.figure(figsize=(12,5))
plt.plot(tR, mis_norm, label="Acceleration-axis misalignment")
plt.plot(tR, R_norm,   label="R(t)")
plt.legend()
plt.title("Acceleration Direction vs R(t)")
plt.tight_layout()
plt.show()

# ============================================================
# ROLLING CORRELATION (MAGNITUDE)
# ============================================================

win = 80
rc_mag = []

for i in range(len(acc_norm) - win):
    c = np.corrcoef(acc_norm[i:i+win], R_norm[i:i+win])[0,1]
    rc_mag.append(c)

plt.figure(figsize=(12,4))
plt.plot(tR[:len(rc_mag)], rc_mag)
plt.axhline(0, linestyle='--')
plt.title("Rolling Correlation: Acceleration Magnitude vs R")
plt.tight_layout()
plt.show()

# ============================================================
# LEAD/LAG TEST (MAGNITUDE)
# ============================================================

lags = np.arange(-50, 51)
lag_corr = []

for L in lags:
    if L < 0:
        c = np.corrcoef(acc_norm[:L], R_norm[-L:])[0,1]
    elif L > 0:
        c = np.corrcoef(acc_norm[L:], R_norm[:-L])[0,1]
    else:
        c = np.corrcoef(acc_norm, R_norm)[0,1]
    lag_corr.append(c)

plt.figure(figsize=(8,4))
plt.plot(lags, lag_corr)
plt.axvline(0, linestyle='--')
plt.title("Lead/Lag: Acceleration Magnitude vs R")
plt.tight_layout()
plt.show()

# ============================================================
# SCATTER PHASE RELATIONSHIP
# ============================================================

plt.figure(figsize=(6,6))
plt.scatter(acc_norm, R_norm, s=12, alpha=0.6)
plt.xlabel("Acceleration magnitude (normalized)")
plt.ylabel("R(t)")
plt.title("Phase Relationship: Acceleration vs R")
plt.tight_layout()
plt.show()
