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=";")

# keep only final solutions
df = df[df["Type"] == "final"].copy()

# build datetime
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

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

# ============================================================
# REMOVE FAST OSCILLATORY ENERGY (CRITICAL STEP)
# 2-year rolling mean suppresses annual + Chandler
# ============================================================

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

# drop edges created by smoothing
mask = np.isfinite(x_s) & np.isfinite(y_s)
x_s = x_s[mask]
y_s = y_s[mask]
t_s = t[mask]

# ============================================================
# REMOVE SECULAR DRIFT (TPW / long-term trend)
# ============================================================

n = len(x_s)
idx = np.arange(n)

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 TO EXTRACT 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]
    
    # axis orientation (NOT direction)
    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)

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

plt.figure(figsize=(12,5))

plt.plot(times, angles, color="black", lw=1.6)

# reference planes
plt.axhline(285 % 180, color='green', linestyle='--', lw=2,
            label='Forcing Axis (~75°W)')
plt.axhline(31, color='blue', linestyle='--', lw=2,
            label='Inertial Plane (~31°E)')
plt.axhline(170 % 180, color='red', linestyle='--', lw=2,
            label='Excursion Plane (~170°E)')

plt.ylabel("Anisotropy Axis (° modulo 180)")
plt.title("Smoothed Polar Motion Anisotropy Axis vs Time")
plt.legend()
plt.tight_layout()
plt.show()

# ============================================================
# ROSE DIAGRAM WITH REGIME OVERLAYS
# ============================================================

theta = np.radians(angles * 2)   # axial doubling

fig = plt.figure(figsize=(7,7))
ax = fig.add_subplot(111, projection='polar')

# histogram
ax.hist(theta, bins=36, alpha=0.75)

ax.set_title("Anisotropy Axis Orientation Density with Regime Sectors")

# ------------------------------------------------------------
# Reference regime angles (in axial space)
# ------------------------------------------------------------

forcing  = 285 % 180
inertial = 31
excursion = 170 % 180

# convert to doubled-angle radians
f = np.radians(forcing * 2)
i = np.radians(inertial * 2)
e = np.radians(excursion * 2)

# ------------------------------------------------------------
# draw regime axes (mirrored)
# ------------------------------------------------------------

for center, color, label in [
    (f, 'green', 'Forcing (~75°W)'),
    (i, 'blue',  'Inertial (~31°E)'),
    (e, 'red',   'Excursion (~170°E)')
]:
    ax.plot([center, center], [0, ax.get_rmax()], color=color, lw=3, label=label)
    ax.plot([center + np.pi, center + np.pi], [0, ax.get_rmax()], color=color, lw=3)


# ------------------------------------------------------------
# draw regime sectors (±15° around each axis), MIRRORED
# ------------------------------------------------------------

sector_width = np.radians(15 * 2)

for center, color in [(f,'green'), (i,'blue'), (e,'red')]:

    # primary sector
    ax.bar(center,
           ax.get_rmax(),
           width=sector_width,
           bottom=0,
           alpha=0.10,
           color=color,
           edgecolor=None)

    # mirrored sector (180° opposite in axial space)
    ax.bar(center + np.pi,
           ax.get_rmax(),
           width=sector_width,
           bottom=0,
           alpha=0.10,
           color=color,
           edgecolor=None)

ax.legend(loc='upper right')
plt.show()

# ============================================================
# MEAN AXIS ESTIMATE (CIRCULAR STATISTICS)
# ============================================================

C = np.mean(np.cos(theta))
S = np.mean(np.sin(theta))

mean_angle = 0.5 * np.degrees(np.arctan2(S, C)) % 180
R = np.sqrt(C**2 + S**2)

print("\nMean anisotropy axis:", round(mean_angle,2), "degrees")
print("Concentration (0–1):", round(R,3))
