from __future__ import annotations
import json
from pathlib import Path
import numpy as np
REPO = Path(__file__).resolve().parents[2]
ROOT = REPO / 'artifacts/baselines/kitsune_2026_04_29'
PROTOCOLS = ('iscxtor_within', 'cicids_within', 'cicddos_within', 'forward_cross', 'reverse_cross')
SEEDS = (42, 43, 44)
AGGS = ('mean', 'max', 'median', 'p90')
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)}
KITSUNE_PAPER = {'iscxtor_within': (0.78, None), 'cicids_within': (0.85, None), 'cicddos_within': (None, None), 'forward_cross': (None, None), 'reverse_cross': (None, 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}
    full = {'protocols': {}}
    for protocol in PROTOCOLS:
        agg_aurocs = {agg: [] for agg in AGGS}
        agg_auprcs = {agg: [] for agg in AGGS}
        seeds_run = []
        for s in SEEDS:
            r = _load(protocol, s)
            if r is None:
                continue
            seeds_run.append(s)
            for agg in AGGS:
                ov = r['overall_by_agg'][agg]
                agg_aurocs[agg].append(ov['auroc'])
                agg_auprcs[agg].append(ov['auprc'])
            for (cls, info) in r.get('per_class_by_agg', {}).get('mean', {}).items():
                per_class_collect[protocol].setdefault(cls, {'n': int(info['_n']), 'aurocs': []})
                per_class_collect[protocol][cls]['aurocs'].append(info['auroc'])
        agg_summary = {}
        for agg in AGGS:
            (m, sd) = _ms(agg_aurocs[agg])
            (ma, sda) = _ms(agg_auprcs[agg])
            agg_summary[agg] = {'auroc_mean': m, 'auroc_std': sd, 'auprc_mean': ma, 'auprc_std': sda}
        full['protocols'][protocol] = {'seeds': seeds_run, 'by_agg': agg_summary}
        best_agg = max(agg_summary, key=lambda a: agg_summary[a]['auroc_mean'])
        rows.append({'protocol': protocol, 'n_seeds': len(seeds_run), 'best_agg': best_agg, 'auroc_mean': agg_summary[best_agg]['auroc_mean'], 'auroc_std': agg_summary[best_agg]['auroc_std'], 'all_aggs': agg_summary})
    lines = ['# Kitsune (Path B) Baseline — On Our 5-Protocol Layout', '', 'Date: 2026-04-29', '', 'Method: KitNET ensemble autoencoder (the ML core of Kitsune).', "**Path B**: feeds our **z-scored 9-d packet features** directly through `KitNET.process()` for the FM+AD grace, then `KitNET.execute()` per packet during eval. **AfterImage's 100-d host/session statistics are skipped** (they require sequential pcap streams which our (B,T,9) tensor abstraction discards). This keeps data usage unified with `eval_new_scores.py`.", 'Train: 5000 source-benign flows → ~75-320k packets (≥ FM+AD=55k grace).', 'Score: per-flow aggregate of per-packet RMSE (mean / max / median / p90).', 'Sampling: same seeds & stratification as `eval_new_scores.py`.', '', '## Headline AUROC (best aggregator per protocol, 3-seed mean ± std)', '', '| Protocol | terminal_norm | Kitsune paper (Shafir reproduction) | **Kitsune Path B (ours)** | best agg | Δ vs paper | Δ vs terminal |', '|---|---:|---:|---:|---|---:|---:|']
    for row in rows:
        p = row['protocol']
        (tn_m, tn_sd) = TERMINAL_NORM[p]
        (kp_m, _) = KITSUNE_PAPER[p]
        (m, sd) = (row['auroc_mean'], row['auroc_std'])
        if np.isnan(m):
            lines.append(f'| {PRETTY[p]} | {tn_m:.4f} | — | (no runs) | — | — | — |')
            continue
        tn_str = f'{tn_m:.4f} ± {tn_sd:.4f}' if tn_sd is not None else f'{tn_m:.4f}'
        kp_str = f'{kp_m:.4f}' if kp_m is not None else '—'
        d_terminal = m - tn_m
        d_paper = m - kp_m if kp_m is not None else None
        d_paper_str = f'{d_paper:+.4f}' if d_paper is not None else '—'
        lines.append(f"| {PRETTY[p]} | {tn_str} | {kp_str} | **{m:.4f} ± {sd:.4f}** | `{row['best_agg']}` | {d_paper_str} | {d_terminal:+.4f} |")
    lines.append('')
    lines.append('## All aggregators (3-seed mean ± std)')
    lines.append('')
    lines.append('| Protocol | mean | max | median | p90 |')
    lines.append('|---|---:|---:|---:|---:|')
    for row in rows:
        cells = [PRETTY[row['protocol']]]
        for agg in AGGS:
            a = row['all_aggs'][agg]
            m = a['auroc_mean']
            if np.isnan(m):
                cells.append('—')
            else:
                cells.append(f"{m:.4f} ± {a['auroc_std']:.4f}")
        lines.append('| ' + ' | '.join(cells) + ' |')
    lines.append('')
    lines.append('## Per-attack (forward + reverse, mean aggregator)')
    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 | Kitsune 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': per_class_collect, 'baselines': {'terminal_norm': TERMINAL_NORM, 'kitsune_paper': KITSUNE_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} best={row['best_agg']:<6s} {row['auroc_mean']:.4f} ± {row['auroc_std']:.4f}")
if __name__ == '__main__':
    main()
