from __future__ import annotations
import json
import re
from pathlib import Path
import numpy as np
from sklearn.metrics import roc_auc_score
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'), ('A+C combo', '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_summary(d: Path) -> dict:
    return json.loads((d / 'phase1_summary.json').read_text())

def _ensemble_sweep(d: Path) -> dict[float, float] | None:
    f = d / 'phase1_scores.npz'
    if not f.exists():
        return None
    z = np.load(f, allow_pickle=True)
    keys = set(z.files)
    if 'val_terminal_norm' not in keys or 'val_disc_nll_total' not in keys:
        return None
    v_tn = z['val_terminal_norm']
    a_tn = z['atk_terminal_norm']
    v_dn = z['val_disc_nll_total']
    a_dn = z['atk_disc_nll_total']

    def zsc(v, a):
        (mu, sd) = (v.mean(), v.std() + 1e-09)
        return ((v - mu) / sd, (a - mu) / sd)
    (v_tn_z, a_tn_z) = zsc(v_tn, a_tn)
    (v_dn_z, a_dn_z) = zsc(v_dn, a_dn)
    out: dict[float, float] = {}
    for alpha in (0.0, 0.25, 0.5, 0.7, 0.8, 0.9, 1.0):
        s_v = alpha * v_tn_z + (1.0 - alpha) * v_dn_z
        s_a = alpha * a_tn_z + (1.0 - alpha) * a_dn_z
        y = np.r_[np.zeros(len(s_v)), np.ones(len(s_a))]
        s = np.r_[s_v, s_a]
        out[alpha] = float(roc_auc_score(y, s))
    return out

def _ensemble_score(d: Path) -> tuple[float, float] | None:
    sweep = _ensemble_sweep(d)
    if sweep is None:
        return None
    best_alpha = max(sweep, key=sweep.get)
    return (sweep[best_alpha], best_alpha)

def _mean_std(vals: list[float]) -> tuple[float, float]:
    arr = np.asarray([v for v in vals 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: dict[str, dict[int, dict]] = {}
    routes_dirs: dict[str, dict[int, Path]] = {}
    for (label, prefix) in ROUTES:
        seeds = _seeds(prefix)
        routes_dirs[label] = seeds
        routes_data[label] = {s: _load_summary(d) for (s, d) in seeds.items()}
    rows: list[str] = []
    rows.append('# Route Comparison Results — CICIoT2023 (multi-seed)')
    rows.append('')
    rows.append('Phase1 eval: AUROC over benign val (5k cap) vs all attacks (10k cap), 3 seeds each.')
    rows.append('')
    rows.append("## Each route's best AUROC (overall)")
    rows.append('')
    rows.append('| Route | Best score | AUROC | Δ vs baseline-best |')
    rows.append('|---|---|---|---|')
    baseline_best = None
    for (label, _) in ROUTES:
        seeds = routes_data[label]
        if not seeds:
            rows.append(f'| {label} | — | — | — |')
            continue
        all_scores: dict[str, list[float]] = {}
        for s in seeds.values():
            for (k, v) in s.get('overall', {}).items():
                all_scores.setdefault(k, []).append(v.get('auroc', float('nan')))
        score_means = {k: _mean_std(v)[0] for (k, v) in all_scores.items()}
        score_means = {k: v for (k, v) in score_means.items() if v == v}
        if not score_means:
            rows.append(f'| {label} | — | — | — |')
            continue
        best_score = max(score_means, key=score_means.get)
        best_val = score_means[best_score]
        if label == 'baseline':
            baseline_best = best_val
            delta_str = '—'
        else:
            delta_str = f'{best_val - baseline_best:+.4f}' if baseline_best else '—'
        std = _mean_std(all_scores[best_score])[1]
        rows.append(f'| {label} | `{best_score}` | {best_val:.4f} ± {std:.4f} | {delta_str} |')
    rows.append('')
    rows.append('## Primary score: `terminal_norm`')
    rows.append('')
    rows.append('| Route | mean ± std | seeds |')
    rows.append('|---|---|---|')
    for (label, _) in ROUTES:
        seeds = routes_data[label]
        if not seeds:
            rows.append(f'| {label} | — | — |')
            continue
        vals = [s['overall'].get('terminal_norm', {}).get('auroc', float('nan')) for s in seeds.values()]
        (m, sd) = _mean_std(vals)
        rows.append(f'| {label} | {m:.4f} ± {sd:.4f} | {sorted(seeds.keys())} |')
    rows.append('')
    rows.append('## Route-specific signature scores (mean ± std, 3 seeds)')
    rows.append('')
    score_groups = [('Route A signature (consistency family)', ['flow_consistency', 'packet_consistency', 'consistency_total', 'causal_surprisal_total', 'causal_surprisal_packet_median']), ('Route B signature (curvature/dynamics)', ['kappa2_speed2norm_packet_median', 'direction_drift_packet_median', 'pna_packet_median', 'curvature_packet']), ('Route C signature (discrete NLL)', ['disc_nll_total', 'disc_nll_ch3', 'disc_nll_ch4', 'disc_nll_ch5', 'disc_nll_ch7'])]
    for (grp_name, scores) in score_groups:
        rows.append(f'### {grp_name}')
        rows.append('')
        rows.append('| Score | ' + ' | '.join((label for (label, _) in ROUTES)) + ' |')
        rows.append('|---' * (1 + len(ROUTES)) + '|')
        for sc in scores:
            cells = [f'`{sc}`']
            for (label, _) in ROUTES:
                seeds = routes_data[label]
                vals = [s['overall'].get(sc, {}).get('auroc', float('nan')) for s in seeds.values()]
                (m, sd) = _mean_std(vals)
                cells.append(f'{m:.4f} ± {sd:.4f}' if m == m else '—')
            rows.append('| ' + ' | '.join(cells) + ' |')
        rows.append('')
    rows.append('## Route C ensemble: α·terminal_norm + (1−α)·disc_nll_total (z-scored)')
    rows.append('')
    c_dirs = routes_dirs.get('C: mixed', {})
    if c_dirs:
        alphas = (0.0, 0.25, 0.5, 0.7, 0.8, 0.9, 1.0)
        rows.append('| α | ' + ' | '.join((f'seed{s}' for s in sorted(c_dirs.keys()))) + ' | mean ± std |')
        rows.append('|---' * (2 + len(c_dirs)) + '|')
        per_alpha: dict[float, list[float]] = {a: [] for a in alphas}
        per_seed_sweeps = {s: _ensemble_sweep(d) or {} for (s, d) in c_dirs.items()}
        for a in alphas:
            cells = [f'{a:.2f}']
            vals = []
            for s in sorted(c_dirs.keys()):
                v = per_seed_sweeps[s].get(a, float('nan'))
                cells.append(f'{v:.4f}')
                vals.append(v)
            (m, sd) = _mean_std(vals)
            cells.append(f'**{m:.4f} ± {sd:.4f}**')
            rows.append('| ' + ' | '.join(cells) + ' |')
        rows.append('')
        rows.append('(α=1.0 = terminal_norm only; α=0.0 = disc_nll only.)')
    rows.append('')
    rows.append('## Per-attack-class AUROC (top 12, terminal_norm)')
    rows.append('')
    if routes_data['baseline']:
        any_summary = next(iter(routes_data['baseline'].values()))
        classes = sorted([(c, d.get('_n', 0)) for (c, d) in any_summary.get('per_class', {}).items() if not c.startswith('_')], 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 classes:
            cells = [cls, f'{int(n)}']
            for (label, _) in ROUTES:
                seeds = routes_data[label]
                vals = [s.get('per_class', {}).get(cls, {}).get('terminal_norm', float('nan')) for s in seeds.values()]
                (m, sd) = _mean_std(vals)
                cells.append(f'{m:.3f}±{sd:.3f}' if m == m 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()
