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

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

ROLL_YEARS = 2.0
WINDOW_DAYS = int(365 * ROLL_YEARS)

PCA_WINDOW = 365 * 3
STEP = 30

# Adhikari convection-driven residual azimuth (deg)
PHI_C_DEG = 259.0

# Inertial plane azimuth (deg) — set to your preferred reference
# (example placeholder; adjust if needed)
PHI_I_DEG = 170.0

BOOTSTRAP_SAMPLES = 2000

# ============================================================
# LOAD DATA
# ============================================================
df = pd.read_csv(CSV_FILE, sep=';')
df = df[df['x_pole'].notna()]
df = df[df['y_pole'].notna()]

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

# ============================================================
# SMOOTH + DETREND
# ============================================================
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

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 = []

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])

    lam = pca.explained_variance_
    R = lam[0] / lam[1]

    angles.append(theta)
    anisotropy.append(R)
    times.append(i)

angles = np.array(angles)
anisotropy = np.array(anisotropy)

# axial representation
theta2 = 2 * angles

# ============================================================
# REFERENCE VECTORS
# ============================================================
phi_c = np.deg2rad(PHI_C_DEG)
phi_i = np.deg2rad(PHI_I_DEG)

vc = np.array([np.cos(2*phi_c), np.sin(2*phi_c)])
vi = np.array([np.cos(2*phi_i), np.sin(2*phi_i)])

va = np.vstack([np.cos(theta2), np.sin(theta2)]).T

alignment_c = np.abs(va @ vc)
alignment_i = np.abs(va @ vi)

# smooth envelopes
alignment_c_s = savgol_filter(alignment_c, 101, 3)
alignment_i_s = savgol_filter(alignment_i, 101, 3)
anisotropy_s = savgol_filter(anisotropy, 101, 3)

# ============================================================
# PHASE OFFSET
# ============================================================
phase_offset = (angles - phi_c + np.pi) % np.pi - np.pi/2

# ============================================================
# CIRCULAR STATISTICS
# ============================================================
def circular_mean(theta):
    return 0.5*np.arctan2(np.mean(np.sin(2*theta)), np.mean(np.cos(2*theta)))

def resultant_length(theta):
    C = np.mean(np.cos(2*theta))
    S = np.mean(np.sin(2*theta))
    return np.sqrt(C**2 + S**2)

theta_mean = circular_mean(angles)
R_bar = resultant_length(angles)

print("\n--- CIRCULAR STATISTICS ---")
print(f"Mean orientation: {np.rad2deg(theta_mean)%180:.2f} deg")
print(f"Resultant length (anisotropy concentration): {R_bar:.3f}")

# ============================================================
# BOOTSTRAP CONFIDENCE FOR ALIGNMENT
# ============================================================
boot_means = []
N = len(angles)

for _ in range(BOOTSTRAP_SAMPLES):
    idx = np.random.randint(0, N, N)
    sample = angles[idx]
    va_s = np.vstack([np.cos(2*sample), np.sin(2*sample)]).T
    align = np.abs(va_s @ vc)
    boot_means.append(np.mean(align))

boot_means = np.array(boot_means)
ci_low, ci_high = np.percentile(boot_means, [2.5, 97.5])

print("\n--- BOOTSTRAP ALIGNMENT ---")
print(f"Mean alignment vs convection: {alignment_c.mean():.3f}")
print(f"95% CI: [{ci_low:.3f}, {ci_high:.3f}]")

# ============================================================
# REGIME-CONDITIONED STATS
# ============================================================
hi_mask = anisotropy > np.percentile(anisotropy, 75)
lo_mask = anisotropy < np.percentile(anisotropy, 25)

print("\n--- REGIME ALIGNMENT ---")
print(f"High anisotropy alignment (conv): {alignment_c[hi_mask].mean():.3f}")
print(f"Low anisotropy alignment (conv):  {alignment_c[lo_mask].mean():.3f}")

print(f"High anisotropy alignment (inert): {alignment_i[hi_mask].mean():.3f}")
print(f"Low anisotropy alignment (inert):  {alignment_i[lo_mask].mean():.3f}")

# ============================================================
# 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

# ============================================================
# PLOTS
# ============================================================

# 1) Alignment envelopes
plt.figure(figsize=(10,6))
plt.plot(times, alignment_c_s)
plt.plot(times, alignment_i_s)
plt.title("Alignment vs Convection and Inertial Axes")
plt.xlabel("Time index")
plt.ylabel("Alignment")
plt.show()

# 2) Anisotropy strength
plt.figure(figsize=(10,6))
plt.plot(times, anisotropy_s)
plt.title("Anisotropy Strength Envelope")
plt.xlabel("Time index")
plt.ylabel("R")
plt.show()

# 3) Phase offset
plt.figure(figsize=(10,6))
plt.plot(times, phase_offset)
plt.title("Phase Offset: θ(t) - φ_c")
plt.xlabel("Time index")
plt.ylabel("Radians")
plt.show()

# 4) Orientation density
plt.figure(figsize=(10,6))
plt.bar(centers, hist, width=2)
plt.axvline(PHI_C_DEG % 180)
plt.axvline(PHI_I_DEG % 180)
plt.title("Orientation Density vs Reference Axes")
plt.xlabel("Degrees")
plt.ylabel("Density")
plt.show()
