"""What we compare to the LIA reference (round 5 copy of
round4/ensemble/plot_lia_max.py). Two figures (volume, area), each with:
- left: the regional time series, the window shaded;
- right: the sum of the per-glacier maxima over the window for the NAT
  members (box + points), the 5 HIST quantile members, and the reference.

The volume reference is anchored on the RGI-date volume (talk_common.py):
compare it with the LIA volume of the runs, whose 2003 volume is the
inversion's. The per-glacier maxima come from filtered_ens/<NAME>/
lia_max_members.csv (preprocess_ens.py).

    python plot_lia_max.py   -> talk/lia_max_{volume,area}_<NAME>.png
"""
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns

from talk_common import (NAME, TALK, FDIR, ids, load, order, colors, label,
                         VOLUME_REFS, AREA_REFS, dfo, ref_lia_volume, REF_DV_ERR,
                         REF_AREA_ERR)

WINDOW = (1800, 1875)
y0, y1 = WINDOW
lm = pd.read_csv(f'{FDIR}/lia_max_members.csv', index_col=0)
nat_lm = lm.loc[lm.kind == 'nat']
hist_lm = lm.loc[lm.kind == 'hist'].set_index('member')


def lia_figure(var, col, ref_lia, ref_err, ylabel, refs, fname):
    series_nat, series_hist = load(var)
    with sns.axes_style('whitegrid'):
        fig, (ax, axr) = plt.subplots(
            1, 2, figsize=(12, 5.5), sharey=True,
            gridspec_kw={'width_ratios': [5, 1.3], 'wspace': 0.04})
        ax.axvspan(y0, y1, color='0.85', alpha=0.5, lw=0,
                   label=f'Window of the per-glacier max ({y0}–{y1})')
        for i, c in enumerate(series_nat.columns):
            ax.plot(series_nat.index, series_nat[c].values, color='C0',
                    alpha=0.15, linewidth=0.8,
                    label='NAT ensemble' if i == 0 else '_nolegend_')
        ax.plot(series_nat.index, series_nat.mean(axis=1), color='navy',
                linewidth=2, label='NAT ensemble mean')
        for m in order:
            ax.plot(series_hist.index, series_hist[m], color=colors[m],
                    linewidth=2, label=label(m))
        for r in refs:
            r(ax, series_hist[order])
        ax.set_xlim([1300, series_nat.index.max()])
        ax.set_xlabel('Year')
        ax.set_ylabel(ylabel)
        ax.legend(loc='lower left', fontsize=8.5)

        v = nat_lm[col].astype(float).values
        axr.boxplot(v, positions=[0], widths=0.5, showfliers=False,
                    medianprops={'color': 'navy', 'linewidth': 2},
                    boxprops={'color': 'C0'}, whiskerprops={'color': 'C0'},
                    capprops={'color': 'C0'})
        jit = np.random.default_rng(0).uniform(-0.15, 0.15, len(v))
        axr.scatter(jit, v, s=8, color='C0', alpha=0.4, lw=0)
        for m, x in zip(order, np.linspace(0.75, 1.25, len(order))):
            axr.scatter([x], [float(hist_lm.loc[m, col])], s=60, color=colors[m],
                        edgecolor='white', linewidth=1, zorder=3)
        axr.errorbar([2], [ref_lia], yerr=[ref_err], fmt='D', ms=8, color='k',
                     capsize=5, elinewidth=1.5, zorder=3)
        axr.set_xticks([0, 1, 2])
        axr.set_xticklabels(['NAT', 'HIST', 'Ref.'])
        axr.set_xlim(-0.6, 2.6)
        axr.set_title(f'Sum of per-glacier\nmax, {y0}–{y1}', fontsize=10)
        fig.suptitle(f'{NAME}: {len(ids)} glaciers, {series_nat.shape[1]} NAT '
                     f'members, {len(order)} HIST members', y=0.98)
        fig.savefig(fname, bbox_inches='tight', dpi=150)
        plt.close(fig)
    h = hist_lm.loc[order, col].astype(float)
    print(f'{col}: NAT median {np.median(v):.0f} [P5–P95 {np.percentile(v, 5):.0f}–'
          f'{np.percentile(v, 95):.0f}], HIST {h.min():.0f}–{h.max():.0f}, '
          f'ref {ref_lia:.0f} ± {ref_err:.0f}  -> {fname}')


lia_figure('volume', f'vol_at_amax_{y0}_{y1}', ref_lia_volume, REF_DV_ERR,
           'Total volume (km$^3$)', VOLUME_REFS, f'{TALK}/lia_max_volume_{NAME}.png')
lia_figure('area', f'area_max_{y0}_{y1}', dfo['Area LIA'], REF_AREA_ERR * dfo['Area LIA'],
           'Total area (km$^2$)', AREA_REFS, f'{TALK}/lia_max_area_{NAME}.png')