"""How much does a small shift of the natural counterfactual move the attribution?

Conditional permutation runs (round7/perm) with the NAT temperature as delivered and shifted
cooler from 1850 (linear ramp 1850-1900, then constant; data/build_nat_shift.py). HIST is
the same in all of them.
(a) regional volume of one 1850 state (STATE, default 083 = median): HIST and the natural
    median (16-84%) of each variant, 1850-2021 (only years the natural runs really have);
(b) the anthropogenic loss A(t) = natural median - HIST of that state;
(c) the 2025 share of the loss since 1850 (Reinthaler + NAT combined, 16-84%), from
    attr_fraction.py (natural runs held at their last 5-yr mean after they end).

    python nat_shift_effect.py  -> figs/nat_shift_effect.png (+ .csv)
"""
import os

import numpy as np
import pandas as pd
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

HERE = os.path.dirname(os.path.abspath(__file__))
R7 = os.path.dirname(HERE)
STATE = os.environ.get('STATE', '083')
NAME = 'ens_v8_adapbs_a5v10'
VARIANTS = [('perm_a5v10', 0.0, '#2e86c1'), ('perm_a5v10_shift-0.2', 0.2, '#7d3c98'),
            ('perm_a5v10_shift-0.5', 0.5, '#4a235a')]
HIST_C, INK2 = '#c0392b', '#52514e'
T_END = 2021

rows, have = [], []
for perm, d, c in VARIANTS:
    f = f'{R7}/perm/filtered/{perm}/volume_km3_perm_reg11.csv'
    tag = '' if perm == 'perm_a5v10' else '_' + perm.replace('perm_a5v10_', '')
    comb = f'{HERE}/figs/attr_fraction_{NAME}{tag}_ref2025_reinthaler_combined.csv'
    if not (os.path.exists(f) and os.path.exists(comb)):
        continue
    v = pd.read_csv(f, index_col=0)
    v.index = v.index.astype(int)
    nat = v[[x for x in v.columns if x.startswith(f'state{STATE}_nat_')]].loc[1850:T_END]
    hist = v[f'state{STATE}_hist'].loc[1850:T_END]
    q = pd.read_csv(comb, index_col=0).share
    a = nat.sub(hist, axis=0)
    have.append((perm, d, c, nat, hist, a, q))
    rows.append(dict(variant=perm, shift_K=d, A_2021_median=a.loc[T_END].median(),
                     A_2015_median=a.loc[2015].median(), share2025_q16=q['q16'],
                     share2025_q50=q['q50'], share2025_q84=q['q84']))
tab = pd.DataFrame(rows).set_index('variant')
tab.round(3).to_csv(f'{HERE}/figs/nat_shift_effect.csv')
print(tab.round(3).to_string())

fig, (ax, bx, cx) = plt.subplots(1, 3, figsize=(15, 4.8),
                                 gridspec_kw={'width_ratios': [2.2, 2.2, 1.2], 'wspace': 0.28})
for perm, d, c, nat, hist, a, q in have:
    lab = 'natural, as delivered' if d == 0 else f'natural, {d:g} K cooler from 1850'
    qq = nat.quantile([0.16, 0.5, 0.84], axis=1).T
    ax.fill_between(qq.index, qq[0.16], qq[0.84], color=c, alpha=0.15, lw=0)
    ax.plot(qq.index, qq[0.5], color=c, lw=2, label=f'{lab} (median, 16–84%)')
    qa = a.quantile([0.16, 0.5, 0.84], axis=1).T
    bx.fill_between(qa.index, qa[0.16], qa[0.84], color=c, alpha=0.15, lw=0)
    bx.plot(qa.index, qa[0.5], color=c, lw=2, label=f'{lab}: {qa.loc[T_END, 0.5]:.0f} km³ in {T_END}')
    cx.errorbar([d], [100 * q['q50']], yerr=[[100 * (q['q50'] - q['q16'])], [100 * (q['q84'] - q['q50'])]],
                fmt='o', color=c, ms=8, capsize=4, lw=1.6)
    cx.annotate(f'{100 * q["q50"]:.0f}%', (d, 100 * q['q50']), xytext=(8, 0),
                textcoords='offset points', va='center', fontsize=9)
ax.plot(hist.index, hist, color=HIST_C, lw=2.2, label='historical')
ax.set_ylabel('Regional volume (km³)')
ax.set_title(f'(a) One 1850 state (ens{STATE}): historical vs natural runs', loc='left', fontsize=10.5)
ax.legend(loc='lower left', fontsize=8.5, frameon=False)
bx.axhline(0, color=INK2, lw=0.8)
bx.set_ylabel('natural − historical volume (km³)')
bx.set_title('(b) Anthropogenic loss A(t), same state', loc='left', fontsize=10.5)
bx.legend(loc='upper left', fontsize=8.5, frameon=False)
for x in (ax, bx):
    x.set_xlim(1850, T_END)
    x.grid(alpha=0.3)
    x.axvspan(1850, 1900, color='0.5', alpha=0.07, lw=0)
cx.set_xlim(-0.12, max(0.62, max(t[1] for t in have) + 0.12))
cx.set_xticks([t[1] for t in VARIANTS])
cx.set_xlabel('natural climate cooler from 1850 (K)')
cx.set_ylabel('anthropogenic share of the loss 1850–2025 (%)')
cx.set_title('(c) Share 2025, Reinthaler + NAT\n(median, 16–84%)', loc='left', fontsize=10.5)
cx.grid(alpha=0.3)
fig.suptitle('A small shift of the natural counterfactual (shaded in a, b: the 1850–1900 ramp) '
             'and the attribution', x=0.01, ha='left', fontsize=11.5)
fig.savefig(f'{HERE}/figs/nat_shift_effect.png', dpi=150, bbox_inches='tight')
print('wrote figs/nat_shift_effect.png')