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
CROSS_DIR = ROOT / 'cross'
NAME_RE = re.compile('^(?P<route>.+?)_seed(?P<seed>\\d+)_to_(?P<target>cicids2017|cicddos2019)$')
ROUTES = [('baseline', 'baseline'), ('A: causal', 'route_a_causal'), ('B: spectral', 'route_b_spectral'), ('C: mixed', 'route_c_mixed'), ('A+C combo', 'route_ac_combo')]
TARGETS = ['cicids2017', 'cicddos2019']
PRIMARY_SCORES = ['terminal_norm', 'terminal_flow', 'terminal_packet', 'flow_consistency', 'packet_consistency', 'consistency_total', 'causal_surprisal_packet_median', 'causal_surprisal_total', 'direction_drift_packet_median', 'pna_packet_median', 'kappa2_speed2norm_packet_median', 'curvature_packet', 'disc_nll_total', 'disc_nll_ch3', 'disc_nll_ch7']

def _collect() -> dict[tuple[str, str], dict[int, dict]]:
    out: dict[tuple[str, str], dict[int, dict]] = {}
    for f in sorted(CROSS_DIR.glob('*.json')):
        m = NAME_RE.match(f.stem)
        if not m:
            continue
        key = (m.group('route'), m.group('target'))
        out.setdefault(key, {})[int(m.group('seed'))] = json.loads(f.read_text())
    return out

def _ensemble_sweep(npz_path: Path) -> dict[float, float] | None:
    if not npz_path.exists():
        return None
    z = np.load(npz_path, allow_pickle=True)
    keys = set(z.files)
    if 'b_terminal_norm' not in keys or 'b_disc_nll_total' not in keys:
        return None
    v_tn = z['b_terminal_norm']
    a_tn = z['a_terminal_norm']
    v_dn = z['b_disc_nll_total']
    a_dn = z['a_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 = {}
    for alpha in (0.0, 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 _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:
    data = _collect()
    rows: list[str] = []
    rows.append('# Cross-Dataset Eval — CICIoT2023 → {CICIDS2017, CICDDoS2019}')
    rows.append('')
    rows.append("All models trained on CICIoT2023 (10K benign), evaluated on each target's")
    rows.append('10K benign + 10K stratified attack. Source-domain norm stats applied.')
    rows.append('3 seeds each. AUROC mean ± std.')
    rows.append('')
    rows.append('## Primary score: `terminal_norm`')
    rows.append('')
    header = '| Route | within-CICIoT2023 (ref) | → CICIDS2017 | → CICDDoS2019 |'
    rows.append(header)
    rows.append('|---|---|---|---|')
    within_fallback = {'baseline': (0.9612, 0.0017), 'A: causal': (0.9636, 0.0006), 'B: spectral': (0.9619, 0.0013), 'C: mixed': (0.9625, 0.0028), 'A+C combo': (0.9587, 0.0017)}
    within_terminal: dict[str, tuple[float, float]] = {}
    for (label, prefix) in ROUTES:
        within_seeds = sorted(ROOT.glob(f'{prefix}_seed*/phase1_summary.json'))
        vals: list[float] = []
        for f in within_seeds:
            try:
                s = json.loads(f.read_text())
                v = s.get('overall', {}).get('terminal_norm', {}).get('auroc')
                if v is not None:
                    vals.append(v)
            except Exception:
                pass
        if vals:
            within_terminal[label] = _mean_std(vals)
        else:
            within_terminal[label] = within_fallback.get(label, (float('nan'), float('nan')))
    for (label, prefix) in ROUTES:
        cells = [label]
        (wm, ws) = within_terminal[label]
        cells.append(f'{wm:.4f} ± {ws:.4f}')
        for tgt in TARGETS:
            seeds = data.get((prefix, tgt), {})
            vals = [s['overall'].get('terminal_norm', {}).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("## Each route's best score per target")
    rows.append('')
    for tgt in TARGETS:
        rows.append(f'### → {tgt}')
        rows.append('')
        rows.append("| Route | Best score | AUROC | Δ (vs same-route's terminal_norm) |")
        rows.append('|---|---|---|---|')
        for (label, prefix) in ROUTES:
            seeds = data.get((prefix, tgt), {})
            if not seeds:
                rows.append(f'| {label} | — | — | — |')
                continue
            score_means: dict[str, float] = {}
            for s in seeds.values():
                for (k, v) in s.get('overall', {}).items():
                    score_means.setdefault(k, []).append(v.get('auroc', float('nan')))
            mean_per_score = {k: _mean_std(v)[0] for (k, v) in score_means.items()}
            mean_per_score = {k: v for (k, v) in mean_per_score.items() if v == v}
            if not mean_per_score:
                rows.append(f'| {label} | — | — | — |')
                continue
            best = max(mean_per_score, key=mean_per_score.get)
            best_v = mean_per_score[best]
            best_sd = _mean_std(score_means[best])[1]
            tn = mean_per_score.get('terminal_norm', float('nan'))
            delta = f'{best_v - tn:+.4f}' if tn == tn else '—'
            rows.append(f'| {label} | `{best}` | {best_v:.4f} ± {best_sd:.4f} | {delta} |')
        rows.append('')
    for tgt in TARGETS:
        rows.append(f'## All key scores → {tgt}')
        rows.append('')
        header = '| Score | ' + ' | '.join((label for (label, _) in ROUTES)) + ' |'
        rows.append(header)
        rows.append('|---' * (1 + len(ROUTES)) + '|')
        for sc in PRIMARY_SCORES:
            cells = [f'`{sc}`']
            for (label, prefix) in ROUTES:
                seeds = data.get((prefix, tgt), {})
                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('')
    for tgt in TARGETS:
        rows.append(f'## Route C ensemble (terminal_norm + disc_nll) → {tgt}')
        rows.append('')
        c_seeds = data.get(('route_c_mixed', tgt), {})
        if c_seeds:
            alphas = (0.0, 0.5, 0.7, 0.8, 0.9, 1.0)
            rows.append('| α | ' + ' | '.join((f'seed{s}' for s in sorted(c_seeds.keys()))) + ' | mean ± std |')
            rows.append('|---' * (2 + len(c_seeds)) + '|')
            seed_sweeps = {}
            for s in c_seeds:
                npz = CROSS_DIR / f'route_c_mixed_seed{s}_to_{tgt}.npz'
                seed_sweeps[s] = _ensemble_sweep(npz) or {}
            for a in alphas:
                cells = [f'{a:.2f}']
                vals = []
                for s in sorted(c_seeds.keys()):
                    v = 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('## Run inventory')
    rows.append('')
    for (label, prefix) in ROUTES:
        for tgt in TARGETS:
            seeds = sorted(data.get((prefix, tgt), {}).keys())
            rows.append(f"- {label} → {tgt}: seeds = {(seeds if seeds else '(none)')}")
    out = ROOT / 'CROSS_RESULTS.md'
    out.write_text('\n'.join(rows) + '\n')
    print(f'[wrote] {out}')
if __name__ == '__main__':
    main()
