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

# -----------------------------
# PARAMETERS
# -----------------------------
CSV_FILE = "finals.all.csv"

# smoothing
ROLL_YEARS = 2.0
WINDOW_DAYS = int(365 * ROLL_YEARS)

# PCA window
PCA_WINDOW = 365 * 3   # 3-year sliding window
STEP = 30              # days between estimates

# Adhikari residual drift azimuth (degrees)
# residual motion ~101°W -> 360-101 = 259° azimuth
PHI_C = np.deg2rad(259.0)

# -----------------------------
# LOAD DATA
# -----------------------------
df = pd.read_csv(CSV_FILE, sep=';')

# Keep only rows with pole data
df = df[df['x_pole'].notna()]
df = df[df['y_pole'].notna()]

x = df['x_pole'].values
y = df['y_pole'].values

# -----------------------------
# SMOOTH TRAJECTORY
# -----------------------------
# rolling mean
x_s = pd.Series(x).rolling(WINDOW_DAYS, center=True, min_periods=1).mean().values
y_s = pd.Series(y).rolling(WINDOW_DAYS, center=True, min_periods=1).mean().values

# remove secular drift (linear detrend)
t = np.arange(len(x_s))
coef_x = np.polyfit(t, x_s, 1)
coef_y = np.polyfit(t, y_s, 1)

x_s -= np.polyval(coef_x, t)
y_s -= np.polyval(coef_y, t)

# -----------------------------
# SLIDING PCA
# -----------------------------
angles = []
times = []
anisotropy_strength = []

for i in range(PCA_WINDOW, len(x_s) - PCA_WINDOW, STEP):
    xs = x_s[i-PCA_WINDOW:i+PCA_WINDOW]
    ys = y_s[i-PCA_WINDOW:i+PCA_WINDOW]

    pts = np.vstack([xs, ys]).T
    pca = PCA(n_components=2)
    pca.fit(pts)

    v = pca.components_[0]

    theta = np.arctan2(v[1], v[0])
    angles.append(theta)
    times.append(i)

    # anisotropy strength
    lam = pca.explained_variance_
    R = lam[0] / lam[1]
    anisotropy_strength.append(R)

angles = np.array(angles)
anisotropy_strength = np.array(anisotropy_strength)

# -----------------------------
# ALIGNMENT WITH CONVECTION AXIS
# -----------------------------
# axial representation (double-angle)
va = np.vstack([np.cos(2*angles), np.sin(2*angles)]).T

vc = np.array([np.cos(2*PHI_C), np.sin(2*PHI_C)])

alignment = np.abs(va @ vc)

# smooth alignment envelope
alignment_smooth = savgol_filter(alignment, 101, 3)
R_smooth = savgol_filter(anisotropy_strength, 101, 3)

# -----------------------------
# ORIENTATION DENSITY
# -----------------------------
theta_deg = np.rad2deg(angles) % 180

hist, bins = np.histogram(theta_deg, bins=72, range=(0,180), density=True)
centers = (bins[:-1] + bins[1:]) / 2

# peak orientation
peak_theta = centers[np.argmax(hist)]
offset = min(abs(peak_theta - np.rad2deg(PHI_C)%180),
             180 - abs(peak_theta - np.rad2deg(PHI_C)%180))

print("\n--- RESULTS ---")
print(f"Peak anisotropy orientation: {peak_theta:.2f}°")
print(f"Convection axis orientation: {np.rad2deg(PHI_C)%180:.2f}°")
print(f"Angular offset: {offset:.2f}°")

# -----------------------------
# REGIME STATISTICS
# -----------------------------
high_R = alignment[R_smooth > np.percentile(R_smooth, 75)]
low_R  = alignment[R_smooth < np.percentile(R_smooth, 25)]

print(f"\nMean alignment (high anisotropy): {high_R.mean():.3f}")
print(f"Mean alignment (low anisotropy):  {low_R.mean():.3f}")

# -----------------------------
# PLOTS
# -----------------------------
plt.figure(figsize=(12,8))

plt.subplot(3,1,1)
plt.plot(times, alignment, alpha=0.4)
plt.plot(times, alignment_smooth, linewidth=2)
plt.title("Alignment vs Adhikari Convection Axis")
plt.ylabel("Alignment")

plt.subplot(3,1,2)
plt.plot(times, anisotropy_strength, alpha=0.4)
plt.plot(times, R_smooth, linewidth=2)
plt.title("Anisotropy Strength")

plt.subplot(3,1,3)
plt.bar(centers, hist, width=2)
plt.axvline(np.rad2deg(PHI_C)%180, linestyle='--')
plt.title("Orientation Density vs Convection Axis")
plt.xlabel("Orientation (deg)")

plt.tight_layout()
plt.show()
