"""Figures from results/fig*.csv. Run with /work/venv-plot/bin/python."""
import csv, matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.lines import Line2D
BLUE, ORANGE, AQUA, YELLOW = '#2a78d6', '#eb6834', '#1baf7a', '#eda100'
INK, MUTED, GRID = '#1f1f1f', '#6b6b6b', '#e6e6e6'
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})
R = lambda f: list(csv.DictReader(open(f'results/{f}.csv')))
fl = lambda r, k: float(r[k])
# Fig 1: reliability
fig, ax = plt.subplots(figsize=(6.4, 5.6))
ax.plot([0.0005, 1], [0.0005, 1], color=MUTED, lw=1, ls='--', zorder=1)
ax.text(0.0016, 0.0011, 'perfect calibration', color=MUTED, fontsize=8.5, rotation=40, ha='left')
for s, col, off in (('all versions', BLUE, 1.0), ('initial forecast', ORANGE, 1.06)):
    rows = [r for r in R('fig1_reliability') if r['series'] == s and fl(r, 'obs') > 0]
    for r in rows:
        ax.plot([fl(r, 'meanP') * off] * 2, [max(fl(r, 'ci05'), 6e-4), fl(r, 'ci95')], color=col, lw=2, solid_capstyle='round', zorder=2)
    ax.scatter([fl(r, 'meanP') * off for r in rows], [fl(r, 'obs') for r in rows], s=44, color=col, edgecolor='white', linewidth=1.2, zorder=3)
ax.set_xscale('log'); ax.set_yscale('log'); ax.set_xlim(5e-4, 1.1); ax.set_ylim(5e-4, 1.1)
ax.set_xlabel('Stated probability of at least one aftershock ≥ M (bin mean)'); ax.set_ylabel('Observed frequency (90% CI, resampling mainshocks)')
ax.grid(color=GRID, lw=0.8); ax.set_axisbelow(True)
_r1 = R('fig1_reliability'); _n = lambda s: sum(int(r['n']) for r in _r1 if r['series'] == s)
ax.legend(handles=[Line2D([], [], marker='o', color=BLUE, lw=2, label=f"All forecast versions ({_n('all versions'):,} statements)"),
                   Line2D([], [], marker='o', color=ORANGE, lw=2, label=f"Initial forecast, issued ≤1 h after the mainshock ({_n('initial forecast'):,})")],
          loc='upper left', frameon=False, fontsize=8.5)
ax.set_title('Magnitudes at or above USGS\'s stated catalog completeness;\n1-day, 1-week and 1-month windows. Initial-forecast bin 0-1%\n(0 observed, 2.7 expected) is off the log axis.', loc='left', fontsize=8.5, color=MUTED)
fig.suptitle('USGS aftershock forecasts: stated vs observed, 2019-2026', x=0.02, ha='left', fontsize=12, color=INK)
fig.tight_layout(); fig.savefig('fig/fig1_reliability.png', dpi=150)
# Fig 2: obs/exp by issue lag, shallow vs deep mainshocks
rows = R('fig2_by_issue_lag'); lags = [r['issue_lag'] for r in rows if r['group'] == 'all']
fig, ax = plt.subplots(figsize=(7.6, 4.4))
for grp, col, off in (('shallow (<70 km)', BLUE, -0.12), ('deep (>=70 km)', ORANGE, 0.12)):
    rr = {r['issue_lag']: r for r in rows if r['group'] == grp}
    for i, k in enumerate(lags):
        r = rr[k]
        ax.plot([i + off] * 2, [fl(r, 'ci05'), fl(r, 'ci95')], color=col, lw=4.5, solid_capstyle='round', alpha=0.45)
        ax.scatter([i + off], [fl(r, 'obs_over_exp')], s=46, color=col, edgecolor='white', linewidth=1.2, zorder=3)
        if grp.startswith('shallow'): ax.text(i + off, fl(r, 'ci05') - 0.05, f"{float(r['observed']):.0f} / {float(r['expected']):.0f}", ha='center', va='top', fontsize=7.5, color=MUTED)
        else: ax.text(i + off + 0.06, fl(r, 'obs_over_exp') + 0.04, f"{float(r['observed']):.0f} / {float(r['expected']):.0f}", ha='left', va='bottom', fontsize=7.5, color=MUTED)
ax.axhline(1, color=INK, lw=1, ls='--'); ax.text(-0.45, 1.02, 'calibrated', fontsize=8.5, color=INK, ha='left', va='bottom')
ax.set_xticks(range(len(lags))); ax.set_xticklabels(lags); ax.set_xlabel('Time between the mainshock and the forecast\'s issue')
ax.set_ylabel('Observed ÷ expected\n(sum of outcomes ÷ sum of probabilities)'); ax.set_ylim(-0.05, 1.75); ax.grid(axis='y', color=GRID, lw=0.8); ax.set_axisbelow(True)
ax.legend(handles=[Line2D([], [], marker='o', color=BLUE, lw=4, label='Shallow mainshocks (<70 km, 221)'), Line2D([], [], marker='o', color=ORANGE, lw=4, label='Deep mainshocks (≥70 km, 29)')],
          loc='upper left', frameon=False, fontsize=8.5)
ax.set_title('Below 1 = fewer aftershocks than forecast. Labels: observed / expected. M ≥ completeness,\nwindows ≤ 1 month; bars are 90% CIs resampling mainshocks.', loc='left', fontsize=8.5, color=MUTED)
fig.suptitle('Updates fix most of the over-forecast for shallow quakes, not for deep ones', x=0.02, ha='left', fontsize=11.5, color=INK)
fig.tight_layout(); fig.savefig('fig/fig2_by_issue_lag.png', dpi=150)
# Fig 3: initial forecast by magnitude and region
LAB = {'CONUS+HI': 'Contiguous\nUS + HI', 'Marianas/Guam': 'Marianas/\nGuam', 'deep (>=70 km)': 'Deep\n(≥70 km)', 'all shallow (<70 km)': 'All shallow\n(<70 km)'}
rows = R('fig3_initial_by_M_region')
fig, axs = plt.subplots(1, 2, figsize=(10.5, 4.3), sharey=True, gridspec_kw={'width_ratios': [4, 6]})
for ax, panel, title in ((axs[0], 'magnitude', 'By aftershock magnitude threshold'), (axs[1], 'region', 'By region, and by mainshock depth')):
    rr = [r for r in rows if r['panel'] == panel]
    for i, r in enumerate(rr):
        ax.plot([i, i], [fl(r, 'ci05'), fl(r, 'ci95')], color=ORANGE, lw=4.5, solid_capstyle='round', alpha=0.45)
        ax.scatter([i], [fl(r, 'obs_over_exp')], s=46, color=ORANGE, edgecolor='white', linewidth=1.2, zorder=3)
        ax.text(i, max(fl(r, 'ci95'), fl(r, 'obs_over_exp')) + 0.06, f"{float(r['observed']):.0f} / {float(r['expected']):.1f}", ha='center', fontsize=8, color=MUTED)
    ax.axhline(1, color=INK, lw=1, ls='--'); ax.set_xticks(range(len(rr))); ax.set_xticklabels([LAB.get(r['group'], r['group']) for r in rr], fontsize=8.5)
    ax.set_title(title, loc='left', fontsize=9.5, color=INK); ax.set_xlim(-0.7, len(rr) - 0.4)
    if panel == 'region': ax.axvline(3.5, color=GRID, lw=1.5); ax.grid(axis='y', color=GRID, lw=0.8); ax.set_axisbelow(True); ax.set_xlim(-0.5, len(rr) - 0.5)
axs[0].set_ylabel('Observed ÷ expected'); axs[0].set_ylim(0, 1.6)
fig.suptitle('Initial forecasts (≤1 h after the mainshock): observed / expected aftershock counts', x=0.02, ha='left', fontsize=11.5, color=INK)
fig.text(0.02, 0.01, 'Labels: observed / expected, summed over 1-day, 1-week and 1-month windows (nested, so not independent). M ≥ completeness only;\nM3+ counts only where USGS states completeness ≤ 3. Marianas/Guam and deep: 0 observed, so no interval is drawn.', fontsize=7.8, color=MUTED)
fig.tight_layout(rect=(0, 0.07, 1, 1)); fig.savefig('fig/fig3_initial_by_M_region.png', dpi=150)
# Fig 4: Brier skill
rows = R('fig4_skill'); wins = ['1 Day', '1 Week', '1 Month', '1 Year']
fig, ax = plt.subplots(figsize=(7.2, 4.2))
for j, (ref, col, off) in enumerate((('generic model, not updated', AQUA, -0.12), ('background rate', YELLOW, 0.12))):
    rr = {r['window']: r for r in rows if r['reference'] == ref}
    for i, w in enumerate(wins):
        r = rr[w]; ax.plot([i + off] * 2, [fl(r, 'ci05'), fl(r, 'ci95')], color=col, lw=4.5, solid_capstyle='round', alpha=0.5)
        ax.scatter([i + off], [fl(r, 'brier_skill')], s=46, color=col, edgecolor='white', linewidth=1.2, zorder=3)
ax.axhline(0, color=INK, lw=1); ax.set_xticks(range(4)); ax.set_xticklabels(wins); ax.set_xlabel('Forecast window')
ax.set_ylabel('Brier skill score of the published forecast'); ax.grid(axis='y', color=GRID, lw=0.8); ax.set_axisbelow(True)
ax.legend(handles=[Line2D([], [], marker='o', color=AQUA, lw=4, label='vs the same model with generic parameters, never updated'),
                   Line2D([], [], marker='o', color=YELLOW, lw=4, label='vs the 10-year pre-mainshock rate in the same circle')],
          loc='lower left', frameon=False, fontsize=8.5)
ax.set_title('Above 0 = published forecast better. All versions, M ≥ completeness;\n90% CIs resampling mainshocks.', loc='left', fontsize=8.5, color=MUTED)
fig.suptitle('Updating helps at every window; 1-year vs background is inconclusive', x=0.02, ha='left', fontsize=11.5, color=INK)
fig.tight_layout(); fig.savefig('fig/fig4_skill.png', dpi=150)
print('ok')