import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.decomposition import PCA
from collections import Counter

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

# ============================================================
# REMOVE FAST OSCILLATIONS (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

mask2 = np.isfinite(x_s) & np.isfinite(y_s)
x_s = x_s[mask2]
y_s = y_s[mask2]
t_s = t[mask2]

# ============================================================
# REMOVE SECULAR DRIFT
# ============================================================

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

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]
    
    # axial orientation (0–180°)
    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)

# ============================================================
# ROSE DIAGRAM (AXIAL)
# ============================================================

theta = np.radians(angles * 2)

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

ax.hist(theta, bins=36)

def mark_axis(deg, color, label):
    th = np.radians(deg * 2)
    ax.plot([th, th], [0, ax.get_ylim()[1]], lw=2, color=color, label=label)

mark_axis(105, 'green', "Forcing Axis (~75°W)")
mark_axis(31,  'blue',  "Inertial Plane (~31°E)")
mark_axis(170 % 180, 'red', "Excursion Plane (~170°E)")

ax.set_title("Polar Motion Anisotropy Axis Density")
ax.legend(loc='upper right')

plt.show()

# ============================================================
# MEAN AXIS + CONCENTRATION
# ============================================================

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

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

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

# ============================================================
# PHASE OCCUPANCY
# ============================================================

forcing_axis  = 105
inertial_axis = 31
excursion_axis = 170 % 180

def ang_dist(a, b):
    d = abs(a - b)
    return min(d, 180 - d)

winners = []

for a in angles:
    dF = ang_dist(a, forcing_axis)
    dI = ang_dist(a, inertial_axis)
    dE = ang_dist(a, excursion_axis)
    
    if dF <= dI and dF <= dE:
        winners.append("forcing")
    elif dI <= dF and dI <= dE:
        winners.append("inertial")
    else:
        winners.append("excursion")

counts = Counter(winners)
total = len(winners)

print("\nPhase occupancy:")
for k in ["forcing","inertial","excursion"]:
    print(f"{k:10s}: {100*counts[k]/total:6.2f} %")

# ============================================================
# TIME SERIES OF DOMINANT AXIS
# ============================================================

label_map = {"forcing":1, "inertial":2, "excursion":3}
yvals = np.array([label_map[w] for w in winners])

plt.figure(figsize=(12,4))
plt.scatter(times, yvals, s=10)

plt.yticks([1,2,3],
           ["Forcing (~75W)", "Inertial (~31E)", "Excursion (~170E)"])

plt.title("Closest Structural Axis vs Time")
plt.xlabel("Year")
plt.tight_layout()
plt.show()

# ============================================================
# ERA SPLIT ANALYSIS
# ============================================================

years = np.array([d.astype('datetime64[Y]').astype(int) + 1970 for d in times])

def occupancy_subset(start, end):
    mask = (years >= start) & (years < end)
    subset = [winners[i] for i in range(len(winners)) if mask[i]]
    
    c = Counter(subset)
    n = len(subset)
    
    print(f"\n{start}-{end}:")
    for k in ["forcing","inertial","excursion"]:
        if n > 0:
            print(f"{k:10s}: {100*c[k]/n:6.2f} %")

occupancy_subset(1973, 1990)
occupancy_subset(1990, 2005)
occupancy_subset(2005, 2030)


# ============================================================
# ROLLING CONCENTRATION R(t)
# ============================================================

def rolling_R(angles, window=60):  # ~5 years if step=30 days
    Rvals = []
    tvals = []
    for i in range(0, len(angles) - window):
        a = angles[i:i+window]
        th = np.radians(a * 2)  # axial
        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, window=60)

plt.figure(figsize=(12,4))
plt.plot(tR, Rseries)
plt.ylabel("Concentration R")
plt.title("Time-varying Anisotropy Strength")
plt.tight_layout()
plt.show()


# ============================================================
# ALIGNMENT ERROR TIME SERIES
# ============================================================

def ang_dist(a, b):
    d = abs(a - b)
    return min(d, 180 - d)

forcing_axis  = 105
inertial_axis = 31
excursion_axis = 170 % 180

dF = np.array([ang_dist(a, forcing_axis) for a in angles])
dI = np.array([ang_dist(a, inertial_axis) for a in angles])
dE = np.array([ang_dist(a, excursion_axis) for a in angles])

plt.figure(figsize=(12,5))
plt.plot(times, dF, label="Δ to Forcing (~75W)")
plt.plot(times, dI, label="Δ to Inertial (~31E)")
plt.plot(times, dE, label="Δ to Excursion (~170E)")

plt.ylabel("Angular misfit (deg)")
plt.title("Alignment Error vs Time")
plt.legend()
plt.tight_layout()
plt.show()


# ============================================================
# SIMPLE CHANGE-POINT SCAN ON R(t)
# ============================================================

# rolling mean difference between halves of the series
mid = len(Rseries)//2
early_mean = np.mean(Rseries[:mid])
late_mean  = np.mean(Rseries[mid:])

print("\nR early mean:", round(early_mean,3))
print("R late  mean:", round(late_mean,3))
print("ΔR:", round(late_mean - early_mean,3))

def epoch_stats(start, end):
    yrs = np.array([d.astype('datetime64[Y]').astype(int) + 1970 for d in tR])
    mask = (yrs >= start) & (yrs < end)
    if mask.sum() > 0:
        print(f"{start}-{end}: mean R =", round(np.mean(Rseries[mask]),3))

epoch_stats(1975,1990)
epoch_stats(1990,2005)
epoch_stats(2005,2024)


# ============================================================
# DOMINANCE MARGIN
# ============================================================

margin = np.minimum.reduce([dF, dI, dE])
second_best = np.sort(np.vstack([dF, dI, dE]), axis=0)[1]
decisiveness = second_best - margin  # larger = clearer winner

plt.figure(figsize=(12,4))
plt.plot(times, decisiveness)
plt.ylabel("Dominance margin (deg)")
plt.title("How Decisive is the Active Axis?")
plt.tight_layout()
plt.show()


