# %% [markdown]
# # Using the water-divide runs: catchments
#
# This notebook shows how to work with the new inventory (RGI2025-hydro v2) at the
# level of individual hydrological catchments, which is what matters for
# hydropower. It gives:
#
# 1. the structure of the inventory (which entity is what);
# 2. a summary table for all catchments (volume, mass loss, peak water);
# 3. a worked example for one aggregate catchment, Hálslón.
#
# All projections are CMIP6 to 2100 (median and 10th-90th percentile range across
# GCMs). The same runs also go to 2300 and exist for CMIP5: see `data/projections/`.

# %%
import geopandas as gpd
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import xarray as xr

# %%
data_dir = "data"
fig_dir = "figures"
table_dir = "tables"

# %% [markdown]
# ## The inventory
#
# The inventory `data/inventory/RGI2025hydro-v7.0-G-06_iceland.shp` follows the RGI7
# format, with a few extra columns:
#
# - `ent_type`: `catchment` (one of your water-divide catchments), `remnant` (ice not
#   covered by any catchment, modelled so that totals stay consistent, but not
#   assigned to a river), or `rgi2025` (unchanged glacier from the first delivery,
#   with the same `rgi_id`).
# - `ws_id` and `catchment`: the ID and name of your water-divide polygon.
# - `icecap`: the ice cap the entity belongs to.
#
# New entities have IDs from `RGI2025-v7.0-G-06-10001` onwards, so they never collide
# with RGI2025 IDs.

# %%
inv = gpd.read_file(f"{data_dir}/inventory/RGI2025hydro-v7.0-G-06_iceland.shp")
inv = inv.set_index("rgi_id")
inv.groupby("ent_type").area_km2.agg(["count", "sum"]).round(1)

# %% [markdown]
# Maps of the new entities, ice cap by ice cap, are in `figures/inventory_maps/`.
# Your original water divides are drawn as a red dashed line on top of the modelled
# entities.
#
# ![Hofsjökull](figures/inventory_maps/Hofsjokull.png)

# %% [markdown]
# ### From your catchments to model entities
#
# Most catchments correspond to exactly one model entity (the one with the same
# `ws_id`). Two of your catchments contain smaller ones and are modelled as sums
# (`data/inventory/catchment_lookup.csv`):
#
# - **Halslon10** (0035) = Kring + illakv + jokulsa + jokulkvisl;
# - **Jökulsá í Lóni** (0041) = Skyndidalsá + the rest of the Jökulsá í Lóni catchment.
#
# Þingvallavatn is modelled with the Þingvallavatn_2 delineation (0023), as you
# chose. It includes 43 km² that are also part of your Hagavatn polygon: in the
# model, this area belongs to Þingvallavatn only.

# %%
lookup = pd.read_csv(f"{data_dir}/inventory/catchment_lookup.csv", dtype={"ws_id": str})

catchment_ids = {}  # catchment name -> list of rgi_ids
for _, row in lookup.iterrows():
    if row.ws_id in ["0035", "0041"]:  # aggregates: sum of sub-catchments
        catchment_ids[row.catchment] = row.rgi_ids.split(";")
    elif row.ws_id in inv.ws_id.values:  # one entity per catchment
        catchment_ids[row.catchment] = list(inv.index[inv.ws_id == row.ws_id])
print(f"{len(catchment_ids)} catchments (incl. 2 aggregates)")
print("Hálslón:", catchment_ids["Halslon10"])

# %% [markdown]
# ## Load the runs

# %%
hist = xr.open_dataset(f"{data_dir}/historical/hydro_v2/run_output_spinup_historical.nc")
proj = xr.open_dataset(f"{data_dir}/projections/hydro_v2/merged_2100_projections_cmip6.nc")

# See notebook 01: four runs lose many small glaciers and are excluded here
exclude = ["CESM2-WACCM_ssp126", "CESM2-WACCM_ssp534-over",
           "MRI-ESM2-0_ssp126", "MRI-ESM2-0_ssp534-over"]
proj = proj.sel(gcm_scenario=[g for g in proj.gcm_scenario.values if g not in exclude])

ssps = ["ssp126", "ssp245", "ssp370", "ssp585"]
ssp_colors = {"ssp126": "steelblue", "ssp245": "gold",
              "ssp370": "darkorange", "ssp585": "firebrick"}
runoff_vars = ["melt_on_glacier", "melt_off_glacier",
               "liq_prcp_on_glacier", "liq_prcp_off_glacier"]


def ssp_runs(da, ssp):
    return da.sel(gcm_scenario=[g for g in da.gcm_scenario.values if g.endswith(f"_{ssp}")])


def catchment_series(ids):
    """Mass (% of 2020) and smoothed annual runoff (Gt/yr) for a list of entities."""
    h = hist.sel(rgi_id=ids)
    p = proj.sel(rgi_id=ids)
    mass_2020 = float(h.mass_kg.sum("rgi_id").sel(time=2020))
    runoff = sum(p[v].clip(min=0) for v in runoff_vars).sum("rgi_id") * 1e-12
    return {
        "volume_2020": float(h.volume.sum("rgi_id").sel(time=2020)) * 1e-9,
        "hist_mass_pct": h.mass_kg.sum("rgi_id") / mass_2020 * 100,
        "mass_pct": p.mass_kg.sum("rgi_id") / mass_2020 * 100,
        "runoff": runoff.rolling(time=11, center=True, min_periods=1).mean(),
    }


# %% [markdown]
# ## Summary table for all catchments
#
# Median across GCMs. "Peak water" is the year of maximum 11-yr smoothed runoff.
# For small catchments the GCM spread is large: use the full time series (next
# section) rather than single numbers when possible.

# %%
rows = []
for name, ids in catchment_ids.items():
    s = catchment_series(ids)
    row = {
        "Catchment": name,
        "Ice cap": inv.loc[ids[0], "icecap"],
        "Area 2025 (km²)": inv.loc[ids, "area_km2"].sum(),
        "Volume 2020 (km³)": s["volume_2020"],
    }
    for ssp in ["ssp245", "ssp585"]:
        mass = ssp_runs(s["mass_pct"], ssp)
        ro = ssp_runs(s["runoff"], ssp)
        row[f"Mass 2050 {ssp} (%)"] = float(mass.sel(time=2050).median())
        row[f"Mass 2100 {ssp} (%)"] = float(mass.sel(time=2100).median())
        row[f"Peak water {ssp}"] = float(ro.idxmax("time").median())
    rows.append(row)

summary = pd.DataFrame(rows).set_index("Catchment").sort_values(["Ice cap", "Area 2025 (km²)"])
summary.to_csv(f"{table_dir}/catchment_summary.csv")
summary.round(1)

# %% [markdown]
# The table is also saved as `tables/catchment_summary.csv`.
#
# ### Caveat: catchments in the accumulation area
#
# OGGM treats each catchment as a glacier with its own flowline. In reality, ice in
# the interior of an ice cap flows across water divides towards the outlet glaciers.
# A catchment that lies mostly in the accumulation area therefore has no realistic
# way to lose its ice in the model, and may even thicken. These catchments stand out
# in the table: they keep almost all of their mass until 2050, and some gain mass.

# %%
interior = summary[summary["Mass 2050 ssp245 (%)"] > 90]
interior[["Ice cap", "Area 2025 (km²)", "Mass 2050 ssp245 (%)",
          "Mass 2100 ssp245 (%)", "Peak water ssp245"]].round(1)

# %% [markdown]
# **grimsvotn** is the clearest case: it gains mass throughout the century under
# ssp245, which is not realistic. The other catchments in this list are the large
# northern Vatnajökull catchments, including all four that make up Hálslón (below).
# They extend high into the accumulation area, so a slower mass loss and a late peak
# water are partly expected, but part of it is likely this artefact, and the runs
# alone cannot separate the two.
#
# For these catchments, treat the mass and peak-water numbers with caution. Their
# runoff in the coming decades is more robust, because it is driven by the surface
# melt and rain in the catchment rather than by ice flow.

# %% [markdown]
# ## Example: Hálslón
#
# The Hálslón catchment (Halslon10) is the sum of four model entities. Any other
# catchment, or any combination of entities, works the same way: replace `ids`.

# %%
ids = catchment_ids["Halslon10"]
s = catchment_series(ids)
print(f"Area {inv.loc[ids, 'area_km2'].sum():.0f} km², volume 2020 {s['volume_2020']:.0f} km³")

fig, axes = plt.subplots(1, 2, figsize=(14, 5))
hs = s["hist_mass_pct"].sel(time=slice(2000, 2026))
axes[0].plot(hs.time, hs, color="black", lw=1.5, label="Historical (CARRA)")
for ssp in ssps:
    for ax, var in zip(axes, ["mass_pct", "runoff"]):
        da = ssp_runs(s[var], ssp)
        c = ssp_colors[ssp]
        ax.plot(da.time, da.median("gcm_scenario"), color=c,
                label=f"{ssp} (n={da.gcm_scenario.size})")
        ax.fill_between(da.time, da.quantile(0.1, "gcm_scenario"),
                        da.quantile(0.9, "gcm_scenario"), color=c, alpha=0.2, lw=0)
axes[0].axhline(100, color="gray", lw=0.7, ls=":")
axes[0].set_ylabel("Mass (% of 2020)")
axes[1].set_ylabel("Annual runoff (Gt yr$^{-1}$, 11-yr smoothing)")
for ax in axes:
    ax.set_xlabel("Year")
    ax.grid(True, alpha=0.3)
    ax.legend(fontsize=9)
fig.suptitle("Hálslón catchment (Kring + illakv + jokulsa + jokulkvisl)", fontsize=13)
plt.tight_layout()
fig.savefig(f"{fig_dir}/halslon_mass_runoff.png", dpi=150, bbox_inches="tight")
plt.show()

# %% [markdown]
# ### Seasonal cycle of runoff
#
# The projection files also store monthly runoff (variables ending in `_monthly`,
# in kg per month). Here: the mean seasonal cycle of Hálslón runoff for three
# decades, under ssp245 (median across GCMs). Runoff is converted to an average
# discharge in m³ s⁻¹.

# %%
ssp = "ssp245"
p = ssp_runs(proj.sel(rgi_id=ids), ssp)
monthly = sum(p[f"{v}_monthly"].clip(min=0) for v in runoff_vars).sum("rgi_id")
seconds_per_month = 365 * 24 * 3600 / 12  # the model uses a 365-day calendar
discharge = monthly / 1000 / seconds_per_month  # kg/month -> m³/s

decades = {"2026-2035": "lightsteelblue", "2056-2065": "steelblue", "2091-2100": "navy"}
fig, ax = plt.subplots(figsize=(9, 5))
for (period, c) in decades.items():
    y0, y1 = [int(y) for y in period.split("-")]
    cycle = discharge.sel(time=slice(y0, y1)).mean("time").median("gcm_scenario")
    ax.plot(p.calendar_month_2d, cycle, color=c, marker="o", label=period)
ax.set_xticks(range(1, 13))
ax.set_xticklabels(["J", "F", "M", "A", "M", "J", "J", "A", "S", "O", "N", "D"])
ax.set_ylabel("Mean discharge (m³ s$^{-1}$)")
ax.set_xlabel("Month")
ax.grid(True, alpha=0.3)
ax.legend()
ax.set_title(f"Hálslón catchment: seasonal cycle of glacier runoff ({ssp}, median)")
plt.tight_layout()
fig.savefig(f"{fig_dir}/halslon_seasonal_cycle.png", dpi=150, bbox_inches="tight")
plt.show()

# %% [markdown]
# Note: as in the first delivery, runoff is computed over a fixed area (the glacier
# outline at the start of the projections). It therefore includes rain and snowmelt
# on terrain that becomes ice-free, but not the rest of the river catchment.