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 POSITION (remove annual + noise)
# ============================================================

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]

# detrend
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)

# ============================================================
# COMPUTE DRIFT VECTOR (dx/dt, dy/dt)
# ============================================================

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

drift_angle = np.degrees(np.arctan2(dy, dx)) % 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 DRIFT ANGLE ONTO PCA TIMES
# ============================================================

drift_interp = np.interp(
    times.astype("int64"),
    t_s.astype("int64"),
    drift_angle
)

# misalignment angle Δθ(t)
dtheta = np.abs(drift_interp - angles)
dtheta = np.minimum(dtheta, 180 - dtheta)

# ============================================================
# ALIGN Δθ WITH R(t)
# ============================================================

dtheta_interp = np.interp(
    tR.astype("int64"),
    times.astype("int64"),
    dtheta
)

# normalize
dtheta_norm = (dtheta_interp - np.mean(dtheta_interp)) / np.std(dtheta_interp)
R_norm      = (Rseries       - np.mean(Rseries))       / np.std(Rseries)

# ============================================================
# GLOBAL CORRELATION
# ============================================================

corr = np.corrcoef(dtheta_norm, R_norm)[0,1]
print("\nDrift misalignment vs R correlation:", round(corr,3))

# ============================================================
# PLOT 1 — TIME SERIES
# ============================================================

plt.figure(figsize=(12,5))
plt.plot(tR, dtheta_norm, label="Drift–axis misalignment")
plt.plot(tR, R_norm, label="Anisotropy strength R(t)")
plt.legend()
plt.title("Drift Alignment vs Anisotropy Strength")
plt.tight_layout()
plt.show()

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

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

# ============================================================
# ROLLING CORRELATION
# ============================================================

corr_window = 80
rcorr = []

for i in range(len(dtheta_norm) - corr_window):
    c = np.corrcoef(
        dtheta_norm[i:i+corr_window],
        R_norm[i:i+corr_window]
    )[0,1]
    rcorr.append(c)

tc = tR[:len(rcorr)]

plt.figure(figsize=(12,4))
plt.plot(tc, rcorr)
plt.axhline(0, linestyle='--')
plt.title("Rolling Correlation: Misalignment vs R")
plt.tight_layout()
plt.show()

# ============================================================
# LEAD/LAG TEST
# ============================================================

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

for L in lags:
    if L < 0:
        c = np.corrcoef(dtheta_norm[:L], R_norm[-L:])[0,1]
    elif L > 0:
        c = np.corrcoef(dtheta_norm[L:], R_norm[:-L])[0,1]
    else:
        c = np.corrcoef(dtheta_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: Drift Alignment vs R")
plt.tight_layout()
plt.show()
