# %% [markdown]
# # Iceland OGGM — Experiment Selection & Validation
#
# Before running the CMIP5/CMIP6 projections, we calibrated OGGM against several
# independent observations (glacier area 2010/2025, Hugonnet et al. 2000-2020 mass
# change, DEM-differenced volume 2021, and Andri's mass-balance maps).
# This was done across a sweep of configuration choices:
#
# - **climate forcing**: ERA5 vs. CARRA reanalysis (1991-2026)
# - **flowline bed shape**: trapezoidal, with the shape parameter lambda ranging from
#   2 (~45° angle of side wall, measured from horizontal) to 16 (~7° angle, much flatter
#   than 45°)
# - **mass-balance models**: monthly temperature-index model with surface type tracking
#   (default) and without tracking (Monthly TI)
#
# **Outcome: CARRA, trapezoidal lambda = 14 was chosen as the delivered configuration.**
# For ERA5, lambda = 8 was the closest contender and is shown throughout for comparison —
# the difference between the two is small. The Monthly TI variant (without snow tracking)
# of each was also tested and is included alongside its parent configuration.
#
# **The choice is a compromise, not a perfect fit.** The lambda sweep (Appendix) shows a
# systematic trade-off: increasing lambda improves the area match (particularly the 2010
# outline, an independent check — 2025 is matched by construction, since it is the
# spinup's initialisation target) but pushes the modelled 2000-2020 volume loss further
# from the Hugonnet estimate; decreasing lambda does the opposite. No single lambda
# simultaneously minimises both mismatches. CARRA lambda=14 and ERA5 lambda=8 are the
# points in the sweep that balance the two best — not the configurations that match
# either observation alone most closely.
#
# This notebook first summarises that decision (see *Decision summary* below), then walks
# through the validation evidence for the shortlisted configurations. The full lambda
# sweep that led to the shortlist is kept in the **Appendix** at the end, for transparency
# and reproducibility — it is not part of the main narrative.
#
# **On data availability:** to keep the delivery package simple and clear, only the
# NetCDF output of the delivered configuration (and, for comparison, the ERA5
# alternative) ships alongside this notebook — not the full ~15-configuration sweep shown
# in the Appendix. That full sweep's data can be made available on request. As shipped,
# this notebook reads from `/home/shared/fmaussion_pschmitt/iceland_projections` on the
# OGGM cluster and will **not run as-is** without access to that path — it is included to
# document the calibration process and decision, not as an executable artefact for the
# recipient.

# %%
import os
from IPython.display import HTML

import numpy as np
import pandas as pd
import xarray as xr
import yaml
import matplotlib.pyplot as plt

# %% [markdown]
# ## Load model output and observations
#
# `CONFIGS` lists every configuration in the calibration sweep (label -> working
# directory suffix). `SHORTLIST` is the subset shown in the main narrative below; the
# rest only reappears in the Appendix.

# %%
project_dir = "/home/shared/fmaussion_pschmitt/iceland_projections"
workflow_dir = f"{project_dir}/complete_workflow"
validation_data_dir = f"{project_dir}/validation/validation_data"

fig_dir = "figures"  # relative to packaging_delivery/, kept self-contained for delivery
os.makedirs(fig_dir, exist_ok=True)

CONFIGS = {
    "era5_2":             "ERA5, trapez. λ=2 (~45°)",
    "era5_4":             "ERA5, trapez. λ=4 (~26°)",
    "era5_8":             "ERA5, trapez. λ=8 (~14°)",
    "era5_14":            "ERA5, trapez. λ=14 (~8°)",
    "era5_16":            "ERA5, trapez. λ=16 (~7°)",
    "era5_8_MonthlyTI":   "ERA5, trapez. λ=8 (~14°), Monthly TI",
    "era5_14_MonthlyTI":  "ERA5, trapez. λ=14 (~8°), Monthly TI",
    "carra_2":            "CARRA, trapez. λ=2 (~45°)",
    "carra_4":            "CARRA, trapez. λ=4 (~26°)",
    "carra_8":            "CARRA, trapez. λ=8 (~14°)",
    "carra_12":           "CARRA, trapez. λ=12 (~9°)",
    "carra_14":           "CARRA, trapez. λ=14 (~8°)",
    "carra_16":           "CARRA, trapez. λ=16 (~7°)",
    "carra_8_MonthlyTI":  "CARRA, trapez. λ=8 (~14°), Monthly TI",
    "carra_14_MonthlyTI": "CARRA, trapez. λ=14 (~8°), Monthly TI",
}

# Delivered configuration (carra_14) + closest contender (era5_8), each with its
# Monthly TI variant. Everything else in CONFIGS is sweep-only, shown in the Appendix.
SHORTLIST = ["carra_14", "carra_14_MonthlyTI", "era5_8", "era5_8_MonthlyTI"]

models = {
    key: xr.open_dataset(
        os.path.join(workflow_dir, f"working_dir_adapted_{key}",
                     "run_output_spinup_historical.nc"))
    for key in CONFIGS
}

# area and volume-change observations, provided by Andri
with open(os.path.join(validation_data_dir, "area_dV_obs.yaml")) as f:
    area_dv_obs = yaml.safe_load(f)

# specific mass-balance observations derived from Andri's SMB maps
ds_smb = xr.open_dataset(os.path.join(validation_data_dir, "smbs_from_maps.nc"))
ds_smb["time_int"] = ds_smb.time_a.dt.year

# %% [markdown]
# ## Helper functions
#
# These operate on the `models` dict and are reused for both the shortlist plots below
# and the full-sweep plots in the Appendix, so the two only differ in which keys they
# are called with.

# %%
rgi_ids_volume_ddem = area_dv_obs["rgi_ids_volume_change_2010_2021"]


def plot_area(ax, keys):
    for key in keys:
        (models[key].sum(dim="rgi_id").area / 1e6).plot(ax=ax, label=CONFIGS[key])
    ax.plot(2010, area_dv_obs["area_2010_km2"], "o", color="k", label="Outline area ~2010")
    ax.plot(2025, area_dv_obs["area_2025_km2"], "o", color="0.4", label="Outline area ~2025")
    ax.set_ylabel("Area (km²)")
    ax.grid(True)
    ax.legend(bbox_to_anchor=[1.01, 0.5], loc="center left")


def plot_mass_change(ax, keys):
    for key in keys:
        ds = models[key].sum(dim="rgi_id")
        (ds.mass_kg - ds.mass_kg.sel(time=2000)).plot(ax=ax, label=CONFIGS[key])
    var, err_var = "hugonnet_adapted_2000_2020_kg", "hugonnet_2000-01-01_2020-01-01_kg_err"
    ax.plot([2000, 2020], [0, area_dv_obs[var]], c="k")
    ax.plot([2020, 2020],
            [area_dv_obs[var] - area_dv_obs[err_var], area_dv_obs[var] + area_dv_obs[err_var]],
            c="k", label="Hugonnet (adapted) 2000-2020 ± error")
    ax.set_ylabel("Mass change since 2000 (kg)")
    ax.grid(True)
    ax.legend(bbox_to_anchor=[1.01, 0.5], loc="center left")


def plot_volume_bed(ax, keys):
    for key in keys:
        models[key].sel(rgi_id=rgi_ids_volume_ddem).volume.sum(dim="rgi_id").plot(
            ax=ax, label=CONFIGS[key])
    ax.plot(2010, area_dv_obs["volume_2010"], "o", color="k", label="Volume 2010 (DEM - bed)")
    ax.plot(2021, area_dv_obs["volume_2021"], "o", color="0.4", label="Volume 2021 (DEM - bed)")
    ax.set_ylabel("Volume (m³)")
    ax.grid(True)
    ax.legend(bbox_to_anchor=[1.01, 0.5], loc="center left")


def plot_volume_change_ddem(ax, keys):
    for key in keys:
        ds_vol = models[key].sel(rgi_id=rgi_ids_volume_ddem).volume.sum(dim="rgi_id")
        (ds_vol - ds_vol.sel(time=2010)).plot(ax=ax, label=CONFIGS[key])
    ax.plot([2010, 2021], [0, area_dv_obs["volume_change_2010_2021"]], c="k",
            label="Volume change from DEM differencing")
    ax.set_ylabel("Volume change since 2010 (m³)")
    ax.grid(True)
    ax.legend(bbox_to_anchor=[1.01, 0.5], loc="center left")


def regional_smb(ds):
    """Area-summed specific mass balance (mm w.e. yr-1) for the glaciers with SMB maps."""
    ds = ds.sel(rgi_id=ds_smb.rgi_id).sum(dim="rgi_id")
    return ds.volume.diff(dim="time", label="lower").reindex(time=ds.time) / ds.area * 900


# Area weights for the regionally-aggregated observed SMB. Using the delivered
# CARRA λ=14 config's glacier areas as the reference; areas barely differ between
# configs at this stage, so the choice of reference has a negligible effect.
regional_smb_weights = models["carra_14"].sel(
    rgi_id=ds_smb.rgi_id, time=ds_smb.time_int).area.fillna(0)
regional_smb_obs = ds_smb.weighted(regional_smb_weights).mean(dim="rgi_id")["ba"] * 1000


def plot_regional_smb(ax, keys):
    for key in keys:
        regional_smb(models[key]).plot(ax=ax, label=CONFIGS[key])
    ax.plot(ds_smb.time_int, regional_smb_obs, c="k", label="Regionally aggregated obs.")
    ax.set_ylabel("Specific MB (mm w.e. yr$^{-1}$)")
    ax.grid(True)
    ax.legend(bbox_to_anchor=[1.01, 0.5], loc="center left")


def smb_stats(model_series, obs_series):
    """Bias, RMSD and correlation of a modelled SMB series against the SMB-map obs."""
    aligned = model_series.sel(time=ds_smb.time_int)
    return {
        "bias": float(np.mean(aligned) - np.mean(obs_series)),
        "rmsd": float(np.sqrt(np.mean((aligned.values - obs_series.values) ** 2))),
        "corr": float(np.corrcoef(aligned.values, obs_series.values)[0, 1]),
    }


def plot_single_glacier_smb(rgi_id, keys):
    fig, ax = plt.subplots()
    for key in keys:
        ds = models[key].sel(rgi_id=rgi_id)
        smb = ds.volume.diff(dim="time", label="lower").reindex(time=ds.time) / ds.area * 900
        smb.plot(ax=ax, label=CONFIGS[key])
    ax.plot(ds_smb.time_int, ds_smb["ba"].sel(rgi_id=rgi_id) * 1000, c="k", label="Observations")
    ax.set_title(str(rgi_id.values))
    ax.set_ylabel("Specific MB (mm w.e. yr$^{-1}$)")
    ax.grid(True)
    ax.legend()
    plt.show()

# %% [markdown]
# ## Decision summary
#
# For each shortlisted configuration: the 2010 area mismatch against the outline
# observations (the independent check — 2025 is matched by construction, see note above,
# so it is not informative here), the 2000-2020 volume-change mismatch against Hugonnet et
# al. 2021 (adapted to the 2025 outlines, see `Landsvirkjun consultancy.md`), and the regional
# SMB fit against Andri's SMB maps. `hugonnet_mismatch_pct` is positive when the model
# loses *more* mass than observed, negative when it loses *less*.

# %%
def config_summary(key):
    ds = models[key].sum(dim="rgi_id")

    # 2025 area is matched by construction (spinup initialisation target); 2010 is the
    # actual independent area check, and the one that trades off against the Hugonnet
    # volume-change match below.
    area_2010 = float(ds.area.sel(time=2010)) / 1e6
    area_2010_mismatch_pct = (
        (area_2010 - area_dv_obs["area_2010_km2"]) / area_dv_obs["area_2010_km2"] * 100)

    # Matching Hugonnets observed mass change
    mass_change_kg = float(ds.mass_kg.sel(time=2020) - ds.mass_kg.sel(time=2000))
    obs_change_kg = area_dv_obs["hugonnet_adapted_2000_2020_kg"]
    hugonnet_mismatch_pct = (mass_change_kg - obs_change_kg) / obs_change_kg * 100

    stats = smb_stats(regional_smb(models[key]), regional_smb_obs)

    return {
        "config": CONFIGS[key],
        "area_2010_km2": area_2010,
        "area_2010_mismatch_pct": area_2010_mismatch_pct,
        "mass_change_2000_2020_Gt": mass_change_kg / 1e12,
        "hugonnet_mismatch_pct": hugonnet_mismatch_pct,
        "regional_smb_bias_mmwe": stats["bias"],
        "regional_smb_rmsd_mmwe": stats["rmsd"],
        "regional_smb_corr": stats["corr"],
    }


decision_table = pd.DataFrame(
    [config_summary(key) for key in SHORTLIST]
).set_index("config").round(2)


def scrollable(df, height=500):
    return HTML(f'<div style="max-height:{height}px; overflow:auto">{df.to_html()}</div>')


scrollable(decision_table)

# %% [markdown]
# **Verdict:** across the full sweep (Appendix), a higher lambda tends to shrink the 2010
# area mismatch but push `hugonnet_mismatch_pct` further positive (more volume loss than
# observed); a lower lambda does the reverse — no single value zeroes both. CARRA λ=14 and
# ERA5 λ=8 are the sweep's best compromise points, matching the Hugonnet volume-change
# record to within ~10% (which corresponds to the uncertainty of the observation) while keeping the 2010 area mismatch under ~2%, with a competitive
# regional SMB fit.

# %% [markdown]
# # Validation for the entire region

# %% [markdown]
# ## Area change from outlines 2010 and 2025

# %%
fig, ax = plt.subplots()
plot_area(ax, SHORTLIST)
ax.set_title("Area evolution — shortlisted configurations")
fig.savefig(f"{fig_dir}/validation_area.png", dpi=150, bbox_inches="tight")

# %% [markdown]
# ## Mass change from Hugonnet

# %%
fig, ax = plt.subplots()
plot_mass_change(ax, SHORTLIST)
ax.set_title("Mass evolution relative to the year 2000 — shortlisted configurations")
fig.savefig(f"{fig_dir}/validation_mass_change.png", dpi=150, bbox_inches="tight")

# %% [markdown]
# # Validation for sub-regions
#
# The following comparisons are restricted to the glaciers for which Andri also provided
# bed topography maps (68 glaciers, `rgi_ids_volume_change_2010_2021`), enabling an
# independent volume/volume-change check beyond the region-wide Hugonnet comparison above.

# %% [markdown]
# ## Total volume from bed maps

# %%
fig, ax = plt.subplots()
plot_volume_bed(ax, SHORTLIST)
ax.set_title("Total volume for glaciers with DEM and bed observations")
fig.savefig(f"{fig_dir}/validation_volume_bed.png", dpi=150, bbox_inches="tight")

# %% [markdown]
# ## Volume change from DEM differencing

# %%
fig, ax = plt.subplots()
plot_volume_change_ddem(ax, SHORTLIST)
ax.set_title("Volume evolution relative to 2010\n(glaciers with DEM and bed observations)")
fig.savefig(f"{fig_dir}/validation_volume_change_ddem.png", dpi=150, bbox_inches="tight")

# %% [markdown]
# ## SMB comparison
#
# ### Regional SMB

# %%
fig, ax = plt.subplots()
plot_regional_smb(ax, SHORTLIST)
ax.set_title("Regionally-aggregated specific mass balance — shortlisted configurations")
fig.savefig(f"{fig_dir}/validation_regional_smb.png", dpi=150, bbox_inches="tight")

# %%
for key in SHORTLIST:
    stats = smb_stats(regional_smb(models[key]), regional_smb_obs)
    print(f"{CONFIGS[key]:45s} bias={stats['bias']:6.1f}  rmsd={stats['rmsd']:6.1f}  "
          f"corr={stats['corr']:.3f}  (mm w.e. yr-1)")

# %% [markdown]
# ### Per-glacier SMB
#
# Same comparison for the four individual glaciers larger than 900 km². Shown for the
# shortlisted configurations only — the full lambda sweep is not repeated here to avoid
# clutter; see the regional comparison and Appendix below for the sweep-wide picture.

# %%
large_glacier_ids = [
    rgi_id for rgi_id in ds_smb.rgi_id
    if models["carra_14"].sel(time=2010, rgi_id=rgi_id).area > 9e8  # > 900 km2
]
for rgi_id in large_glacier_ids:
    plot_single_glacier_smb(rgi_id, SHORTLIST)

# %% [markdown]
# # Appendix: full calibration sweep
#
# The plots below repeat the region-wide and sub-regional comparisons above, but across
# every configuration in the sweep (`CONFIGS`), not just the shortlist. This is the full
# evidence base behind the decision summarised at the top of this notebook — kept here for
# transparency and reproducibility rather than as part of the main narrative.

# %% [markdown]
# ## Area change — full sweep

# %%
fig, ax = plt.subplots(figsize=(8, 5))
plot_area(ax, CONFIGS.keys())
ax.set_title("Area evolution — full calibration sweep")

# %% [markdown]
# ## Mass change from Hugonnet — full sweep

# %%
fig, ax = plt.subplots(figsize=(8, 5))
plot_mass_change(ax, CONFIGS.keys())
ax.set_title("Mass evolution relative to the year 2000 — full calibration sweep")

# %% [markdown]
# ## Total volume from bed maps — full sweep

# %%
fig, ax = plt.subplots(figsize=(8, 5))
plot_volume_bed(ax, CONFIGS.keys())
ax.set_title("Total volume for glaciers with DEM and bed observations — full sweep")

# %% [markdown]
# ## Volume change from DEM differencing — full sweep

# %%
fig, ax = plt.subplots(figsize=(8, 5))
plot_volume_change_ddem(ax, CONFIGS.keys())
ax.set_title("Volume evolution relative to 2010 — full calibration sweep")

# %% [markdown]
# ## Regional SMB — full sweep

# %%
fig, ax = plt.subplots(figsize=(8, 5))
plot_regional_smb(ax, CONFIGS.keys())
ax.set_title("Regionally-aggregated specific mass balance — full calibration sweep")

# %%
for key in CONFIGS:
    stats = smb_stats(regional_smb(models[key]), regional_smb_obs)
    print(f"{CONFIGS[key]:45s} bias={stats['bias']:6.1f}  rmsd={stats['rmsd']:6.1f}  "
          f"corr={stats['corr']:.3f}  (mm w.e. yr-1)")