"""Compare the round 4a runs (BE + LMR forcing variants) with each other and with
round4's lam8_jjas_dyn (the attribution run), on their common glaciers.

    cd round4a && python compare.py

Same definitions as round4/sensitivity/lia_max.py (per-glacier max over the
LIA window, summed; area_min_h; Hugonnet per glacier), for two windows.
Also prints the static calibration and the dyn calibration convergence, and
writes compare.csv and figs/volume_area.png.
"""
import os

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

from oggm import cfg, utils

HERE = os.path.dirname(os.path.abspath(__file__))
R4 = os.path.join(HERE, '..', 'round4', 'sensitivity', 'output')
RUNS = {
    'v6 no shift': os.path.join(HERE, 'output', 'lam8_be_lmr21_dyn'),
    'v6 -0.14 K': os.path.join(HERE, 'output',
                               'lam8_be_lmr21_off1750-2000_dyn'),
    'v6 -0.38 K': os.path.join(HERE, 'output',
                               'lam8_be_lmr21_off1750-1779_dyn'),
    'r4 lam8_jjas_dyn': os.path.join(R4, 'lam8_jjas_dyn'),
}
V6_EXPS = {'v6 no shift': 'lam8_be_lmr21',
           'v6 -0.14 K': 'lam8_be_lmr21_off1750-2000',
           'v6 -0.38 K': 'lam8_be_lmr21_off1750-1779'}
STATS = os.path.join(HERE, '..', 'lmr_gdirs_v6_202610', 'exps', '{}',
                     'RGI62', 'b_160', 'L3', 'summary',
                     'glacier_statistics_11.csv')
WINDOWS = [slice(1800, 1875), slice(1800, 1900)]
RHO = 900.

cfg.initialize(logging_level='CRITICAL')
ref = pd.read_csv(os.path.join(HERE, '..', 'round1', 'filtered_update',
                               'Area_volume_change_LIA_Alps.csv'),
                  encoding='ISO-8859-1', index_col=0, skiprows=[1]).loc[17]

# --- Static calibration of the v6 builds
print('Static calibration (area-weighted, 11-01):')
rows = {}
for lab, e in V6_EXPS.items():
    s = pd.read_csv(STATS.format(e), index_col=0, low_memory=False)
    s = s.loc[s.rgi_subregion == '11-01']
    ok = s.error_task.isnull()
    w = s.loc[ok, 'rgi_area_km2']
    rows[lab] = {'failed': int((~ok).sum()),
                 'failed_km2': s.loc[~ok, 'rgi_area_km2'].sum(),
                 'prcp_fac': np.average(s.loc[ok, 'prcp_fac'], weights=w),
                 'melt_f': np.average(s.loc[ok, 'melt_f'], weights=w),
                 'melt_f_median': s.loc[ok, 'melt_f'].median(),
                 'temp_bias': np.average(s.loc[ok, 'temp_bias'], weights=w),
                 'inv_vol_km3': s.loc[ok, 'inv_volume_km3'].sum()}
    if lab == 'v6 no shift':
        stats = s
        fails = s.loc[~ok, 'error_task'].value_counts()
print(pd.DataFrame(rows).T.round(2).to_string())
print('failed tasks (no shift):', fails.to_dict(), '\n')

# --- Dyn calibration convergence
print('Dyn calibration (area-weighted share, %):')
rows = {}
area = stats.rgi_area_km2
for lab, d in RUNS.items():
    f = os.path.join(d, 'dyn_calib_reg11.csv')
    c = pd.read_csv(f, index_col=0)
    c = c.loc[c.index.isin(area.index)]
    a = area.reindex(c.index)
    r = (a.groupby(c.used_spinup_option).sum() / a.sum() * 100).to_dict()
    r = {k.replace('dynamic melt_f calibration ', ''): v for k, v in r.items()}
    tol = (c.dmdtda_mismatch_dynamic_calibration.abs() <=
           c.dmdtda_dynamic_calibration_given_error)
    r['within tol'] = a[tol].sum() / a.sum() * 100
    mf = c.melt_f.fillna(c.melt_f_before_dynamic_calibration)
    ok = mf.notnull() & a.notnull()
    r['melt_f mean'] = np.average(mf[ok], weights=a[ok])
    r['melt_f moved'] = (mf - c.melt_f_before_dynamic_calibration).abs()[ok] \
        .mean()
    r['iter median'] = c.run_dynamic_melt_f_calibration_iterations.median()
    rows[lab] = r
print(pd.DataFrame(rows).T.round(1).to_string(na_rep=''), '\n')

for lab in V6_EXPS:
    c = pd.read_csv(os.path.join(RUNS[lab], 'dyn_calib_reg11.csv'),
                    index_col=0)
    sr = c.run_dynamic_melt_f_calibration_stop_reason.dropna()
    sr = sr.str.slice(0, 60).value_counts()
    print(f'stop reasons, {lab}:')
    print(sr.head(6).to_string(), '\n')

# --- LIA and present day on the common glaciers
dss = {}
for lab, d in RUNS.items():
    with xr.open_dataset(os.path.join(d, 'lmr_mira_hist_reg11.nc')) as ds:
        dss[lab] = ds[['volume', 'area_min_h']].load()
common = stats.index.values
for ds in dss.values():
    ok = ds.rgi_id.values[~ds.volume.isel(time=0).isnull().values]
    common = np.intersect1d(common, ok)
rgi_common = stats.rgi_area_km2.loc[common].sum()
hug = utils.get_geodetic_mb_dataframe()
hug = hug.loc[hug.period == '2000-01-01_2020-01-01'].reindex(common)
obs_dm = float((hug.dmdtda * hug.area).sum()) * 20 * 1e-9

rows = {}
for lab, ds in dss.items():
    ds = ds.sel(rgi_id=common)
    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
    r = {}
    for w in WINDOWS:
        lia = ds.sel(time=w)
        tag = f'{w.start}-{w.stop}'
        r[f'area LIA {tag}'] = float(lia.area_min_h.max('time').sum()) * 1e-6
        r[f'vol LIA {tag}'] = float(lia.volume.max('time').sum()) * 1e-9
    r.update({'vol 1000': v(1000), 'vol 1500': v(1500), 'vol 1749': v(1749),
              'vol 1850': v(1850), 'vol 1900': v(1900),
              'area 2003': a(2003), 'vol 2015': v(2015),
              'dA 03-15 %': (a(2015) / a(2003) - 1) * 100,
              'dM 00-20 Gt': (v(2020) - v(2000)) * RHO * 1e-3})
    rows[lab] = r
df = pd.DataFrame(rows).T
obs = {f'area LIA {w.start}-{w.stop}': ref['Area LIA'] for w in WINDOWS}
obs.update({f'vol LIA {w.start}-{w.stop}': ref['Volume LIA'] for w in WINDOWS})
obs.update({'area 2003': rgi_common, 'vol 2015': ref['Volume 2015'],
            'dA 03-15 %': (ref['Area 2015'] / stats.rgi_area_km2.sum() - 1)
            * 100, 'dM 00-20 Gt': obs_dm})
df.loc['observed'] = pd.Series(obs)
df.to_csv(os.path.join(HERE, 'compare.csv'))
print(f'Common glaciers: {len(common)} of {len(stats)} '
      f'({rgi_common / stats.rgi_area_km2.sum():.1%} of the 11-01 area)')
print(df.round(1).T.to_string(na_rep=''))

# --- Figure: regional volume and area, from FIG_START
FIG_START = 1500
FIG_LABELS = {'v6 no shift': 'BE from 1750, LMR v2.1 before (no shift)',
              'v6 -0.14 K': 'BE from 1750, LMR v2.1 before (−0.14 K)',
              'v6 -0.38 K': 'BE from 1750, LMR v2.1 before (−0.38 K)',
              'r4 lam8_jjas_dyn': 'Current default run'}
fig, axs = plt.subplots(2, 1, figsize=(11, 8), sharex=True)
colors = ['#2a78d6', '#1baf7a', '#eb6834', '#52514e', '#b0afa8']
for (lab, ds), c in zip(dss.items(), colors):
    ds = ds.sel(rgi_id=common).sel(time=slice(FIG_START, None))
    ls = '--' if lab.startswith('r4') else '-'
    axs[0].plot(ds.time, ds.volume.sum('rgi_id') * 1e-9, color=c, ls=ls,
                lw=1.5, label=FIG_LABELS[lab])
    axs[1].plot(ds.time, ds.area_min_h.sum('rgi_id') * 1e-6, color=c, ls=ls,
                lw=1.5, label=FIG_LABELS[lab])
axs[0].set_ylabel('Volume (km³)')
axs[1].set_ylabel('Area (km², ice > 2 m)')
for ax in axs:
    ax.axvspan(1800, 1875, color='#d0cfc8', alpha=0.35, lw=0)
    ax.axvline(1750, color='#8a8982', lw=1, ls='--')
    ax.axvline(1850, color='k', lw=1)
    ax.grid(axis='y', color='#e6e5e0', lw=0.6)
    for sp in ['top', 'right']:
        ax.spines[sp].set_visible(False)
axs[0].legend(frameon=False, fontsize=9, loc='lower left')
axs[0].set_title(f'Alps (RGI 11-01), {len(common)} glaciers. Dashed line: '
                 'switch from LMR to BE (1750); black line: 1850; shaded: '
                 'LIA window 1800–1875', loc='left', fontsize=10)
axs[1].set_xlim(FIG_START, 2025)
fig.tight_layout()
os.makedirs(os.path.join(HERE, 'figs'), exist_ok=True)
fig.savefig(os.path.join(HERE, 'figs', 'volume_area.png'), dpi=150)