# %% [markdown]
# # Iceland Glacier Volume Projections
#
# This notebook plots total Iceland glacier volume from 2000 to 2100 (and 2300),
# combining the CARRA-forced historical run (2000-2026, black line) with the CMIP5 and
# CMIP6 ensemble projections (2026 onwards). All volumes are expressed as a percentage
# of the 2020 volume from the CARRA historical run.
#
# Shaded bands show the 10th-90th percentile range across GCMs; solid lines show
# the median.
#
# This notebook ships with its own copy of all the data it needs, under `data/` and
# `03_amoc_index/` in this same delivery folder — it does not depend on the OGGM
# cluster and can be run standalone.

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

# %% [markdown]
# ## Load data

# %%
data_dir  = "data"
fig_dir   = "figures"
os.makedirs(fig_dir, exist_ok=True)

proj_dir = f"{data_dir}/projections/projections_output_carra"
hist_path = f"{data_dir}/historical/working_dir_adapted_carra_14/run_output_spinup_historical.nc"

# projection ensembles
ds_cmip6      = xr.open_dataset(f"{proj_dir}/merged_2100_projections_cmip6.nc")
ds_cmip5      = xr.open_dataset(f"{proj_dir}/merged_2100_projections_cmip5.nc")
ds_cmip6_2300 = xr.open_dataset(f"{proj_dir}/merged_2300_projections_cmip6.nc")
ds_cmip5_2300 = xr.open_dataset(f"{proj_dir}/merged_2300_projections_cmip5.nc")

# historical run (CARRA-forced, single trajectory, integer year index, 1991-2026)
ds_hist = xr.open_dataset(hist_path)

# %% [markdown]
# ## Normalise to 2020
#
# All volumes are expressed as % of the 2020 total Iceland volume from the CARRA
# historical run. This removes the absolute volume baseline and makes it easy to read
# the fraction of ice remaining.

# %%
vol_hist = ds_hist.volume.sum(dim="rgi_id") / 1e9   # km³, integer year index
vol_2020 = float(vol_hist.sel(time=2020))            # reference value
print(f"2020 total Iceland glacier volume (CARRA): {vol_2020:.0f} km³")

# Historical 2000-2026 as % of 2020. Include 2026 so the line meets the
# projection ensemble at their shared starting point (no gap).
vol_hist_plot = vol_hist.sel(time=slice(2000, 2026)) / vol_2020 * 100

# Projected volumes (% of 2020) — float year index starting at 2026
vol_cmip6      = ds_cmip6["volume"].sum(dim="rgi_id")      / 1e9 / vol_2020 * 100
vol_cmip5      = ds_cmip5["volume"].sum(dim="rgi_id")      / 1e9 / vol_2020 * 100
vol_cmip6_2300 = ds_cmip6_2300["volume"].sum(dim="rgi_id") / 1e9 / vol_2020 * 100
vol_cmip5_2300 = ds_cmip5_2300["volume"].sum(dim="rgi_id") / 1e9 / vol_2020 * 100

# %% [markdown]
# ## Scenario colours
#
# Used throughout the notebook for the CMIP6 (SSP) and CMIP5 (RCP) scenarios.

# %%
scenario_colors = {
    "ssp126":      "steelblue",
    "ssp245":      "gold",
    "ssp370":      "darkorange",
    "ssp534-over": "mediumpurple",
    "ssp585":      "firebrick",
    "rcp26":       "steelblue",
    "rcp45":       "gold",
    "rcp60":       "darkorange",
    "rcp85":       "firebrick",
}

# %% [markdown]
# ## Projections to 2100

# %%
ssp_scenarios = ["ssp126", "ssp245", "ssp370", "ssp585"]
rcp_scenarios = ["rcp26", "rcp45", "rcp60", "rcp85"]

fig, axes = plt.subplots(1, 2, figsize=(14, 5), sharey=True)

# --- left panel: CMIP6 (SSP scenarios) ---
ax = axes[0]
ax.plot(vol_hist_plot.time, vol_hist_plot, color="black", lw=1.5, label="Historical (CARRA)")
for scen in ssp_scenarios:
    mask = np.array([s.endswith(f"_{scen}") for s in vol_cmip6.gcm_scenario.values])
    if mask.sum() == 0:
        continue
    vol_scen = vol_cmip6.isel(gcm_scenario=mask)
    med = vol_scen.quantile(0.50, dim="gcm_scenario")
    p10 = vol_scen.quantile(0.10, dim="gcm_scenario")
    p90 = vol_scen.quantile(0.90, dim="gcm_scenario")
    c = scenario_colors[scen]
    ax.plot(vol_cmip6.time, med, color=c, label=f"{scen}  (n={mask.sum()})")
    ax.fill_between(vol_cmip6.time, p10, p90, color=c, alpha=0.2)
ax.axhline(100, color="gray", lw=0.7, ls=":")
ax.set_xlabel("Year")
ax.set_ylabel("Volume (% of 2020)")
ax.set_title("CMIP6")
ax.legend(fontsize=9)
ax.grid(True, alpha=0.3)

# --- right panel: CMIP5 (RCP scenarios) ---
ax = axes[1]
ax.plot(vol_hist_plot.time, vol_hist_plot, color="black", lw=1.5, label="Historical (CARRA)")
for scen in rcp_scenarios:
    mask = np.array([s.endswith(f"_{scen}") for s in vol_cmip5.gcm_scenario.values])
    if mask.sum() == 0:
        continue
    vol_scen = vol_cmip5.isel(gcm_scenario=mask)
    med = vol_scen.quantile(0.50, dim="gcm_scenario")
    p10 = vol_scen.quantile(0.10, dim="gcm_scenario")
    p90 = vol_scen.quantile(0.90, dim="gcm_scenario")
    c = scenario_colors[scen]
    ax.plot(vol_cmip5.time, med, color=c, label=f"{scen}  (n={mask.sum()})")
    ax.fill_between(vol_cmip5.time, p10, p90, color=c, alpha=0.2)
ax.axhline(100, color="gray", lw=0.7, ls=":")
ax.set_xlabel("Year")
ax.set_title("CMIP5")
ax.legend(fontsize=9)
ax.grid(True, alpha=0.3)

fig.suptitle("Iceland total glacier volume (2000-2100)", fontsize=13)
plt.tight_layout()
fig.savefig(f"{fig_dir}/volume_projections_2100.png", dpi=150, bbox_inches="tight")
plt.show()

# %% [markdown]
# ## Projections to 2300
#
# Fewer models run to 2300. CMIP6 covers ssp126, ssp534-over, and ssp585 only.

# %%
ssp_scenarios_2300 = ["ssp126", "ssp534-over", "ssp585"]
rcp_scenarios_2300 = ["rcp26", "rcp45", "rcp60", "rcp85"]

fig, axes = plt.subplots(1, 2, figsize=(14, 5), sharey=True)

# --- left panel: CMIP6 (SSP scenarios) ---
ax = axes[0]
ax.plot(vol_hist_plot.time, vol_hist_plot, color="black", lw=1.5, label="Historical (CARRA)")
for scen in ssp_scenarios_2300:
    mask = np.array([s.endswith(f"_{scen}") for s in vol_cmip6_2300.gcm_scenario.values])
    if mask.sum() == 0:
        continue
    vol_scen = vol_cmip6_2300.isel(gcm_scenario=mask)
    med = vol_scen.quantile(0.50, dim="gcm_scenario")
    p10 = vol_scen.quantile(0.10, dim="gcm_scenario")
    p90 = vol_scen.quantile(0.90, dim="gcm_scenario")
    c = scenario_colors[scen]
    ax.plot(vol_cmip6_2300.time, med, color=c, label=f"{scen}  (n={mask.sum()})")
    ax.fill_between(vol_cmip6_2300.time, p10, p90, color=c, alpha=0.2)
ax.axhline(100, color="gray", lw=0.7, ls=":")
ax.set_xlabel("Year")
ax.set_ylabel("Volume (% of 2020)")
ax.set_title("CMIP6")
ax.legend(fontsize=9)
ax.grid(True, alpha=0.3)

# --- right panel: CMIP5 (RCP scenarios) ---
ax = axes[1]
ax.plot(vol_hist_plot.time, vol_hist_plot, color="black", lw=1.5, label="Historical (CARRA)")
for scen in rcp_scenarios_2300:
    mask = np.array([s.endswith(f"_{scen}") for s in vol_cmip5_2300.gcm_scenario.values])
    if mask.sum() == 0:
        continue
    vol_scen = vol_cmip5_2300.isel(gcm_scenario=mask)
    med = vol_scen.quantile(0.50, dim="gcm_scenario")
    p10 = vol_scen.quantile(0.10, dim="gcm_scenario")
    p90 = vol_scen.quantile(0.90, dim="gcm_scenario")
    c = scenario_colors[scen]
    ax.plot(vol_cmip5_2300.time, med, color=c, label=f"{scen}  (n={mask.sum()})")
    ax.fill_between(vol_cmip5_2300.time, p10, p90, color=c, alpha=0.2)
ax.axhline(100, color="gray", lw=0.7, ls=":")
ax.set_xlabel("Year")
ax.set_title("CMIP5")
ax.legend(fontsize=9)
ax.grid(True, alpha=0.3)

fig.suptitle("Iceland total glacier volume (2000-2300)", fontsize=13)
plt.tight_layout()
fig.savefig(f"{fig_dir}/volume_projections_2300.png", dpi=150, bbox_inches="tight")
plt.show()

# %% [markdown]
# ## Projections by global temperature level
#
# Runs from CMIP5 and CMIP6 are pooled and grouped by end-of-century global mean
# temperature change relative to pre-industrial (2071-2100 average, IPCC AR6
# definition). Temperature levels classification data from Lilian Schuster / OGGM standard projections.

# %%
temp_dir = f"{data_dir}/supporting/temperature_levels"
temp_col = "global_temp_ch_2071-2100_preindustrial"

df_temp_cmip6 = pd.read_csv(
    f"{temp_dir}/Global_mean_temp_deviation_2071_2100_2081_2100_2281_2300_rel_1850_1900_cmip6_gcms.csv",
    index_col=0)
df_temp_cmip5 = pd.read_csv(
    f"{temp_dir}/Global_mean_temp_deviation_2071_2100_2081_2100_2281_2300_rel_1850_1900_cmip5_gcms.csv",
    index_col=0)

# temperature level definitions (from Patrick Schmitt)
temp_levels = ["+1.5°C", "+2.0°C", "+3.0°C", "+4.0°C"]
temp_lower  = [1.2,  1.75, 2.5, 3.5]
temp_upper  = [1.75, 2.5,  3.5, 4.5]
temp_colors = {
    "+1.5°C": "steelblue",
    "+2.0°C": "gold",
    "+3.0°C": "darkorange",
    "+4.0°C": "firebrick",
}

# %%
proj_time = vol_cmip6.time.values  # float, 2026-2100, same for CMIP5

fig, ax = plt.subplots(figsize=(10, 5))
ax.plot(vol_hist_plot.time, vol_hist_plot, color="black", lw=1.5, label="Historical (CARRA)")

for level, lo, hi in zip(temp_levels, temp_lower, temp_upper):
    runs = []
    for vol, df_temp in [(vol_cmip6, df_temp_cmip6), (vol_cmip5, df_temp_cmip5)]:
        for i, gcm_scen in enumerate(vol.gcm_scenario.values):
            if gcm_scen not in df_temp.index:
                continue
            t = df_temp.loc[gcm_scen, temp_col]
            if pd.notna(t) and lo <= t < hi:
                runs.append(vol.isel(gcm_scenario=i).values)

    if not runs:
        continue
    arr = np.array(runs)  # (n_runs, n_time)
    med = np.median(arr, axis=0)
    p10 = np.percentile(arr, 10, axis=0)
    p90 = np.percentile(arr, 90, axis=0)
    c = temp_colors[level]
    ax.plot(proj_time, med, color=c, label=f"{level}  (n={len(runs)})")
    ax.fill_between(proj_time, p10, p90, color=c, alpha=0.2)

ax.axhline(100, color="gray", lw=0.7, ls=":")
ax.set_xlabel("Year")
ax.set_ylabel("Volume (% of 2020)")
ax.legend(fontsize=9)
ax.grid(True, alpha=0.3)
fig.suptitle("Iceland total glacier volume by global temperature level (2000-2100)",
             fontsize=13)
plt.tight_layout()
fig.savefig(f"{fig_dir}/volume_projections_temp_levels.png", dpi=150, bbox_inches="tight")
plt.show()

# %% [markdown]
# ## Projections by AMOC index (high/low delta_I)
#
# Classify runs by the computed AMOC index groups (high_index vs low_index) based on
# delta_I (change in warming hole from 2005-2014 to 2090-2100). See
# `03_amoc_index/amoc_index_table_README.txt` for the full methodology.

# %%
df_amoc = pd.read_csv("03_amoc_index/amoc_index_table.csv")

# Map AMOC groups to gcm_scenario coordinates
# gcm_scenario format: "MODEL_SCENARIO" (e.g., "BCC-CSM2-MR_ssp126")
amoc_map = {}
for _, row in df_amoc.iterrows():
    key = f"{row['model']}_{row['experiment']}"
    amoc_map[key] = row['amoc_group']

amoc_colors = {
    "low_index":     "firebrick",
    "high_index":    "steelblue",
    "neutral_index": "lightgray",
}
amoc_labels = {
    "low_index":     "Strong AMOC decline",
    "high_index":    "Weak AMOC decline",
    "neutral_index": "Intermediate",
}

# %%
proj_time = vol_cmip6.time.values

fig, ax = plt.subplots(figsize=(10, 5))
ax.plot(vol_hist_plot.time, vol_hist_plot, color="black", lw=1.5, label="Historical (CARRA)")

for group in ["low_index", "high_index", "neutral_index"]:
    runs = []
    for i, gcm_scen in enumerate(vol_cmip6.gcm_scenario.values):
        if gcm_scen in amoc_map and amoc_map[gcm_scen] == group:
            runs.append(vol_cmip6.isel(gcm_scenario=i).values)

    if not runs:
        continue
    arr = np.array(runs)  # (n_runs, n_time)
    med = np.median(arr, axis=0)
    p10 = np.percentile(arr, 10, axis=0)
    p90 = np.percentile(arr, 90, axis=0)
    c = amoc_colors[group]
    lw = 1.5 if group != "neutral_index" else 1
    alpha = 0.2 if group != "neutral_index" else 0.1
    ax.plot(proj_time, med, color=c, label=f"{amoc_labels[group]}  (n={len(runs)})", lw=lw)
    ax.fill_between(proj_time, p10, p90, color=c, alpha=alpha)

ax.axhline(100, color="gray", lw=0.7, ls=":")
ax.set_xlabel("Year")
ax.set_ylabel("Volume (% of 2020)")
ax.legend(fontsize=10)
ax.grid(True, alpha=0.3)
fig.suptitle("Iceland total glacier volume by AMOC index (2000-2100)", fontsize=13)
plt.tight_layout()
fig.savefig(f"{fig_dir}/volume_projections_amoc_index.png", dpi=150, bbox_inches="tight")
plt.show()

# %% [markdown]
# ## Projections by Weijer AMOC streamfunction (external validation)
#
# Runs classified by the Weijer et al. (2020) AMOC streamfunction groups (strong_amoc vs
# weak_amoc), i.e. not dependent on our self-made temperature classification. This covers 12 of 19 models; the rest are excluded.

# %%
weijer_map = {}
for _, row in df_amoc.iterrows():
    if pd.notna(row['amoc_group_weijer']) and row['amoc_group_weijer'] != "unknown":
        key = f"{row['model']}_{row['experiment']}"
        weijer_map[key] = row['amoc_group_weijer']

weijer_colors = {
    "strong_amoc": "firebrick",
    "weak_amoc":   "steelblue",
}
weijer_labels = {
    "strong_amoc": "Historically strong AMOC",
    "weak_amoc":   "Historically weak AMOC",
}

# %%
fig, ax = plt.subplots(figsize=(10, 5))
ax.plot(vol_hist_plot.time, vol_hist_plot, color="black", lw=1.5, label="Historical (CARRA)")

for group in ["strong_amoc", "weak_amoc"]:
    runs = []
    for i, gcm_scen in enumerate(vol_cmip6.gcm_scenario.values):
        if gcm_scen in weijer_map and weijer_map[gcm_scen] == group:
            runs.append(vol_cmip6.isel(gcm_scenario=i).values)

    if not runs:
        continue
    arr = np.array(runs)  # (n_runs, n_time)
    med = np.median(arr, axis=0)
    p10 = np.percentile(arr, 10, axis=0)
    p90 = np.percentile(arr, 90, axis=0)
    c = weijer_colors[group]
    ax.plot(proj_time, med, color=c, label=f"{weijer_labels[group]}  (n={len(runs)})")
    ax.fill_between(proj_time, p10, p90, color=c, alpha=0.2)

ax.axhline(100, color="gray", lw=0.7, ls=":")
ax.set_xlabel("Year")
ax.set_ylabel("Volume (% of 2020)")
ax.legend(fontsize=10)
ax.grid(True, alpha=0.3)
fig.suptitle("Iceland total glacier volume by Weijer AMOC streamfunction (2000-2100)",
             fontsize=13)
plt.tight_layout()
fig.savefig(f"{fig_dir}/volume_projections_weijer_index.png", dpi=150, bbox_inches="tight")
plt.show()

# %% [markdown]
# ## Projections by AMOC index — glacier runoff
#
# The plots above track glacier *volume*. Glacier *runoff* — melt water and rain
# leaving the glacierised area — is arguably more directly relevant to downstream water
# resources (e.g. hydropower reservoirs) than volume itself. The runoff definition, unit
# conversion, and 11-year centred smoothing below follow the same approach as the OGGM-EDU
# tutorial notebook on glacier water resources:
# <https://edu-notebooks.oggm.org/oggm-edu/glacier_water_resources_projections.html>
#
# Runoff is defined as on-glacier melt + off-glacier melt (of snow that fell within the
# glacier outline but on now ice-free terrain) + on-glacier and off-glacier liquid
# precipitation, each clipped at zero before summing (a small number of grid cells can
# otherwise go slightly negative due to numerical noise in the mass-balance model).
#
# These hydro variables are only stored in the CMIP5/CMIP6 projection files, not in the
# historical run, so there is no historical (pre-2026) trace in this plot — unlike the
# volume plots above.

# %%
runoff_vars = ["melt_off_glacier", "melt_on_glacier",
               "liq_prcp_off_glacier", "liq_prcp_on_glacier"]

# annual total runoff (Gt yr-1) for the whole of Iceland, per GCM/scenario run
runoff_cmip6 = sum(ds_cmip6[v].clip(min=0) for v in runoff_vars).sum(dim="rgi_id") / 1e12

fig, ax = plt.subplots(figsize=(10, 5))

for group in ["low_index", "high_index", "neutral_index"]:
    runs = []
    for i, gcm_scen in enumerate(runoff_cmip6.gcm_scenario.values):
        if gcm_scen in amoc_map and amoc_map[gcm_scen] == group:
            # 11-year centred rolling mean, same smoothing as the OGGM-EDU tutorial
            smoothed = runoff_cmip6.isel(gcm_scenario=i).to_pandas().rolling(
                window=11, center=True, min_periods=1).mean()
            runs.append(smoothed.values)

    if not runs:
        continue
    arr = np.array(runs)  # (n_runs, n_time)
    med = np.nanmedian(arr, axis=0)
    p10 = np.nanpercentile(arr, 10, axis=0)
    p90 = np.nanpercentile(arr, 90, axis=0)
    c = amoc_colors[group]
    lw = 1.5 if group != "neutral_index" else 1
    alpha = 0.2 if group != "neutral_index" else 0.1
    ax.plot(runoff_cmip6.time, med, color=c, label=f"{amoc_labels[group]}  (n={len(runs)})", lw=lw)
    ax.fill_between(runoff_cmip6.time, p10, p90, color=c, alpha=alpha)

ax.set_xlabel("Year")
ax.set_ylabel("Annual glacier runoff (Gt yr$^{-1}$)")
ax.legend(fontsize=10)
ax.grid(True, alpha=0.3)
fig.suptitle("Iceland total glacier runoff by AMOC index (11-yr centred smoothing)",
             fontsize=13)
plt.tight_layout()
fig.savefig(f"{fig_dir}/runoff_projections_amoc_index.png", dpi=150, bbox_inches="tight")
plt.show()

# %% [markdown]
# This is for all of Iceland, hence quite averaged. To plot basin runoff,
# find the IDs of the glaciers you want to select and run:

# %%
sel_ids = ['RGI2025-v7.0-G-06-00570', 'RGI2025-v7.0-G-06-00571']  # random selection

# %%
# annual total runoff (Gt yr-1) for the selected glaciers, per GCM/scenario run
runoff_cmip6 = sum(ds_cmip6[v].sel(rgi_id=sel_ids).clip(min=0) for v in runoff_vars).sum(dim="rgi_id") / 1e12

fig, ax = plt.subplots(figsize=(10, 5))

for group in ["low_index", "high_index", "neutral_index"]:
    runs = []
    for i, gcm_scen in enumerate(runoff_cmip6.gcm_scenario.values):
        if gcm_scen in amoc_map and amoc_map[gcm_scen] == group:
            # 11-year centred rolling mean, same smoothing as the OGGM-EDU tutorial
            smoothed = runoff_cmip6.isel(gcm_scenario=i).to_pandas().rolling(
                window=11, center=True, min_periods=1).mean()
            runs.append(smoothed.values)

    if not runs:
        continue
    arr = np.array(runs)  # (n_runs, n_time)
    med = np.nanmedian(arr, axis=0)
    p10 = np.nanpercentile(arr, 10, axis=0)
    p90 = np.nanpercentile(arr, 90, axis=0)
    c = amoc_colors[group]
    lw = 1.5 if group != "neutral_index" else 1
    alpha = 0.2 if group != "neutral_index" else 0.1
    ax.plot(runoff_cmip6.time, med, color=c, label=f"{amoc_labels[group]}  (n={len(runs)})", lw=lw)
    ax.fill_between(runoff_cmip6.time, p10, p90, color=c, alpha=alpha)

ax.set_xlabel("Year")
ax.set_ylabel("Annual glacier runoff (Gt yr$^{-1}$)")
ax.legend(fontsize=10)
ax.grid(True, alpha=0.3)
fig.suptitle("Selected glaciers total glacier runoff by AMOC index (11-yr centred smoothing)",
             fontsize=13)
plt.tight_layout()
plt.show();

# %%