from __future__ import annotations
import json
import re
from pathlib import Path
import numpy as np
ROOT = Path(__file__).resolve().parent
SEED_RE = re.compile('_seed(\\d+)$')
EXISTING_SOTA = {'ISCXTor2016 (NonTor → Tor)': {'shafir_baseline': 0.8731, 'shafir_ref': 'Table VI', 'ours_existing': (0.9945, 0.0011), 'ours_score': 'terminal_norm', 'sigma': 0.1, 'ac_prefix': 'route_ac_combo_iscxtor2016'}, 'CICIDS2017 within (Shafir 10k/10k)': {'shafir_baseline': 0.9303, 'shafir_ref': 'Table VII', 'ours_existing': (0.9858, 0.0021), 'ours_score': 'terminal_norm', 'sigma': 0.6, 'ac_prefix': 'route_ac_combo_cicids2017'}, 'CICDDoS2019 within': {'shafir_baseline': 0.93, 'shafir_ref': 'Table IX, row 1', 'ours_existing': (0.996, 0.001), 'ours_score': 'terminal_norm', 'sigma': 0.1, 'ac_prefix': 'route_ac_combo_cicddos2019'}, 'CICIoT2023 within (multi-seed)': {'shafir_baseline': None, 'shafir_ref': None, 'ours_existing': (0.9612, 0.0017), 'ours_score': 'terminal_norm', 'sigma': 0.1, 'ac_prefix': 'route_ac_combo_ciciot2023'}}

def _seeds(prefix: str) -> dict[int, Path]:
    out = {}
    for d in sorted(ROOT.glob(f'{prefix}_seed*')):
        m = SEED_RE.search(d.name)
        if m and (d / 'phase1_summary.json').exists():
            out[int(m.group(1))] = d
    return out

def _load(d: Path) -> dict:
    return json.loads((d / 'phase1_summary.json').read_text())

def _mean_std(vs: list[float]) -> tuple[float, float]:
    arr = np.asarray([v for v in vs if v == v], dtype=np.float64)
    if arr.size == 0:
        return (float('nan'), float('nan'))
    return (float(arr.mean()), float(arr.std()))

def main() -> None:
    rows: list[str] = []
    rows.append('# SOTA Comparison: A+C combo vs existing UnifiedCFM')
    rows.append('')
    rows.append('All 4 datasets, 3 seeds each, within-dataset Shafir 10K/10K protocol.')
    rows.append('Existing UnifiedCFM uses Phase-2 consistency loss (λ_flow=λ_packet=0.3).')
    rows.append('A+C combo uses Mixed_CFM (continuous CFM + DFM) + causal-packet attention,')
    rows.append('**no Phase-2 consistency loss**. lambda_disc=1.0, sigma=0.1, use_ot=True.')
    rows.append('')
    rows.append("## Headline: A+C combo's best score per dataset")
    rows.append('')
    rows.append('| Dataset | Shafir 2026 | Existing UnifiedCFM (SOTA) | A+C combo `terminal_norm` | A+C combo `terminal_packet` | A+C combo `disc_nll_total` | A+C best | New SOTA? |')
    rows.append('|---|---|---|---|---|---|---|---|')
    for (label, meta) in EXISTING_SOTA.items():
        seeds = _seeds(meta['ac_prefix'])
        shafir_str = f"{meta['shafir_baseline']:.4f}" if meta['shafir_baseline'] else '—'
        (existing_m, existing_sd) = meta['ours_existing']
        existing_str = f'{existing_m:.4f} ± {existing_sd:.4f}'
        if not seeds:
            rows.append(f'| {label} | {shafir_str} | {existing_str} | (running) | — | — | — | — |')
            continue
        vals_term = [_load(d).get('overall', {}).get('terminal_norm', {}).get('auroc', float('nan')) for d in seeds.values()]
        vals_pkt = [_load(d).get('overall', {}).get('terminal_packet', {}).get('auroc', float('nan')) for d in seeds.values()]
        vals_disc = [_load(d).get('overall', {}).get('disc_nll_total', {}).get('auroc', float('nan')) for d in seeds.values()]
        (m_t, s_t) = _mean_std(vals_term)
        (m_p, s_p) = _mean_std(vals_pkt)
        (m_d, s_d) = _mean_std(vals_disc)
        (best_score, best_m, best_sd) = ('terminal_norm', m_t, s_t)
        if m_p > best_m:
            (best_score, best_m, best_sd) = ('terminal_packet', m_p, s_p)
        if m_d > best_m:
            (best_score, best_m, best_sd) = ('disc_nll_total', m_d, s_d)
        beats = '✅' if best_m > existing_m else '❌'
        rows.append(f'| {label} | {shafir_str} | {existing_str} | {m_t:.4f} ± {s_t:.4f} | {m_p:.4f} ± {s_p:.4f} | {m_d:.4f} ± {s_d:.4f} | `{best_score}` {best_m:.4f} ± {best_sd:.4f} | {beats} {best_m - existing_m:+.4f} |')
    rows.append('')
    rows.append('## Per-dataset full scoring')
    rows.append('')
    score_keys = ['terminal_norm', 'terminal_flow', 'terminal_packet', 'disc_nll_total', 'disc_nll_ch3', 'disc_nll_ch4', 'disc_nll_ch5', 'disc_nll_ch7']
    for (label, meta) in EXISTING_SOTA.items():
        rows.append(f'### {label}')
        rows.append('')
        seeds = _seeds(meta['ac_prefix'])
        if not seeds:
            rows.append('(not yet completed)\n')
            continue
        rows.append('| Score | mean ± std | seeds |')
        rows.append('|---|---|---|')
        for sc in score_keys:
            vals = [_load(d).get('overall', {}).get(sc, {}).get('auroc', float('nan')) for d in seeds.values()]
            (m, sd) = _mean_std(vals)
            if m == m:
                rows.append(f'| `{sc}` | {m:.4f} ± {sd:.4f} | {sorted(seeds.keys())} |')
        rows.append('')
    out = ROOT / 'SOTA_COMPARISON.md'
    out.write_text('\n'.join(rows) + '\n')
    print(f'[wrote] {out}')
if __name__ == '__main__':
    main()
