"""LIA maximum and present-day fit per experiment, against observations.

    cd round4/sensitivity && python lia_max.py

Each glacier may have reached its LIA maximum at a different time, so the
LIA state is the per-glacier max of volume and area over LIA_PERIOD, summed
over the glaciers. All experiments are compared on the same glaciers: the
ones which ran in every experiment (some tiny ones grow out of their domain
and have no output, and which ones depends on the experiment). Areas are
`area_min_h` (ice thicker than 2 m), as in round3.

Goodness of fit of the present day, on the same glaciers:
- the initial area (year 1, the RGI-date flowlines) against the RGI area;
- the area in 2003 against the RGI area (RGI dates are mostly 2003);
- the relative area change 2003 -> 2015, against Reinthaler 2015 / RGI;
- the mass change 2000-2020, against Hugonnet (per glacier, summed).

Writes lia_max.csv and prints one table per sensitivity axis.
"""
import os

import numpy as np
import pandas as pd
import xarray as xr

from oggm import cfg, utils

LIA_PERIOD = slice(1800, 1900)
RHO = 900.  # ice density for volume -> mass, as OGGM
V5_STATS = ('/home/www/fmaussion/runs/lmr_2025/lmr_gdirs_v5_202610/exps/'
            '{}/RGI62/b_160/L3/summary/glacier_statistics_11.csv')
EXPS = ['default', 'pf1.5', 'pf3.5', 'mf3', 'mf7', 'mf9', 'centerlines',
        'lambda6', 'lam8', 'lam6_ib', 'lam8_ib', 'lam8_mf7', 'lam8_ib_mf7',
        'lam8_ib_dyn20', 'lam8_ib_dynbug', 'lam8_dynbug',
        'lam8_ib_mf7_dynbug', 'lam6_ib_dynbug',
        'lam8_ib_dyn', 'lam8_dyn', 'lam8_ib_mf7_dyn', 'lam6_ib_dyn']
# Each axis: the experiments along it, default included, and their label
AXES = {
    'prcp_fac': {'pf1.5': 1.5, 'default': 2.5, 'pf3.5': 3.5},
    'melt_f': {'mf3': 3, 'default': 5, 'mf7': 7, 'mf9': 9},
    'flowlines': {'default': 'elev. bands', 'centerlines': 'centerlines'},
    'trapezoid_lambdas': {'default': 2, 'lambda6': 6, 'lam8': 8},
    # towards the best run: lambda x inversion volume x melt_f
    'best_search': {'lambda6': 'l6', 'lam6_ib': 'l6 ib', 'lam8': 'l8',
                    'lam8_ib': 'l8 ib', 'lam8_mf7': 'l8 mf7',
                    'lam8_ib_mf7': 'l8 ib mf7'},
    # per-glacier dynamic melt_f calibration, maxiter 40. dynbug: with the
    # OGGM step-length bug (melt_f could move only ~+-3.5); dyn20: idem,
    # maxiter 20; dyn: bug fixed (third hunk of the OGGM patch)
    'dyn_calib': {'lam8_ib': 'l8 ib', 'lam8_ib_dyn20': 'l8 ib dyn20',
                  'lam8_ib_dynbug': 'l8 ib dynbug', 'lam8_ib_dyn': 'l8 ib dyn',
                  'lam8': 'l8', 'lam8_dynbug': 'l8 dynbug',
                  'lam8_dyn': 'l8 dyn', 'lam8_ib_mf7': 'l8 ib mf7',
                  'lam8_ib_mf7_dynbug': 'l8 ib mf7 dynbug',
                  'lam8_ib_mf7_dyn': 'l8 ib mf7 dyn', 'lam6_ib': 'l6 ib',
                  'lam6_ib_dynbug': 'l6 ib dynbug', 'lam6_ib_dyn': 'l6 ib dyn'},
}

cfg.initialize(logging_level='CRITICAL')

# Reinthaler & Paul, all Alps (the row round3 uses)
ref = pd.read_csv('../../round1/filtered_update/Area_volume_change_LIA_Alps.csv',
                  encoding='ISO-8859-1', index_col=0, skiprows=[1]).loc[17]

exps = [e for e in EXPS
        if os.path.exists(f'output/{e}/lmr_mira_hist_reg11.nc')]
missing = [e for e in EXPS if e not in exps]

dss = {}
for e in exps:
    with xr.open_dataset(f'output/{e}/lmr_mira_hist_reg11.nc') as ds:
        dss[e] = ds[['volume', 'area_min_h']].load()

# The Alps only: 11-02 (Pyrenees) is outside the LMR grid, and not in the
# Reinthaler reference either
stats = pd.read_csv(V5_STATS.format('default'), index_col=0, low_memory=False)
stats = stats.loc[stats.rgi_subregion == '11-01']
rgi_area = stats.rgi_area_km2

# The glaciers which ran in every experiment
common = stats.index.values
for e, ds in dss.items():
    ok = ds.rgi_id.values[~ds.volume.isel(time=0).isnull().values]
    common = np.intersect1d(common, ok)
n_all = len(stats)

# Observations on the common glaciers
rgi_area_common = rgi_area.loc[common].sum()
rgi_area_1101 = rgi_area.sum()
hug = utils.get_geodetic_mb_dataframe()
hug = hug.loc[hug.period == '2000-01-01_2020-01-01'].reindex(common)
# dmdtda is in m w.e. yr-1, area in m2: x 1000 kg m-3 x 20 yr x 1e-12 Gt/kg
obs_dm = float((hug.dmdtda * hug.area).sum()) * 20 * 1e-9
obs_dm_err = float(np.sqrt(((hug.err_dmdtda * hug.area) ** 2).sum())) * 20 * 1e-9

OBS = {
    'area_lia': ref['Area LIA'], 'vol_lia': ref['Volume LIA'],
    'area_ratio': ref['Area LIA'] / ref['Area 2003'],
    'vol_ratio': ref['Volume LIA'] / ref['Volume 2015'],
    'area_y1': rgi_area_common, 'area_2003': rgi_area_common,
    # relative change 2003 -> 2015: Reinthaler 2015 against the RGI area
    'dA_03_15_%': (ref['Area 2015'] / rgi_area_1101 - 1) * 100,
    'dM_00_20_Gt': obs_dm,
}

rows = {}
for e, ds in dss.items():
    ds = ds.sel(rgi_id=common)
    lia = ds.sel(time=LIA_PERIOD)
    a = lambda y: float(ds.area_min_h.sel(time=y).sum()) * 1e-6
    v = lambda y: float(ds.volume.sel(time=y).sum()) * 1e-9
    peak_yr = lia.area_min_h.idxmax('time')
    w = lia.area_min_h.max('time').values
    order = np.argsort(peak_yr.values)
    cw = np.cumsum(w[order]) / w.sum()
    r = {
        'area_lia': float(lia.area_min_h.max('time').sum()) * 1e-6,
        'vol_lia': float(lia.volume.max('time').sum()) * 1e-9,
        'area_y1': float(ds.area_min_h.isel(time=0).sum()) * 1e-6,
        'area_2003': a(2003), 'area_2015': a(2015),
        'vol_2015': v(2015),
        'dA_03_15_%': (a(2015) / a(2003) - 1) * 100,
        'dM_00_20_Gt': (v(2020) - v(2000)) * RHO * 1e-3,
        'peak_yr_median': int(peak_yr.values[order][np.searchsorted(cw, 0.5)]),
        'peak_at_end_%': float(w[peak_yr.values >= LIA_PERIOD.stop - 4].sum()
                               / w.sum() * 100),
    }
    r['area_ratio'] = r['area_lia'] / r['area_2003']
    r['vol_ratio'] = r['vol_lia'] / r['vol_2015']
    rows[e] = r
df = pd.DataFrame(rows).T
df.loc['observed'] = pd.Series(OBS)
df.to_csv('lia_max.csv')

print(f'LIA max = per-glacier max over {LIA_PERIOD.start}-{LIA_PERIOD.stop}, '
      f'summed. Common glaciers: {len(common)} of {n_all} in 11-01 '
      f'({rgi_area_common / rgi_area.sum():.1%} of the 11-01 RGI area), in '
      f'{len(exps)} experiments' + (f' (missing: {missing})' if missing else ''))
print(f'Observed: Reinthaler Alps LIA {OBS["area_lia"]:.0f} km2 / '
      f'{OBS["vol_lia"]:.0f} km3, LIA/2003 area {OBS["area_ratio"]:.2f}, '
      f'LIA/2015 volume {OBS["vol_ratio"]:.2f}; RGI area {rgi_area_common:.0f} km2; '
      f'area change 2003-2015 {OBS["dA_03_15_%"]:.1f}% (Reinthaler 2015 / RGI 11-01); '
      f'Hugonnet 2000-2020 {obs_dm:.1f} +- {obs_dm_err:.1f} Gt\n')

fit = ['area_y1', 'area_2003', 'dA_03_15_%', 'dM_00_20_Gt']
lia_cols = ['area_lia', 'area_ratio', 'vol_lia', 'vol_ratio',
            'peak_yr_median', 'peak_at_end_%']
for axis, members in AXES.items():
    members = {e: v for e, v in members.items() if e in df.index}
    if len(members) < 2:
        continue
    t = df.loc[list(members) + ['observed'], fit + lia_cols].copy()
    t.insert(0, axis, list(members.values()) + ['obs'])
    print(f'--- {axis} ---')
    print(t.to_string(float_format=lambda x: f'{x:.2f}', na_rep=''))
    print()