from __future__ import annotations
import json
import re
from collections import defaultdict
from pathlib import Path
import numpy as np
ROOT = Path(__file__).resolve().parent
SEED_RE = re.compile('_seed(\\d+)$')
ROUTES = [('baseline', 'baseline_ciciot2023'), ('A: causal', 'route_a_causal_ciciot2023'), ('B: spectral', 'route_b_spectral_ciciot2023'), ('C: mixed', 'route_c_mixed_ciciot2023')]
PRIMARY_SCORES = ['terminal_norm', 'terminal_flow', 'terminal_packet', 'causal_surprisal_packet_median', 'causal_surprisal_packet_max', 'causal_surprisal_total', 'consistency_total', 'flow_consistency', 'packet_consistency', 'kappa2_speed2norm_packet_median', 'direction_drift_packet_median', 'pna_packet_median', 'disc_nll_total', 'disc_nll_ch2', 'disc_nll_ch3', 'disc_nll_ch4', 'disc_nll_ch5', 'disc_nll_ch6', 'disc_nll_ch7']

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

def _mean_std(values: list[float]) -> tuple[float, float]:
    arr = np.asarray([v for v in values 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:
    routes_data = {label: _collect(prefix) for (label, prefix) in ROUTES}
    rows = []
    rows.append('# Route Comparison Results — CICIoT2023')
    rows.append('')
    rows.append('All routes trained on CICIoT2023 with the protocol locked in `PROTOCOL.md`. ')
    rows.append('Numbers are AUROC over benign val (10k cap) vs all attacks (10k cap), ')
    rows.append('3 seeds each. ± std across seeds.')
    rows.append('')
    rows.append('## Overall AUROC by score')
    rows.append('')
    header = '| Score | ' + ' | '.join((label for (label, _) in ROUTES)) + ' |'
    sep = '|---' * (1 + len(ROUTES)) + '|'
    rows.append(header)
    rows.append(sep)
    for score in PRIMARY_SCORES:
        cells = [f'`{score}`']
        for (label, _) in ROUTES:
            seeds = routes_data[label]
            if not seeds:
                cells.append('—')
                continue
            vals = [summary.get('overall', {}).get(score, {}).get('auroc', float('nan')) for summary in seeds.values()]
            (mean, std) = _mean_std(vals)
            cells.append(f'{mean:.4f} ± {std:.4f}' if mean == mean else '—')
        rows.append('| ' + ' | '.join(cells) + ' |')
    rows.append('')
    rows.append('## Per-attack-class `terminal_norm` AUROC (top 12 by support)')
    rows.append('')
    seed_dicts = list(routes_data['baseline'].values())
    if seed_dicts:
        all_classes: dict[str, float] = {}
        for s in seed_dicts:
            for (cls, cls_data) in s.get('per_class', {}).items():
                if cls.startswith('_'):
                    continue
                n = cls_data.get('_n', 0.0)
                all_classes[cls] = max(all_classes.get(cls, 0.0), n)
        ranked = sorted(all_classes.items(), key=lambda kv: -kv[1])[:12]
        header = '| Class | n | ' + ' | '.join((label for (label, _) in ROUTES)) + ' |'
        sep = '|---' * (2 + len(ROUTES)) + '|'
        rows.append(header)
        rows.append(sep)
        for (cls, n) in ranked:
            cells = [cls, f'{int(n)}']
            for (label, _) in ROUTES:
                seeds = routes_data[label]
                if not seeds:
                    cells.append('—')
                    continue
                vals = [summary.get('per_class', {}).get(cls, {}).get('terminal_norm', float('nan')) for summary in seeds.values()]
                (mean, std) = _mean_std(vals)
                cells.append(f'{mean:.3f} ± {std:.3f}' if mean == mean else '—')
            rows.append('| ' + ' | '.join(cells) + ' |')
    rows.append('')
    rows.append('## Run inventory')
    rows.append('')
    for (label, prefix) in ROUTES:
        seeds = sorted(routes_data[label].keys())
        rows.append(f"- **{label}** (`{prefix}_seed*`): seeds = {(seeds if seeds else '(none yet)')}")
    out = ROOT / 'RESULTS.md'
    out.write_text('\n'.join(rows) + '\n')
    print(f'[wrote] {out}')
if __name__ == '__main__':
    main()
