import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.signal import savgol_filter
from scipy.optimize import curve_fit
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 + DERIVATIVE
# ==============================

# Smooth area curve to suppress noise
area_smooth = savgol_filter(area, window_length=21, polyorder=3)

# First derivative (rate of change)
dA_dt = np.gradient(area_smooth, time)

# ==============================
# DETECT TRANSITIONS (CRITICAL)
# ==============================

# Find strongest negative slopes (collapse events)
threshold = np.percentile(dA_dt, 5)  # bottom 5% slopes

transition_indices = np.where(dA_dt < threshold)[0]

# Group into clusters (separate events)
clusters = []
current_cluster = [transition_indices[0]]

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

clusters.append(current_cluster)

# Take central point of each cluster
transition_points = [int(np.mean(cluster)) for cluster in clusters]

# ==============================
# COMPUTE ALPHA VALUES
# ==============================

alphas = []
alpha_times = []

for i in range(len(transition_points) - 1):
    idx1 = transition_points[i]
    idx2 = transition_points[i + 1]

    A1 = area_smooth[idx1]
    A2 = area_smooth[idx2]

    if A1 > 0:
        alpha = A2 / A1
        alphas.append(alpha)
        alpha_times.append(time[idx2])

alphas = np.array(alphas)

# ==============================
# EXPONENTIAL FIT
# ==============================

def exp_model(t, A0, k):
    return A0 * np.exp(-k * (time[0] - t))

params, _ = curve_fit(exp_model, time, area_smooth, p0=[1, 0.1])
exp_fit = exp_model(time, *params)

# ==============================
# GEOMETRIC MODEL (FIT α)
# ==============================

alpha_mean = np.mean(alphas)

geom = np.zeros_like(time)
geom[0] = 1.0

for i in range(1, len(time)):
    geom[i] = geom[i - 1] * alpha_mean

# ==============================
# ERROR METRICS
# ==============================

rmse_exp = np.sqrt(np.mean((area_smooth - exp_fit)**2))
rmse_geom = np.sqrt(np.mean((area_smooth - geom)**2))

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

print("\n=== ALPHA EXTRACTION RESULTS ===")
print(f"Detected transitions: {len(transition_points)}")
print(f"Alpha values: {alphas}")
print(f"Mean alpha: {alpha_mean:.3f}")
print(f"Std alpha: {np.std(alphas):.3f}")

print("\n=== MODEL COMPARISON ===")
print(f"RMSE (Exponential): {rmse_exp:.4f}")
print(f"RMSE (Geometric):   {rmse_geom:.4f}")

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

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

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

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

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

# Model comparison
plt.figure(figsize=(8,5))
plt.plot(time, area_smooth, label="Observed")
plt.plot(time, exp_fit, label="Exponential")
plt.plot(time, geom, label="Geometric (mean α)")

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

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