import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.signal import savgol_filter
import os

OUT_DIR = "outputs"
os.makedirs(OUT_DIR, exist_ok=True)

# ==============================
# LOAD DATA
# ==============================

df = pd.read_csv("outputs/area_timeseries.csv")

time = df["time_ka"].values
area = df["area_norm"].values

# Ensure descending time (30 → 0)
if time[0] < time[-1]:
    time = time[::-1]
    area = area[::-1]

# ==============================
# SMOOTH DATA
# ==============================

area_smooth = savgol_filter(area, window_length=21, polyorder=3)

# ==============================
# RESTRICT TO ACTIVE REGION
# ==============================

# Only analyze where system is actually evolving
active_mask = area_smooth < 0.95

time_active = time[active_mask]
area_active = area_smooth[active_mask]

# ==============================
# CURVATURE (SECOND DERIVATIVE)
# ==============================

dA_dt = np.gradient(area_active, time_active)
d2A_dt2 = np.gradient(dA_dt, time_active)

# ==============================
# DETECT TRUE TRANSITIONS
# ==============================

# Find peaks in curvature magnitude
threshold = np.percentile(np.abs(d2A_dt2), 90)

candidate_idx = np.where(np.abs(d2A_dt2) > threshold)[0]

# Cluster nearby points into single transitions
clusters = []
if len(candidate_idx) > 0:
    current = [candidate_idx[0]]

    for idx in candidate_idx[1:]:
        if idx - current[-1] < 5:
            current.append(idx)
        else:
            clusters.append(current)
            current = [idx]

    clusters.append(current)

# Take center of clusters
transition_idx = [int(np.mean(c)) for c in clusters]

# Limit to max 3 transitions (physical expectation)
if len(transition_idx) > 3:
    # keep strongest curvature peaks only
    strengths = [np.max(np.abs(d2A_dt2[c])) for c in clusters]
    order = np.argsort(strengths)[-3:]
    transition_idx = [transition_idx[i] for i in order]
    transition_idx = sorted(transition_idx)

# ==============================
# EXTRACT STATE VALUES
# ==============================

states = []

# Add pre-collapse state
states.append(area_active[0])

for idx in transition_idx:
    states.append(area_active[idx])

# Add final state
states.append(area_active[-1])

states = np.array(states)

# ==============================
# COMPUTE ALPHA (PAIRWISE)
# ==============================

alphas = []

for i in range(len(states) - 1):
    if states[i] > 0:
        alpha = states[i + 1] / states[i]
        alphas.append(alpha)

alphas = np.array(alphas)

# ==============================
# BUILD PIECEWISE GEOMETRIC MODEL
# ==============================

geom = np.ones_like(area_smooth)

# Assign values stepwise
if len(states) >= 2:
    thresholds = [time_active[idx] for idx in transition_idx]

    for i, t in enumerate(time):
        if t > thresholds[0]:
            geom[i] = states[0]
        elif len(thresholds) > 1 and t > thresholds[1]:
            geom[i] = states[1]
        elif len(thresholds) > 2 and t > thresholds[2]:
            geom[i] = states[2]
        else:
            geom[i] = states[-1]

# ==============================
# OUTPUT RESULTS
# ==============================

print("\n=== CORRECTED ALPHA EXTRACTION ===")
print(f"Detected transitions: {len(transition_idx)}")

print("\nState values:")
for i, s in enumerate(states):
    print(f"A{i}: {s:.3f}")

print("\nAlpha values:")
for i, a in enumerate(alphas):
    print(f"alpha_{i+1}: {a:.3f}")

print(f"\nMean alpha: {np.mean(alphas):.3f}")
print(f"Std alpha: {np.std(alphas):.3f}")

# Save alpha
pd.DataFrame({
    "alpha": alphas
}).to_csv(f"{OUT_DIR}/alpha_values_corrected.csv", index=False)

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

# Plot transitions
plt.figure(figsize=(8,5))
plt.plot(time, area, alpha=0.3, label="Raw")
plt.plot(time, area_smooth, label="Smoothed")

for idx in transition_idx:
    plt.axvline(time_active[idx], linestyle="--")

plt.gca().invert_xaxis()
plt.xlabel("Time (ka BP)")
plt.ylabel("Normalized Area")
plt.title("Corrected Transition Detection")
plt.legend()
plt.savefig(f"{OUT_DIR}/transitions_corrected.png", dpi=200)

# Plot geometric model
plt.figure(figsize=(8,5))
plt.plot(time, area_smooth, label="Observed")
plt.plot(time, geom, label="Piecewise Geometric")

plt.gca().invert_xaxis()
plt.xlabel("Time (ka BP)")
plt.ylabel("Normalized Area")
plt.title("Corrected Geometric Model")
plt.legend()
plt.savefig(f"{OUT_DIR}/geometric_corrected.png", dpi=200)

print(f"\nOutputs saved to: {OUT_DIR}")