# %% [markdown]
# # GLAMBIE v17a_TIModel — Comparison Plots
#
# Two comparisons:
#   A. glambie_submission_fixed_geometry_v17a_TIModel vs
#      glambie_submission_fixed_geometry_v163_TIModel
#      (same "fixed geometry" method, old vs new OGGM source/calibration)
#   B. glambie_submission_spinup_geometry_v17a_TIModel vs
#      glambie_submission_fixed_geometry_v17a_TIModel
#      (same source, fixed vs dynamic geometry)

# %%
import numpy as np
import pandas as pd
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

REGIONS = list(range(1, 20))
CMAP = plt.cm.tab20
colors = {r: CMAP(i / len(REGIONS)) for i, r in enumerate(REGIONS)}
reg_labels = {r: f"RGI{r:02d}" for r in REGIONS}
labels = [reg_labels[r] for r in REGIONS] + ["Global"]
bar_colors = [colors[r] for r in REGIONS] + ["#333333"]
x = np.arange(len(labels))


def load(fp):
    df = pd.read_csv(fp)
    df["time"] = pd.to_datetime(df["start_date"], dayfirst=True)
    df["region_id"] = df["region_id"].astype(int)
    return df


def annual_by_region(df):
    """Annual specific MB (m w.e./yr) per region, and annual area (km2)."""
    ann_mb = (
        df.groupby(["region_id", df["time"].dt.year])["glacier_change_observed"]
        .sum()
        .unstack(level=0)
    )
    ann_area = (
        df.groupby(["region_id", df["time"].dt.year])["glacier_area_reference_start"]
        .first()
        .unstack(level=0)
    )
    return ann_mb, ann_area


def cumulative_gt_by_region(ann_mb, ann_area):
    """Cumulative mass change (Gt) per region: sum(specific_mb * area * 1e-3)."""
    mass_gt = ann_mb * ann_area * 1e-3
    return mass_gt.cumsum()


def global_mass_gt(ann_mb, ann_area):
    mass_gt = ann_mb * ann_area * 1e-3
    return mass_gt.sum(axis=1)  # sum across regions per year


# %%
fixed_v17a = load("glambie_submission_fixed_geometry_v17a_TIModel.csv")
fixed_v163 = load("glambie_submission_fixed_geometry_v163_TIModel.csv")
spinup_v17a = load("glambie_submission_spinup_geometry_v17a_TIModel.csv")

ann_mb_v17a, ann_area_v17a = annual_by_region(fixed_v17a)
ann_mb_v163, ann_area_v163 = annual_by_region(fixed_v163)
ann_mb_spin, ann_area_spin = annual_by_region(spinup_v17a)

cum_v17a = cumulative_gt_by_region(ann_mb_v17a, ann_area_v17a)
cum_v163 = cumulative_gt_by_region(ann_mb_v163, ann_area_v163)
cum_spin = cumulative_gt_by_region(ann_mb_spin, ann_area_spin)

glob_v17a = global_mass_gt(ann_mb_v17a, ann_area_v17a).cumsum()
glob_v163 = global_mass_gt(ann_mb_v163, ann_area_v163).cumsum()
glob_spin = global_mass_gt(ann_mb_spin, ann_area_spin).cumsum()

# %% [markdown]
# ## A1 — Global cumulative mass change: v17a_TIModel vs v163_TIModel (fixed geometry)

# %%
fig, ax = plt.subplots(figsize=(9, 5))
ax.plot(glob_v163.index, glob_v163.values, color="0.4", lw=2, label="fixed-geometry v163_TIModel")
ax.plot(glob_v17a.index, glob_v17a.values, color="C3", lw=2, label="fixed-geometry v17a_TIModel")
ax.axhline(0, color="k", lw=0.5, ls="--")
ax.set_ylabel("Cumulative mass change (Gt)")
ax.set_title("Global cumulative mass change 1975-2025 — fixed geometry, v163 vs v17a source")
ax.legend()
fig.tight_layout()
fig.savefig("plot_A1_global_v17a_vs_v163.png", dpi=150)

# %% [markdown]
# ## A2 — Cumulative mass change by region (Gt, 1975-2025): v17a vs v163

# %%
end_v17a = [cum_v17a[r].iloc[-1] for r in REGIONS] + [glob_v17a.iloc[-1]]
end_v163 = [cum_v163[r].iloc[-1] for r in REGIONS] + [glob_v163.iloc[-1]]

fig, ax = plt.subplots(figsize=(13, 5))
w = 0.35
ax.bar(x - w / 2, end_v163, width=w, color="0.6", label="v163_TIModel")
ax.bar(x + w / 2, end_v17a, width=w, color="C3", alpha=0.85, label="v17a_TIModel")
ax.axhline(0, color="k", lw=0.5, ls="--")
ax.set_xticks(x)
ax.set_xticklabels(labels, rotation=45, ha="right")
ax.set_ylabel("Cumulative mass change 1975-2025 (Gt)")
ax.set_title("Fixed geometry: cumulative mass change by region — v17a vs v163")
ax.legend()
ax.grid(axis="y", alpha=0.3)
fig.tight_layout()
fig.savefig("plot_A2_by_region_v17a_vs_v163.png", dpi=150)

# %% [markdown]
# ## A2b — Cumulative mass-change DIFFERENCE only (v17a - v163), by region
# (separate scale from A2 so small regions aren't dwarfed by RGI01/RGI19/etc.)

# %%
diff_end_a = [b - a for a, b in zip(end_v163, end_v17a)]
fig, ax = plt.subplots(figsize=(13, 5))
ax.bar(x, diff_end_a, color=bar_colors, alpha=0.85)
ax.axhline(0, color="k", lw=0.5, ls="--")
ax.set_xticks(x)
ax.set_xticklabels(labels, rotation=45, ha="right")
ax.set_ylabel("Δ cumulative mass change, v17a - v163 (Gt)")
ax.set_title("Fixed geometry: v17a vs v163 difference by region (own scale)")
ax.grid(axis="y", alpha=0.3)
fig.tight_layout()
fig.savefig("plot_A2b_diff_only_v17a_vs_v163.png", dpi=150)

# %% [markdown]
# ## A3 — Annual specific MB difference (v17a - v163), by region

# %%
diff_a = ann_mb_v17a - ann_mb_v163
fig, ax = plt.subplots(figsize=(14, 5))
for r in REGIONS:
    ax.plot(diff_a.index, diff_a[r], color=colors[r], lw=0.9, alpha=0.8, label=reg_labels[r])
ax.axhline(0, color="k", lw=0.5, ls="--")
ax.set_ylabel("Δ specific MB, v17a - v163 (m w.e. yr⁻¹)")
ax.set_title("Fixed geometry: annual specific-MB difference per region (v17a_TIModel - v163_TIModel)")
ax.legend(ncol=4, fontsize=7, loc="lower left")
fig.tight_layout()
fig.savefig("plot_A3_annual_diff_v17a_vs_v163.png", dpi=150)

# %% [markdown]
# ## B1 — Global cumulative mass change: spinup-geometry vs fixed-geometry (both v17a_TIModel)

# %%
fig, ax = plt.subplots(figsize=(9, 5))
ax.plot(glob_v17a.index, glob_v17a.values, color="C3", lw=2, label="fixed-geometry v17a_TIModel")
ax.plot(glob_spin.index, glob_spin.values, color="C0", lw=2, label="spinup-geometry v17a_TIModel")
ax.axhline(0, color="k", lw=0.5, ls="--")
ax.set_ylabel("Cumulative mass change (Gt)")
ax.set_title("Global cumulative mass change 1975-2025 — fixed vs spinup (varying) geometry")
ax.legend()
fig.tight_layout()
fig.savefig("plot_B1_global_spinup_vs_fixed.png", dpi=150)

# %% [markdown]
# ## B2 — Cumulative mass change by region (Gt, 1975-2025): spinup vs fixed

# %%
end_spin = [cum_spin[r].iloc[-1] for r in REGIONS] + [glob_spin.iloc[-1]]

fig, ax = plt.subplots(figsize=(13, 5))
ax.bar(x - w / 2, end_v17a, width=w, color="C3", alpha=0.85, label="fixed-geometry")
ax.bar(x + w / 2, end_spin, width=w, color="C0", alpha=0.85, label="spinup-geometry")
ax.axhline(0, color="k", lw=0.5, ls="--")
ax.set_xticks(x)
ax.set_xticklabels(labels, rotation=45, ha="right")
ax.set_ylabel("Cumulative mass change 1975-2025 (Gt)")
ax.set_title("v17a_TIModel: cumulative mass change by region — fixed vs spinup geometry")
ax.legend()
ax.grid(axis="y", alpha=0.3)
fig.tight_layout()
fig.savefig("plot_B2_by_region_spinup_vs_fixed.png", dpi=150)

# %% [markdown]
# ## B2b — Cumulative mass-change DIFFERENCE only (spinup - fixed), by region
# (separate scale from B2 so small regions like RGI16 aren't dwarfed by RGI01/RGI19)

# %%
diff_end_b = [b - a for a, b in zip(end_v17a, end_spin)]
fig, ax = plt.subplots(figsize=(13, 5))
ax.bar(x, diff_end_b, color=bar_colors, alpha=0.85)
ax.axhline(0, color="k", lw=0.5, ls="--")
ax.set_xticks(x)
ax.set_xticklabels(labels, rotation=45, ha="right")
ax.set_ylabel("Δ cumulative mass change, spinup - fixed (Gt)")
ax.set_title("v17a_TIModel: spinup vs fixed geometry difference by region (own scale)")
ax.grid(axis="y", alpha=0.3)
fig.tight_layout()
fig.savefig("plot_B2b_diff_only_spinup_vs_fixed.png", dpi=150)

# %% [markdown]
# ## B2c — Cumulative mass-change difference, spinup vs fixed (% of fixed-geometry total)
# (highlights regions with a large *relative* shift, e.g. RGI16's sign flip, even
# though its absolute Gt difference is small next to RGI01)

# %%
pct_diff_b = [100 * (b - a) / abs(a) for a, b in zip(end_v17a, end_spin)]
YLIM = 30  # RGI16's fixed-geometry baseline (+2.58 Gt) is near-zero, so its
           # % difference (-593%) is mathematically correct but off-scale for
           # every other region - clip and annotate it instead.
fig, ax = plt.subplots(figsize=(13, 5))
ax.bar(x, np.clip(pct_diff_b, -YLIM, YLIM), color=bar_colors, alpha=0.85)
for xi, v in zip(x, pct_diff_b):
    if abs(v) > YLIM:
        ax.annotate(f"{v:+.0f}%\n(near-zero baseline)", (xi, np.sign(v) * YLIM),
                    ha="center", va="bottom" if v > 0 else "top", fontsize=7)
ax.axhline(0, color="k", lw=0.5, ls="--")
ax.set_ylim(-YLIM * 1.15, YLIM * 1.15)
ax.set_xticks(x)
ax.set_xticklabels(labels, rotation=45, ha="right")
ax.set_ylabel("Δ cumulative mass change, spinup vs fixed (% of |fixed| total)")
ax.set_title("v17a_TIModel: spinup vs fixed geometry — relative difference by region")
ax.grid(axis="y", alpha=0.3)
fig.tight_layout()
fig.savefig("plot_B2c_pct_diff_spinup_vs_fixed.png", dpi=150)

# %% [markdown]
# ## B3 — Annual specific MB difference (spinup - fixed), by region

# %%
diff_b = ann_mb_spin - ann_mb_v17a
fig, ax = plt.subplots(figsize=(14, 5))
for r in REGIONS:
    ax.plot(diff_b.index, diff_b[r], color=colors[r], lw=0.9, alpha=0.8, label=reg_labels[r])
ax.axhline(0, color="k", lw=0.5, ls="--")
ax.set_ylabel("Δ specific MB, spinup - fixed (m w.e. yr⁻¹)")
ax.set_title("v17a_TIModel: annual specific-MB difference per region (spinup - fixed geometry)")
ax.legend(ncol=4, fontsize=7, loc="lower left")
fig.tight_layout()
fig.savefig("plot_B3_annual_diff_spinup_vs_fixed.png", dpi=150)

# %% [markdown]
# ## B4 — Regional area evolution: dynamic (spinup) vs constant (fixed)

# %%
area_change_pct = (ann_area_spin.iloc[-1] - ann_area_spin.iloc[0]) / ann_area_spin.iloc[0] * 100

fig, ax = plt.subplots(figsize=(14, 5))
for r in REGIONS:
    norm = ann_area_spin[r] / ann_area_spin[r].iloc[0] * 100
    ax.plot(norm.index, norm.values, color=colors[r], lw=1.1, alpha=0.85, label=reg_labels[r])
ax.axhline(100, color="k", lw=0.5, ls="--")
ax.set_ylabel("Glacier area (% of 1975 area)")
ax.set_title("Spinup-geometry: dynamic regional glacier area, normalized to 1975 (fixed-geometry stays at 100% by definition)")
ax.legend(ncol=4, fontsize=7, loc="lower left")
fig.tight_layout()
fig.savefig("plot_B4_area_evolution.png", dpi=150)

# %% [markdown]
# ## Summary tables (printed for the report)

# %%
summary_a = pd.DataFrame({
    "region": labels,
    "v163_TIModel_Gt": [round(v, 2) for v in end_v163],
    "v17a_TIModel_Gt": [round(v, 2) for v in end_v17a],
    "diff_Gt": [round(a - b, 2) for a, b in zip(end_v17a, end_v163)],
})
print("=== Comparison A: fixed-geometry v17a vs v163 (cumulative Gt, 1975-2025) ===")
print(summary_a.to_string(index=False))
print()

summary_b = pd.DataFrame({
    "region": labels,
    "fixed_v17a_Gt": [round(v, 2) for v in end_v17a],
    "spinup_v17a_Gt": [round(v, 2) for v in end_spin],
    "diff_Gt": [round(a - b, 2) for a, b in zip(end_spin, end_v17a)],
    "area_change_pct_1975_2025": [round(area_change_pct[r], 2) if r in REGIONS else round((glob_spin.iloc[-1] - glob_spin.iloc[0]), 2) for r in REGIONS] + [np.nan],
})
print("=== Comparison B: spinup vs fixed geometry, both v17a (cumulative Gt, 1975-2025) ===")
print(summary_b.to_string(index=False))

summary_a.to_csv("comparison_A_v17a_vs_v163.csv", index=False)
summary_b.to_csv("comparison_B_spinup_vs_fixed.csv", index=False)
print("\nWrote comparison_A_v17a_vs_v163.csv and comparison_B_spinup_vs_fixed.csv")