# %% [markdown]
# # GLAMBIE Fixed-Geometry Submission — Diagnostic Plots
#
# Four plots based on `glambie_submission_fixed_geometry.csv`:
# 1. Monthly MB time series (spaghetti by region)
# 2. Mean annual cycle (spaghetti by region)
# 3. Annual mass loss rate per region and globally (Gt yr⁻¹, 2000–2019),
#    compared against Hugonnet et al. 2021 `dmdt`.
# 4. Specific MB rate per region and globally (m w.e. yr⁻¹, 2000–2019),
#    compared against Hugonnet et al. 2021 `dmdtda`.

# %%
import numpy as np
import pandas as pd
import matplotlib
matplotlib.use("Agg")   # remove when running interactively in Jupyter
import matplotlib.pyplot as plt
from oggm import cfg, utils

cfg.initialize(logging_level="WARNING")

# %%
df = pd.read_csv("glambie_submission_fixed_geometry.csv")
df["time"] = pd.to_datetime(df["start_date"], dayfirst=True)
df["region_id"] = df["region_id"].astype(int)

REGIONS = sorted(df["region_id"].unique())
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}

area = df.groupby("region_id")["glacier_area_reference_start"].first()  # km²

# %% [markdown]
# ## 1 — Monthly MB time series

# %%
fig, ax = plt.subplots(figsize=(14, 5))
for rid, grp in df.groupby("region_id"):
    ax.plot(
        grp["time"], grp["glacier_change_observed"],
        color=colors[rid], lw=0.6, alpha=0.75, label=reg_labels[rid],
    )
ax.axhline(0, color="k", lw=0.5, ls="--")
ax.set_ylabel("Specific MB (m w.e. month⁻¹)")
ax.set_title("Monthly glacier-wide specific MB — OGGM fixed geometry, all RGI regions")
ax.legend(ncol=4, fontsize=7, loc="lower left")
fig.tight_layout()
fig.savefig("plot_monthly_ts.png", dpi=150)
plt.show()

# %% [markdown]
# ## 2 — Mean annual cycle

# %%
df["month"] = df["time"].dt.month
cycle = (
    df.groupby(["region_id", "month"])["glacier_change_observed"]
    .mean()
    .unstack(level=0)
)

MONTH_LABELS = ["J", "F", "M", "A", "M", "J", "J", "A", "S", "O", "N", "D"]

fig, ax = plt.subplots(figsize=(9, 5))
for rid in REGIONS:
    ax.plot(
        range(1, 13), cycle[rid],
        color=colors[rid], marker="o", ms=3, lw=1.2, label=reg_labels[rid],
    )
ax.axhline(0, color="k", lw=0.5, ls="--")
ax.set_xticks(range(1, 13))
ax.set_xticklabels(MONTH_LABELS)
ax.set_ylabel("Mean specific MB (m w.e. month⁻¹)")
ax.set_title("Mean annual cycle of specific MB — 1975–2025")
ax.legend(ncol=4, fontsize=7)
fig.tight_layout()
fig.savefig("plot_annual_cycle.png", dpi=150)
plt.show()

# %% [markdown]
# ## Hugonnet et al. 2021 reference (2000–2019)

# %%
hug_all = utils.get_geodetic_mb_dataframe(regional=True)
period = "2000-01-01_2020-01-01"

hug = hug_all[
    (hug_all["period"] == period) & (hug_all["reg"].isin(REGIONS))
].copy()
hug.index = hug["reg"].astype(int)

# Global aggregate pre-computed by Hugonnet (reg=21)
hug_glob_row = hug_all[(hug_all["period"] == period) & (hug_all["reg"] == 21)].iloc[0]

# %% [markdown]
# ## OGGM estimates for 2000–2019 (Hugonnet period)

# %%
mask = df["time"].dt.year.between(2000, 2019)
df_hp = df[mask]

# dmdtda: mean annual specific MB (mwe yr⁻¹)
dmdtda_our = (
    df_hp.groupby(["region_id", df_hp["time"].dt.year])["glacier_change_observed"]
    .sum()                  # sum 12 months → mwe yr⁻¹
    .groupby(level=0)
    .mean()
)

# dmdt: annual mass loss rate (Gt yr⁻¹)
dmdt_our = dmdtda_our * area * 1e-3

# Global: area-weighted specific MB and sum of regional mass rates
total_area = area.sum()
dmdtda_global_our = (dmdtda_our * area).sum() / total_area
dmdt_global_our = dmdt_our.sum()

# %% [markdown]
# ## 3 — Annual mass loss rate by region (Gt yr⁻¹, 2000–2019)

# %%
labels = [reg_labels[r] for r in REGIONS] + ["Global"]
bar_colors = [colors[r] for r in REGIONS] + ["#333333"]
x = np.arange(len(labels))

our_dmdt = [-dmdt_our[r] for r in REGIONS] + [-dmdt_global_our]
hug_dmdt = [-hug.loc[r, "dmdt"] for r in REGIONS] + [-hug_glob_row["dmdt"]]
hug_dmdt_err = [hug.loc[r, "err_dmdt"] for r in REGIONS] + [hug_glob_row["err_dmdt"]]

fig, ax = plt.subplots(figsize=(13, 5))
ax.bar(x, our_dmdt, color=bar_colors, alpha=0.8, zorder=2, label="OGGM 2000–2019")
ax.errorbar(
    x, hug_dmdt, yerr=hug_dmdt_err,
    fmt="o", color="k", ms=5, lw=1.2, capsize=3, zorder=3,
    label="Hugonnet et al. 2021 (2000–2019)",
)
ax.set_yscale("log")
ax.set_xticks(x)
ax.set_xticklabels(labels, rotation=45, ha="right")
ax.set_ylabel("Mass loss rate (Gt yr⁻¹)")
ax.set_title("Annual mass loss rate — OGGM fixed geometry vs Hugonnet et al. 2021 (2000–2019)")
ax.legend(fontsize=9)
ax.grid(axis="y", which="both", alpha=0.3, zorder=1)
fig.tight_layout()
fig.savefig("plot_mass_loss_by_region.png", dpi=150)
plt.show()

# %% [markdown]
# ## 4 — Specific MB rate by region (m w.e. yr⁻¹, 2000–2019)
#
# Three series:
# - Bars: OGGM fixed-geometry (2000–2019 average)
# - Black dots: Hugonnet et al. 2021 official regional `dmdtda` (includes
#   extrapolation to unmeasured glacier area via `is_cor` corrected values)
# - Grey dots: per-glacier Hugonnet `dmdtda` aggregated with Hugonnet area
#   weights (no extrapolation — same glaciers, same aggregation method as OGGM)

# %%
# Per-glacier Hugonnet aggregated with Hugonnet area weights
gdf_all = utils.get_geodetic_mb_dataframe()
gdf_p = gdf_all[gdf_all["period"] == "2000-01-01_2020-01-01"].copy()
gdf_p["reg"] = gdf_p["reg"].astype(int)
hug_pergla = gdf_p.groupby("reg").apply(
    lambda g: pd.Series({
        "dmdtda": (g["dmdtda"] * g["area"]).sum() / g["area"].sum(),
    }),
    include_groups=False,
).reindex(REGIONS)
# Global: area-weighted mean across all 19 regions
_all = gdf_p[gdf_p["reg"].isin(REGIONS)]
hug_pergla_global = ((_all["dmdtda"] * _all["area"]).sum() / _all["area"].sum())

# %%
our_dmdtda = [dmdtda_our[r] for r in REGIONS] + [dmdtda_global_our]
hug_dmdtda = [hug.loc[r, "dmdtda"] for r in REGIONS] + [hug_glob_row["dmdtda"]]
hug_dmdtda_err = [hug.loc[r, "err_dmdtda"] for r in REGIONS] + [hug_glob_row["err_dmdtda"]]
hug_pg_dmdtda = [hug_pergla.loc[r, "dmdtda"] for r in REGIONS] + [hug_pergla_global]

fig, ax = plt.subplots(figsize=(13, 5))
ax.bar(x, our_dmdtda, color=bar_colors, alpha=0.8, zorder=2, label="OGGM 2000–2019")
ax.plot(
    x, hug_pg_dmdtda,
    "s", color="0.5", ms=5, zorder=3, label="Hugonnet per-glacier (area-weighted, no extrapolation)",
)
ax.errorbar(
    x, hug_dmdtda, yerr=hug_dmdtda_err,
    fmt="o", color="k", ms=5, lw=1.2, capsize=3, zorder=4,
    label="Hugonnet et al. 2021 regional (with area extrapolation)",
)
ax.axhline(0, color="k", lw=0.5, ls="--")
ax.set_xticks(x)
ax.set_xticklabels(labels, rotation=45, ha="right")
ax.set_ylabel("Specific MB rate (m w.e. yr⁻¹)")
ax.set_title("Specific MB rate — OGGM fixed geometry vs Hugonnet et al. 2021 (2000–2019)")
ax.legend(fontsize=8)
ax.grid(axis="y", alpha=0.3, zorder=1)
fig.tight_layout()
fig.savefig("plot_dmdtda_by_region.png", dpi=150)
plt.show()