"""Figures for the hurricane-outlook audit. Run with /work/venv-plot/bin/python."""
import csv, collections
import matplotlib; matplotlib.use('Agg')
import matplotlib.pyplot as plt, matplotlib.ticker
BLUE, ORANGE, AQUA, YELLOW = '#2a78d6', '#eb6834', '#1baf7a', '#eda100'
INK, MUTED, GRID = '#222222', '#6b6b6b', '#e4e4e4'
plt.rcParams.update({'font.size': 10, 'axes.edgecolor': MUTED, 'axes.labelcolor': INK, 'xtick.color': MUTED,
                     'ytick.color': MUTED, 'axes.spines.top': False, 'axes.spines.right': False})
obs = {int(r['year']): r for r in csv.DictReader(open('../data/season_counts.csv'))}
noaa = list(csv.DictReader(open('../data/noaa_outlooks.csv')))
# Fig 1: ranges vs outcome
fig, axs = plt.subplots(3, 1, figsize=(10, 8.2), sharex=True)
for ax, k, lab in zip(axs, ('ns', 'h', 'mh'), ('Named storms', 'Hurricanes', 'Major hurricanes')):
    ax.axvspan(2000.5, 2007.5, color='#f3f3f3', zorder=0)
    for r in noaa:
        y = int(r['year']); off = -0.17 if r['month'] == 'May' else 0.17
        lo, hi = int(r[k + '_lo']), int(r[k + '_hi']); v = int(obs[y][k])
        col = BLUE if r['month'] == 'May' else AQUA
        ax.plot([y + off, y + off], [lo, hi], color=col, lw=4.5, solid_capstyle='round', zorder=2)
    ys = list(range(2001, 2026)); vs = [int(obs[y][k]) for y in ys]
    ax.scatter(ys, vs, s=34, color=INK, zorder=3, edgecolor='white', linewidth=1)
    ax.yaxis.set_major_locator(matplotlib.ticker.MaxNLocator(integer=True)); ax.set_ylabel(lab); ax.grid(axis='y', color=GRID, lw=0.8); ax.set_axisbelow(True)
axs[0].set_title('Shaded: 2001-2007, "likely" ranges with no probability stated. 2008: 60-70% / 67% stated; from 2009: "70% probability for each range".', loc='left', fontsize=9, color=MUTED)
from matplotlib.lines import Line2D
axs[0].legend(handles=[Line2D([], [], color=BLUE, lw=4.5, label='NOAA May range'), Line2D([], [], color=AQUA, lw=4.5, label='NOAA August range'),
                       Line2D([], [], marker='o', color=INK, lw=0, label='Observed (HURDAT2)')], loc='upper right', bbox_to_anchor=(1, 1.32), frameon=False, fontsize=9, ncol=3)
axs[-1].set_xticks(range(2001, 2026, 2))
fig.tight_layout(rect=(0, 0, 1, 0.94))
fig.suptitle('NOAA Atlantic hurricane outlook ranges and what happened, 2001-2025', x=0.06, y=0.985, ha='left', fontsize=12, color=INK); fig.savefig('../fig/fig1_noaa_ranges.png', dpi=150)
# Fig 2: coverage by era
import numpy as np
sc = list(csv.DictReader(open('../data/noaa_scored.csv')))
fig, ax = plt.subplots(figsize=(7.5, 4))
vars_ = [('ns', 'Named storms'), ('h', 'Hurricanes'), ('mh', 'Major hurricanes'), ('ace', 'ACE')]
for j, (era, sel, col) in enumerate((('2001-2007 ("likely")', lambda d: int(d['year']) < 2008, ORANGE), ('2008-2025 (stated 60-70%)', lambda d: int(d['year']) >= 2008, BLUE))):
    for i, (k, lab) in enumerate(vars_):
        rs = [d[k + '_res'] for d in sc if sel(d) and d[k + '_res']]
        c = sum(x == 'in' for x in rs) / len(rs)
        x = i + (j - 0.5) * 0.36
        ax.bar(x, c, width=0.32, color=col, label=era if i == 0 else None)
        lab_ = f'{sum(x == "in" for x in rs)}/{len(rs)}'
        if c > 0.08: ax.text(x, c - 0.05, lab_, ha='center', fontsize=8.5, color='white')
        else: ax.text(x, c + 0.02, lab_, ha='center', fontsize=8.5, color=INK)
ax.axhline(0.7, color=INK, lw=1.2, ls='--'); ax.text(3.62, 0.715, 'stated 70%', fontsize=8.5, color=INK, ha='right')
ax.set_xticks(range(4)); ax.set_xticklabels([v[1] for v in vars_]); ax.set_ylim(0, 1); ax.set_ylabel('Share of outcomes inside range')
ax.legend(frameon=False, loc='upper left', fontsize=9); ax.grid(axis='y', color=GRID, lw=0.8); ax.set_axisbelow(True)
ax.set_title('How often NOAA\'s ranges contained the outcome (May + August outlooks)', loc='left', fontsize=11, color=INK)
fig.tight_layout(); fig.savefig('../fig/fig2_coverage.png', dpi=150)
# Fig 3: skill by issue window
sk = list(csv.DictReader(open('../data/skill_ci.csv')))
order = ['Dec-Feb', 'Mar-Apr', 'May-Jun', 'Jul', 'Aug']
iss = [('CSU', BLUE), ('TSR', ORANGE), ('NOAA', AQUA), ('NCSU', YELLOW)]
fig, ax = plt.subplots(figsize=(8, 4.4))
for j, (name, col) in enumerate(iss):
    for r in sk:
        if r['issuer'] != name: continue
        x = order.index(r['window']) + (j - 1.5) * 0.14
        e, lo, hi = float(r['era1995']), float(r['era1995_lo90']), float(r['era1995_hi90'])
        ax.plot([x, x], [max(lo, -1), hi], color=col, lw=2)
        ax.scatter([x], [max(e, -1)], s=46, color=col, edgecolor='white', linewidth=1.2, zorder=3, label=name if r['window'] in ('Mar-Apr', 'May-Jun') and name in ('NCSU', 'NOAA') or r['window'] == 'Dec-Feb' else None)
ax.axhline(0, color=INK, lw=1)
ax.set_xticks(range(len(order))); ax.set_xticklabels(['Dec-Feb\n(before season)', 'Mar-Apr', 'May-Jun', 'Jul', 'Aug'])
ax.set_ylabel('MSE skill vs mean of 1995..prior year'); ax.set_ylim(-1, 0.9)
h, l = ax.get_legend_handles_labels(); seen = {}
for hh, ll in zip(h, l): seen.setdefault(ll, hh)
ax.legend([seen[n] for n, _ in iss if n in seen], [n for n, _ in iss if n in seen], frameon=False, loc='upper left', fontsize=9, ncol=4)
ax.text(2.5, -0.04, 'below 0: worse than the plain average of seasons since 1995', fontsize=8.5, color=MUTED, ha='center', va='top')
ax.grid(axis='y', color=GRID, lw=0.8); ax.set_axisbelow(True)
ax.set_title('Seasonal-count forecast skill by issue month, 2001-2025 (NS, H, MH pooled; 90% bootstrap CI)', loc='left', fontsize=10.5, color=INK)
fig.tight_layout(); fig.savefig('../fig/fig3_skill_by_month.png', dpi=150)
print('ok')