from __future__ import annotations
import json
from pathlib import Path
import numpy as np
REPO = Path(__file__).resolve().parents[2]
ROOT = REPO / 'artifacts/baselines/shafir_nf_2026_04_29'
PROTOCOLS = ('iscxtor_within', 'cicids_within', 'cicddos_within', 'forward_cross', 'reverse_cross')
SEEDS = (42, 43, 44)
TERMINAL_NORM = {'iscxtor_within': (0.9945, 0.0011), 'cicids_within': (0.9858, 0.0021), 'cicddos_within': (0.996, 0.001), 'forward_cross': (0.9109, 0.0032), 'reverse_cross': (0.5999, None)}
SHAFIR_PAPER = {'iscxtor_within': (0.8731, None), 'cicids_within': (0.9303, None), 'cicddos_within': (0.93, None), 'forward_cross': (0.89, None), 'reverse_cross': (0.93, None)}
PRETTY = {'iscxtor_within': 'ISCXTor2016 within', 'cicids_within': 'CICIDS2017 within (σ=0.6)', 'cicddos_within': 'CICDDoS2019 within', 'forward_cross': 'IDS2017→DDoS2019 forward', 'reverse_cross': 'DDoS2019→IDS2017 reverse'}

def _load(protocol, seed):
    p = ROOT / f'{protocol}_seed{seed}.json'
    if not p.exists():
        return None
    return json.loads(p.read_text())

def _ms(vals):
    arr = np.asarray([v for v in vals if v is not None and (not np.isnan(v))], dtype=np.float64)
    if len(arr) == 0:
        return (float('nan'), float('nan'))
    return (float(arr.mean()), float(arr.std(ddof=1)) if len(arr) > 1 else 0.0)

def main():
    rows = []
    per_class_collect = {p: {} for p in PROTOCOLS}
    for protocol in PROTOCOLS:
        (aurocs, auprcs, t_train) = ([], [], [])
        for s in SEEDS:
            r = _load(protocol, s)
            if r is None:
                continue
            aurocs.append(r['overall']['neg_log_prob']['auroc'])
            auprcs.append(r['overall']['neg_log_prob']['auprc'])
            t_train.append(r.get('t_train_sec', 0.0))
            for (cls, info) in r.get('per_class', {}).items():
                per_class_collect[protocol].setdefault(cls, {'n': int(info['_n']), 'aurocs': []})
                per_class_collect[protocol][cls]['aurocs'].append(info['auroc'])
        (m, sd) = _ms(aurocs)
        (ma, sda) = _ms(auprcs)
        (tt, _) = _ms(t_train)
        rows.append({'protocol': protocol, 'n_seeds': len(aurocs), 'auroc_mean': m, 'auroc_std': sd, 'auprc_mean': ma, 'auprc_std': sda, 't_train_sec_mean': tt})
    lines = ['# Shafir 2026 NF Baseline — On Our 5-Protocol Layout', '', 'Date: 2026-04-29', '', "Method: Shafir's official `pzflow.Flow` (single basic NF).", 'Features: our **20-d canonical packet-derived flow features** (`common.data_contract.CANONICAL_FLOW_FEATURE_NAMES`), z-scored with the **same source training stats** that the Unified_CFM checkpoint uses.', 'Train cap: 10,000 source-benign samples (Shafir paper protocol).', 'Optimizer: SGD lr=1e-3, 100 epochs (Shafir paper defaults).', 'Sampling: same seeds & stratification as `eval_new_scores.py`.', '', '## Headline AUROC (3-seed mean ± std)', '', '| Protocol | terminal_norm (ours) | Shafir NF — paper | **Shafir NF — our features** | Δ vs paper | Δ vs terminal_norm |', '|---|---:|---:|---:|---:|---:|']
    for row in rows:
        p = row['protocol']
        (tn_m, tn_sd) = TERMINAL_NORM[p]
        (sp_m, _) = SHAFIR_PAPER[p]
        (m, sd) = (row['auroc_mean'], row['auroc_std'])
        if np.isnan(m):
            lines.append(f'| {PRETTY[p]} | {tn_m:.4f} | {sp_m:.4f} | (no runs yet) | — | — |')
            continue
        d_paper = m - sp_m
        d_terminal = m - tn_m
        tn_str = f'{tn_m:.4f} ± {tn_sd:.4f}' if tn_sd is not None else f'{tn_m:.4f}'
        lines.append(f'| {PRETTY[p]} | {tn_str} | {sp_m:.4f} | **{m:.4f} ± {sd:.4f}** | {d_paper:+.4f} | {d_terminal:+.4f} |')
    lines.append('')
    lines.append('## Per-protocol stats')
    lines.append('')
    lines.append('| Protocol | n_seeds | AUPRC mean ± std | Train time (s, mean) |')
    lines.append('|---|---:|---:|---:|')
    for row in rows:
        p = row['protocol']
        (m, sd) = (row['auprc_mean'], row['auprc_std'])
        if np.isnan(m):
            continue
        lines.append(f"| {PRETTY[p]} | {row['n_seeds']} | {m:.4f} ± {sd:.4f} | {row['t_train_sec_mean']:.1f} |")
    lines.append('')
    lines.append('## Per-attack (forward + reverse)')
    for protocol in ('forward_cross', 'reverse_cross'):
        lines.append(f'\n### {PRETTY[protocol]}')
        d = per_class_collect[protocol]
        if not d:
            lines.append('(no runs)')
            continue
        lines.append('| attack | n | Shafir NF AUROC mean ± std |')
        lines.append('|---|---:|---:|')
        for cls in sorted(d):
            n = d[cls]['n']
            (m, sd) = _ms(d[cls]['aurocs'])
            lines.append(f'| `{cls}` | {n} | {m:.4f} ± {sd:.4f} |')
    out = ROOT / 'summary.md'
    out.write_text('\n'.join(lines))
    summary_json = {'rows': rows, 'per_class': {p: {cls: {'n': v['n'], **dict(zip(['mean', 'std'], _ms(v['aurocs'])))} for (cls, v) in dd.items()} for (p, dd) in per_class_collect.items()}, 'baselines': {'terminal_norm': TERMINAL_NORM, 'shafir_paper': SHAFIR_PAPER}}
    (ROOT / 'summary.json').write_text(json.dumps(summary_json, indent=2))
    print(f'[saved] {out}')
    print(f"[saved] {ROOT / 'summary.json'}")
    print()
    for row in rows:
        if not np.isnan(row['auroc_mean']):
            print(f"  {PRETTY[row['protocol']]:<34s} {row['auroc_mean']:.4f} ± {row['auroc_std']:.4f}")
if __name__ == '__main__':
    main()
