"""One HIST member against the NAT ensemble, for the talk (round 5 copy of
round4/ensemble/plot_talk_member.py). For volume and area, four versions on
the same axes:
- <var>_<NAME>_<MEMBER>:          HIST MEMBER + obs, NAT up to 1850
- <var>_<NAME>_nat_only:          NAT only, up to 1850
- <var>_<NAME>_<MEMBER>_natfull:  HIST MEMBER + obs, NAT to the end of each
                                  member (2006-2021; the NAT mean is over the
                                  members available each year)
- <var>_<NAME>_hist5_natfull:     the same with the 5 HIST quantile members
                                  (Hugonnet starts from their mean in 2000)

    MEMBER=ens010 python plot_talk_member.py
"""
import os

import matplotlib.pyplot as plt
import seaborn as sns

from talk_common import (NAME, TALK, load, order, colors, label,
                         VOLUME_REFS, AREA_REFS)

MEMBER = os.environ.get('MEMBER', 'ens010')
NAT_END = 1850


def figure(var, ylabel, refs):
    nat_all, hist_all = load(var)
    hist = hist_all[MEMBER]
    ylim = None
    for nat_only, nat_full, all_hist in ((False, False, False), (True, False, False),
                                         (False, True, False), (False, True, True)):
        nat = nat_all if nat_full else nat_all.loc[:NAT_END]
        with sns.axes_style('whitegrid'):
            fig, ax = plt.subplots(figsize=(10, 5.5))
            for i, c in enumerate(nat.columns):
                ax.plot(nat.index, nat[c].values, color='C0', alpha=0.15, linewidth=0.8,
                        label='NAT ensemble' if i == 0 else '_nolegend_')
            ax.plot(nat.index, nat.mean(axis=1), color='navy', linewidth=3,
                    label='NAT ensemble mean')
            if all_hist:
                for m in order:
                    ax.plot(hist_all.index, hist_all[m].values, color=colors[m],
                            linewidth=1.8, label=label(m))
            elif not nat_only:
                ax.plot(hist.index, hist.values, color='C3', linewidth=2,
                        label=f'HIST {MEMBER}')
            if not nat_only:
                for r in refs:
                    r(ax, hist_all[order] if all_hist else hist)
            ax.set_xlim([1300, hist.index.max()])
            if ylim is None:
                ylim = ax.get_ylim()
            ax.set_ylim(ylim)
            ax.set_xlabel('Year')
            ax.set_ylabel(ylabel)
            ax.legend(loc='lower left', fontsize=9 if all_hist else 10,
                      ncol=2 if all_hist else 1)
            name = (f'{var}_{NAME}_nat_only' if nat_only else
                    f'{var}_{NAME}_hist5_natfull' if all_hist else
                    f'{var}_{NAME}_{MEMBER}_natfull' if nat_full else
                    f'{var}_{NAME}_{MEMBER}')
            fig.savefig(f'{TALK}/{name}.png', bbox_inches='tight', dpi=150)
            plt.close(fig)
            print(f'Written: {TALK}/{name}.png')


figure('volume', 'Total volume (km$^3$)', VOLUME_REFS)
figure('area', 'Total area (km$^2$)', AREA_REFS)