"""Questions A and B of the D&A follow-up doc (comms/, v261103), model only (no observations).

A (the "envelope" / IPCC view): the change of every HIST run vs the change of every NAT
  run, each measured from its own 1820-1850 mean (HIST strand i == NAT member i before
  1850, so HIST i branches off NAT i's own 1850 state, as in the sketch):
      dN_j = N_j(t) - mean_1820-1850 N_j,   dH_i = H_i(t) - mean_1820-1850 H_i
      f_A = 1 - median_j dN_j / median_i dH_i
  range: f_ij = 1 - dN_j / dH_i over all (HIST i, NAT j) pairs, 16-84%.
B (the conditional view, so far): for each 1850 state S (ens023/007/083/169/118), every
  NAT climate j from S (permutation runs, perm_a5v10):
      f_j(t | S) = (N_j(t | S) - H(t | S)) / (X_S - H(t | S)),  median and 16-84% over j.
At REF_YEAR (default 2025). NAT runs end 2005-2020: for REF_YEAR after that, each run is
held at its last 5-yr mean (as attr_fraction.py); for earlier years only the runs that
reach it are used.

    [REF_YEAR=2025] python question_a_vs_b.py -> figs/question_a_vs_b_ref<REF_YEAR>.{png,csv}
"""
import os

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec
from scipy.stats import gaussian_kde

HERE = os.path.dirname(os.path.abspath(__file__))
R7 = os.path.dirname(HERE)
ENS = 'ens_v8_adapbs_a5v10'
PERM = os.environ.get('PERM', 'perm_a5v10')
REF = int(os.environ.get('REF_YEAR', 2025))
BASE = (1820, 1850)
STATES = ['023', '007', '083', '169', '118']   # 1850 NAT volume quantiles 0/25/50/75/100%
NAT_LAST = 5
X0 = 1650
NAT_C, NAT_D, HIST_C, HIST_D, PUR, INK = '#5dade2', '#1b4f72', '#e6705f', '#922b21', '#8e44ad', '#52514e'


def read(f):
    df = pd.read_csv(f, index_col=0)
    df.index = df.index.astype(int)
    return df


def at_ref(df):
    """Values at REF: held at the last NAT_LAST-yr mean after a run ends, if REF is past all ends."""
    if REF <= max(df[c].last_valid_index() for c in df):
        return df.loc[REF].dropna()
    return df.apply(lambda x: x.dropna().iloc[-NAT_LAST:].mean())


def q(x):
    return np.nanquantile(x, [0.16, 0.5, 0.84])


vn = read(f'{R7}/filtered_ens/{ENS}/volume_km3_nat_ens_reg11.csv')
vh = read(f'{R7}/filtered_ens/{ENS}/volume_km3_hist_ens_reg11.csv')
assert (vh.loc[:1849] - vn[vh.columns].loc[:1849]).abs().max().max() == 0

# --- A -----------------------------------------------------------------------------
dN = vn - vn.loc[BASE[0]:BASE[1]].mean()
dH = vh - vh.loc[BASE[0]:BASE[1]].mean()
dn_r, dh_r = at_ref(dN), dH.loc[REF]
fA = 1 - dn_r.median() / dh_r.median()
pairs = (1 - dn_r.values[None, :] / dh_r.values[:, None]).ravel()
fA_q = q(pairs)
rows = [dict(question='A', state='all runs', n_nat=len(dn_r), n_hist=len(dh_r),
             dN_median=dn_r.median(), dH_median=dh_r.median(),
             f_median_of_medians=fA, f_q16=fA_q[0], f_q50=fA_q[1], f_q84=fA_q[2])]

# --- B -----------------------------------------------------------------------------
pv = read(f'{R7}/perm/filtered/{PERM}/volume_km3_perm_reg11.csv')
B = {}
for s in STATES:
    nat = pv[[c for c in pv.columns if c.startswith(f'state{s}_nat_')]]
    hist = pv[f'state{s}_hist']
    xs = hist.loc[1850]
    assert abs(xs - vh[f'ens{s}'].loc[1850]) < 1e-3
    n_r, h_r = at_ref(nat), hist.loc[REF]
    f = (n_r - h_r) / (xs - h_r)
    B[s] = dict(nat=nat, hist=hist, xs=xs, n_r=n_r, h_r=h_r, f=f)
    fq = q(f)
    rows.append(dict(question='B', state=f'ens{s}', n_nat=len(n_r), n_hist=1, X_S=xs,
                     N_median=n_r.median(), H=h_r, A_median=(n_r - h_r).median(), L=xs - h_r,
                     f_q16=fq[0], f_q50=fq[1], f_q84=fq[2]))
tab = pd.DataFrame(rows).set_index(['question', 'state'])
os.makedirs(f'{HERE}/figs', exist_ok=True)
out = f'{HERE}/figs/question_a_vs_b_ref{REF}'
tab.round(3).to_csv(out + '.csv')
with pd.option_context('display.width', 250, 'display.max_columns', 20):
    print(tab.round(3).to_string())


# --- figure ------------------------------------------------------------------------
def kde_side(ax, vals, color, label=None, flip=False, scale=1.0):
    vals = np.asarray(vals)
    lo, hi = ax.get_ylim()
    y = np.linspace(lo, hi, 400)
    if vals.std() < 1e-6:
        ax.axhline(vals[0], color=color, lw=2, label=label)
        return
    d = gaussian_kde(vals)(y)
    d = scale * d / d.max()
    x = -d if flip else d
    ax.fill_betweenx(y, 0, x, color=color, alpha=0.35, lw=0)
    ax.plot(x, y, color=color, lw=1.5, label=label)


def side_style(ax, flip=False):
    ax.set_xlim((-1.15, 0) if flip else (0, 1.15))
    ax.set_xticks([])
    for sp in ('top', 'right', 'bottom') if not flip else ('top', 'left', 'bottom'):
        ax.spines[sp].set_visible(False)
    ax.tick_params(axis='y', labelleft=False, left=False, right=False)


def lines_style(ax):
    ax.grid(alpha=0.3)
    ax.axvline(1850, color='0.3', lw=1, ls=':')
    ax.set_xlim(X0, REF)


fig = plt.figure(figsize=(17, 10.5))
gs = GridSpec(2, 8, figure=fig, width_ratios=[0.4, 2.3, 0.5, 0.75, 2.3, 0.5, 1.15, 2.0],
              wspace=0.0, hspace=0.32)
held = REF > int(vn.apply(lambda x: x.last_valid_index()).max())
ref_lab = f'{REF}' + (' (NAT held at their last 5-yr mean)' if held else '')

# (a) volume, like the sketch
ap = fig.add_subplot(gs[0, 0])
a = fig.add_subplot(gs[0, 1], sharey=ap)
ar = fig.add_subplot(gs[0, 2], sharey=ap)
yr = vn.loc[X0:REF]
a.plot(yr.index, yr.values, color=NAT_C, lw=0.4, alpha=0.25)
a.plot(vh.loc[1850:REF].index, vh.loc[1850:REF].values, color=HIST_C, lw=0.4, alpha=0.3)
a.plot(yr.index, yr.median(axis=1).where(yr.notnull().sum(axis=1) >= 100), color=NAT_D, lw=2,
       label=f'NAT ({vn.shape[1]} runs), median')
a.plot(vh.loc[1850:REF].index, vh.loc[1850:REF].median(axis=1), color=HIST_D, lw=2,
       label=f'HIST ({vh.shape[1]} runs, each from its NAT\nrun\'s own 1850 state), median')
a.set_ylim(40, 290)
lines_style(a)
a.tick_params(axis='y', labelleft=False)
a.set_title('(a) Question A: all runs, volume (the sketch)', loc='left', fontsize=10.5)
a.legend(loc='lower left', fontsize=8, frameon=False)
kde_side(ap, vn.loc[1850], PUR, flip=True)
side_style(ap, flip=True)
ap.tick_params(axis='y', labelleft=True, left=True)
ap.set_ylabel('Regional glacier volume (km³)')
ap.text(-0.05, 275, '1850\nstates', color=PUR, ha='right', va='top', fontsize=8.5)
kde_side(ar, at_ref(vn), NAT_D)
kde_side(ar, vh.loc[REF], HIST_D)
side_style(ar)
ar.text(0.05, 285, str(REF), ha='left', va='top', fontsize=8.5, color=INK)

# (b) volume change from each run's 1820-1850 mean
b = fig.add_subplot(gs[0, 4])
br = fig.add_subplot(gs[0, 5], sharey=b)
y2 = dN.loc[X0:REF]
b.plot(y2.index, y2.values, color=NAT_C, lw=0.4, alpha=0.25)
b.plot(dH.loc[1850:REF].index, dH.loc[1850:REF].values, color=HIST_C, lw=0.4, alpha=0.3)
b.plot(y2.index, y2.median(axis=1).where(y2.notnull().sum(axis=1) >= 100), color=NAT_D, lw=2)
b.plot(dH.loc[1850:REF].index, dH.loc[1850:REF].median(axis=1), color=HIST_D, lw=2)
b.axvspan(*BASE, color='0.5', alpha=0.12, lw=0)
b.axhline(0, color=INK, lw=0.8)
b.set_ylim(-150, 90)
lines_style(b)
b.set_ylabel('ΔV from the run\'s own 1820–1850 mean (km³)')
b.set_title('(b) Question A: volume change ΔV', loc='left', fontsize=10.5)
kde_side(br, dn_r, NAT_D)
kde_side(br, dh_r, HIST_D)
side_style(br)
for v, c, lab in [(dn_r.median(), NAT_D, 'NAT'), (dh_r.median(), HIST_D, 'HIST')]:
    br.annotate(f'{lab} median\n{v:+.0f} km³', (0.02, v), xytext=(0.25, v + (14 if lab == 'NAT' else -20)),
                fontsize=8, color=c, va='center')
b.annotate('', xy=(REF - 3, dh_r.median()), xytext=(REF - 3, dn_r.median()),
           arrowprops=dict(arrowstyle='<->', color='k', lw=1.3))
b.text(REF - 6, 0.5 * (dn_r.median() + dh_r.median()),
       f'anthropogenic\n{dn_r.median() - dh_r.median():.0f} km³', ha='right', va='center', fontsize=8.5)
b.text(1838, 80, 'baseline', ha='center', fontsize=8, color=INK)

# (c) f_A
c = fig.add_subplot(gs[0, 7])
pp = pairs[(pairs > -0.5) & (pairs < 2.5)]
cn, _, _ = c.hist(100 * pp, bins=np.arange(-50, 251, 5), color='0.6', alpha=0.7)
c.set_ylim(0, 1.7 * cn.max())
c.axvline(100, color='k', lw=0.6, ymax=0.62, label='100%')
c.axvline(100 * fA_q[1], color='k', lw=1.6, ymax=0.62, label='median')
for v in (fA_q[0], fA_q[2]):
    c.axvline(100 * v, color='k', lw=1.2, ls=':', ymax=0.62, label='16–84%' if v == fA_q[0] else None)
c.legend(loc='center left', bbox_to_anchor=(0.0, 0.5), fontsize=8, frameon=False)
c.set_xlim(-10, 220)
c.set_yticks([])
c.set_xlabel(f'1 − ΔN_j / ΔH_i (%), all {len(dn_r)}×{len(dh_r)} pairs')
c.set_title(f'(c) Question A: share in {REF}', loc='left', fontsize=10.5)
c.text(0.97, 0.95, f'f_A = {100 * fA:.0f}%\n(16–84%: {100 * fA_q[0]:.0f}–{100 * fA_q[2]:.0f})',
       transform=c.transAxes, ha='right', va='top', fontsize=12, fontweight='bold')
c.text(0.97, 0.80, f'NAT median ΔV {dn_r.median():+.0f} km³\nHIST median ΔV {dh_r.median():+.0f} km³',
       transform=c.transAxes, ha='right', va='top', fontsize=8.5, color=INK)
for sp in ('top', 'right', 'left'):
    c.spines[sp].set_visible(False)

# (d), (e) question B, smallest and largest 1850 state
for k, (s, cols) in enumerate([('023', (slice(0, 2), 2)), ('118', (4, 5))]):
    d = fig.add_subplot(gs[1, cols[0]])
    dr = fig.add_subplot(gs[1, cols[1]], sharey=d)
    st = B[s]
    pre = vh[f'ens{s}'].loc[X0:1850]
    nat = st['nat'].loc[:REF]
    d.plot(nat.index, nat.values, color=NAT_C, lw=0.4, alpha=0.25)
    d.plot(nat.index, nat.median(axis=1).where(nat.notnull().sum(axis=1) >= 100), color=NAT_D, lw=2,
           label=f'NAT from this state ({nat.shape[1]} climates), median')
    d.plot(pre.index, pre.values, color=INK, lw=1.8, label='history of this state (NAT climate)')
    d.plot(st['hist'].loc[1850:REF].index, st['hist'].loc[1850:REF], color=HIST_D, lw=2.2,
           label='HIST from this state')
    d.plot([1850], [st['xs']], 'o', color='k', ms=7, zorder=5)
    d.annotate(f'X_S = {st["xs"]:.0f} km³', (1850, st['xs']), xytext=(-8, 10), textcoords='offset points',
               ha='right', fontsize=9)
    d.set_ylim(40, 290)
    lines_style(d)
    d.set_ylabel('Regional glacier volume (km³)')
    small = 'smallest' if s == '023' else 'largest'
    fq_s = q(st['f'])
    d.set_title(f'({"de"[k]}) Question B: ens{s}, {small} 1850 state: f = {100 * fq_s[1]:.0f}% '
                f'({100 * fq_s[0]:.0f}–{100 * fq_s[2]:.0f})', loc='left', fontsize=10.5)
    if k == 0:
        d.legend(loc='lower left', fontsize=8, frameon=False)
    kde_side(dr, st['n_r'], NAT_D)
    dr.axhline(st['h_r'], color=HIST_D, lw=2)
    side_style(dr)

# (f) the shares side by side
f_ = fig.add_subplot(gs[1, 7])
labs, xs_ = [], []
items = [('A', fA_q, '0.35')]
for s in STATES:
    items.append((f'{s}\n{B[s]["xs"]:.0f}', q(B[s]['f']), HIST_C))
for i, (lab, v, col) in enumerate(items):
    f_.errorbar([i], [100 * v[1]], yerr=[[100 * (v[1] - v[0])], [100 * (v[2] - v[1])]], fmt='o',
                color=col, ms=7, capsize=3, lw=1.5)
    f_.text(i + 0.12, 100 * v[1], f'{100 * v[1]:.0f}', fontsize=8.5, va='center')
    labs.append(lab)
f_.axhline(100, color='0.4', lw=0.8, ls='--')
f_.axvline(0.5, color='0.6', lw=0.8)
f_.set_xticks(range(len(items)))
f_.set_xticklabels(labs, fontsize=8)
f_.set_xlabel('A  |  B: 1850 state (ens) and X_S (km³)', fontsize=9)
f_.set_xlim(-0.6, len(items) - 0.4)
f_.set_ylabel(f'Anthropogenic share, {REF} (%)')
f_.set_title('(f) Share, A vs B (median, 16–84%)', loc='left', fontsize=10.5)
f_.grid(axis='y', alpha=0.3)

fig.suptitle(f'Question A (envelopes: all HIST vs all NAT changes) vs question B (conditional on the '
             f'1850 state), regional volume, model only, {ref_lab}', x=0.01, ha='left', fontsize=12)
fig.savefig(out + '.png', dpi=150, bbox_inches='tight')
print('wrote', out + '.png')