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 (same as before)
# ============================================================

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 = x_s - np.polyval(np.polyfit(idx, x_s, 1), idx)
y_s = y_s - np.polyval(np.polyfit(idx, y_s, 1), idx)

# ============================================================
# ROLLING PCA — ANGLES + ECCENTRICITY
# ============================================================

window = 365 * 3
step   = 30

angles = []
ecc    = []
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
    
    # eigenvalue ratio -> eccentricity proxy
    lam1, lam2 = pca.explained_variance_
    e = np.sqrt(lam1 / lam2)
    
    angles.append(angle)
    ecc.append(e)
    times.append(t_s[i + window//2])

angles = np.array(angles)
ecc    = np.array(ecc)
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)

# ============================================================
# ALIGN ECCENTRICITY TO R(t) TIME BASE
# ============================================================

ecc_interp = np.interp(
    tR.astype("int64"),
    times.astype("int64"),
    ecc
)

# normalize
ecc_norm = (ecc_interp - np.mean(ecc_interp)) / np.std(ecc_interp)
R_norm   = (Rseries    - np.mean(Rseries))    / np.std(Rseries)

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

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

# ============================================================
# PLOT 1 — ECCENTRICITY vs R(t)
# ============================================================

plt.figure(figsize=(12,5))
plt.plot(tR, ecc_norm, label="Trajectory eccentricity")
plt.plot(tR, R_norm, label="Anisotropy strength R(t)")
plt.legend()
plt.title("Eccentricity vs Anisotropy Strength")
plt.tight_layout()
plt.show()

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

corr_window = 80
rcorr = []

for i in range(len(ecc_norm) - corr_window):
    c = np.corrcoef(
        ecc_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: Eccentricity vs R")
plt.tight_layout()
plt.show()

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

plt.figure(figsize=(6,6))
plt.scatter(ecc_norm, R_norm, s=12, alpha=0.6)
plt.xlabel("Eccentricity (normalized)")
plt.ylabel("R(t)")
plt.title("Phase Relationship")
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(ecc_norm[:L], R_norm[-L:])[0,1]
    elif L > 0:
        c = np.corrcoef(ecc_norm[L:], R_norm[:-L])[0,1]
    else:
        c = np.corrcoef(ecc_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: Eccentricity vs R")
plt.tight_layout()
plt.show()
