from __future__ import annotations
import json
from pathlib import Path
import numpy as np
from sklearn.covariance import LedoitWolf, OAS, GraphicalLassoCV
from sklearn.metrics import roc_auc_score
ROOT = Path(__file__).resolve().parent
CROSS_DIR = ROOT / 'cross'
WITHIN_DATASETS = ['iscxtor2016', 'cicids2017', 'cicddos2019', 'ciciot2023']
CROSS_TARGETS = ['cicids2017', 'cicddos2019']
SEEDS = [42, 43, 44]

def _aggregators(val_S: np.ndarray, test_S_list: list[np.ndarray]) -> dict[str, list[np.ndarray]]:
    val_S = np.nan_to_num(val_S, nan=0.0, posinf=1000000.0, neginf=-1000000.0)
    test_S_list = [np.nan_to_num(t, nan=0.0, posinf=1000000.0, neginf=-1000000.0) for t in test_S_list]
    mu = val_S.mean(axis=0)
    sigma = val_S.std(axis=0) + 1e-09
    K = val_S.shape[1]
    cov_emp = np.cov(val_S, rowvar=False)
    inv_cov_plain = np.linalg.inv(cov_emp + 0.001 * np.eye(K))
    lw = LedoitWolf().fit(val_S)
    inv_cov_lw = np.linalg.inv(lw.covariance_ + 1e-09 * np.eye(K))
    oas = OAS().fit(val_S)
    inv_cov_oas = np.linalg.inv(oas.covariance_ + 1e-09 * np.eye(K))

    def _max_abs_z(S):
        return np.abs((S - mu) / sigma).max(axis=1)

    def _max_pos_z(S):
        return ((S - mu) / sigma).max(axis=1)

    def _mahal_factory(inv_cov):

        def f(S):
            d = S - mu
            return np.einsum('ni,ij,nj->n', d, inv_cov, d)
        return f
    out: dict[str, list[np.ndarray]] = {}
    for (tag, fn) in [('max_abs_z', _max_abs_z), ('max_pos_z', _max_pos_z), ('mahal_plain', _mahal_factory(inv_cov_plain)), ('mahal_lw', _mahal_factory(inv_cov_lw)), ('mahal_oas', _mahal_factory(inv_cov_oas))]:
        out[tag] = [fn(t) for t in test_S_list]
    return out
SCORE_SUBSETS = {'all': None, 'terminal3': ['terminal_norm', 'terminal_flow', 'terminal_packet'], 'disc7': ['disc_nll_total', 'disc_nll_ch2', 'disc_nll_ch3', 'disc_nll_ch4', 'disc_nll_ch5', 'disc_nll_ch6', 'disc_nll_ch7']}

def _evaluate(npz: Path, val_prefix: str, atk_prefix: str) -> dict:
    z = np.load(npz, allow_pickle=True)
    all_keys = sorted([k.replace(val_prefix, '') for k in z.files if k.startswith(val_prefix) and (not k.endswith('labels'))])
    out: dict = {'n_val': None, 'n_atk': None}
    for (subset_name, subset_keys) in SCORE_SUBSETS.items():
        if subset_keys is None:
            keys = all_keys
        else:
            keys = [k for k in subset_keys if k in all_keys]
            if len(keys) < 2:
                continue
        val_S = np.stack([z[f'{val_prefix}{k}'] for k in keys], axis=1)
        atk_S = np.stack([z[f'{atk_prefix}{k}'] for k in keys], axis=1)
        (n_val, n_atk) = (val_S.shape[0], atk_S.shape[0])
        out['n_val'] = n_val
        out['n_atk'] = n_atk
        y = np.r_[np.zeros(n_val), np.ones(n_atk)]
        aggs = _aggregators(val_S, [val_S, atk_S])
        for (tag, (v_agg, a_agg)) in aggs.items():
            s = np.r_[v_agg, a_agg]
            s = np.nan_to_num(s, nan=0.0, posinf=1000000000000.0, neginf=-1000000000000.0)
            try:
                auc = float(roc_auc_score(y, s))
            except ValueError:
                auc = float('nan')
            out[f'auc_{tag}_{subset_name}'] = auc
    out['auc_max_abs_z'] = out.get('auc_max_abs_z_all')
    out['auc_max_pos_z'] = out.get('auc_max_pos_z_all')
    out['auc_mahal_plain'] = out.get('auc_mahal_plain_all')
    out['auc_mahal_lw'] = out.get('auc_mahal_lw_all')
    out['auc_mahal_oas'] = out.get('auc_mahal_oas_all')
    val_S = np.stack([z[f'{val_prefix}{k}'] for k in all_keys], axis=1)
    atk_S = np.stack([z[f'{atk_prefix}{k}'] for k in all_keys], axis=1)
    val_S = np.nan_to_num(val_S, nan=0.0, posinf=1000000000000.0, neginf=-1000000000000.0)
    atk_S = np.nan_to_num(atk_S, nan=0.0, posinf=1000000000000.0, neginf=-1000000000000.0)
    y = np.r_[np.zeros(val_S.shape[0]), np.ones(atk_S.shape[0])]
    per_score = {}
    for (i, k) in enumerate(all_keys):
        s = np.r_[val_S[:, i], atk_S[:, i]]
        s = np.nan_to_num(s, nan=0.0, posinf=1000000000000.0, neginf=-1000000000000.0)
        a1 = roc_auc_score(y, s)
        per_score[k] = max(a1, 1 - a1)
    best_score = max(per_score, key=per_score.get)
    out['auc_best_fixed'] = per_score[best_score]
    out['best_fixed_name'] = best_score
    out['auc_term_norm'] = per_score.get('terminal_norm', float('nan'))
    out['auc_term_pkt'] = per_score.get('terminal_packet', float('nan'))
    out['auc_disc_total'] = per_score.get('disc_nll_total', float('nan'))
    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:
    rows: list[str] = []
    rows.append('# Score-vector auto-selection: max-of-|z| / Mahalanobis vs fixed scores')
    rows.append('')
    rows.append('Aggregators are fit on **benign val only** (no attack labels). All numbers')
    rows.append('are 3-seed mean ± std on A+C combo (Mixed_CFM + causal-packet attention).')
    rows.append('')
    rows.append('Note on fairness: `auc_best_fixed` is selection-biased (picks per-dataset best')
    rows.append('score post-hoc on test set). `max_abs_z` and `mahalanobis` are NOT — they only')
    rows.append('use benign val to fit aggregator parameters.')
    rows.append('')
    rows.append("## Within-dataset(A+C combo on each dataset's own benign/attack)")
    rows.append('')
    rows.append('| Dataset | term_norm | best fixed | max-\\|z\\| (all) | mahal-OAS (all) | **mahal-OAS (term3)** | **mahal-OAS (disc7)** |')
    rows.append('|---|---|---|---|---|---|---|')
    for ds in WITHIN_DATASETS:
        rows_per_seed: list[dict] = []
        for s in SEEDS:
            md = ROOT / f'route_ac_combo_{ds}_seed{s}'
            npz = md / 'phase1_scores.npz'
            if not npz.exists():
                continue
            rows_per_seed.append(_evaluate(npz, 'val_', 'atk_'))
        if not rows_per_seed:
            rows.append(f'| {ds} | (no data) | | | | | |')
            continue

        def col(field):
            (m, sd) = _mean_std([r[field] for r in rows_per_seed])
            return f'{m:.4f} ± {sd:.4f}'
        rows.append(f"| {ds} | {col('auc_term_norm')} | {col('auc_best_fixed')} | {col('auc_max_abs_z_all')} | {col('auc_mahal_oas_all')} | **{col('auc_mahal_oas_terminal3')}** | **{col('auc_mahal_oas_disc7')}** |")
    rows.append('')
    rows.append('## Cross-dataset(A+C combo trained on CICIoT2023 → eval on target)')
    rows.append('')
    rows.append('| Target | term_norm | best fixed | max-\\|z\\| (all) | mahal-OAS (all) | **mahal-OAS (term3)** | **mahal-OAS (disc7)** |')
    rows.append('|---|---|---|---|---|---|---|')
    for tgt in CROSS_TARGETS:
        rows_per_seed: list[dict] = []
        for s in SEEDS:
            npz = CROSS_DIR / f'route_ac_combo_seed{s}_to_{tgt}.npz'
            if not npz.exists():
                continue
            rows_per_seed.append(_evaluate(npz, 'b_', 'a_'))
        if not rows_per_seed:
            rows.append(f'| {tgt} | (no data) | | | | | |')
            continue

        def col(field):
            (m, sd) = _mean_std([r[field] for r in rows_per_seed])
            return f'{m:.4f} ± {sd:.4f}'
        rows.append(f"| {tgt} | {col('auc_term_norm')} | {col('auc_best_fixed')} | {col('auc_max_abs_z_all')} | {col('auc_mahal_oas_all')} | **{col('auc_mahal_oas_terminal3')}** | **{col('auc_mahal_oas_disc7')}** |")
    rows.append('')
    rows.append('## Best-fixed-score winner per setup')
    rows.append('')
    rows.append('| Setup | seed42 | seed43 | seed44 |')
    rows.append('|---|---|---|---|')
    for ds in WITHIN_DATASETS:
        cells = [f'within {ds}']
        for s in SEEDS:
            npz = ROOT / f'route_ac_combo_{ds}_seed{s}/phase1_scores.npz'
            if not npz.exists():
                cells.append('—')
                continue
            r = _evaluate(npz, 'val_', 'atk_')
            cells.append(f"{r['best_fixed_name']} ({r['auc_best_fixed']:.4f})")
        rows.append('| ' + ' | '.join(cells) + ' |')
    for tgt in CROSS_TARGETS:
        cells = [f'cross→{tgt}']
        for s in SEEDS:
            npz = CROSS_DIR / f'route_ac_combo_seed{s}_to_{tgt}.npz'
            if not npz.exists():
                cells.append('—')
                continue
            r = _evaluate(npz, 'b_', 'a_')
            cells.append(f"{r['best_fixed_name']} ({r['auc_best_fixed']:.4f})")
        rows.append('| ' + ' | '.join(cells) + ' |')
    out = ROOT / 'SCORE_ROUTER.md'
    out.write_text('\n'.join(rows) + '\n')
    print(f'[wrote] {out}')
if __name__ == '__main__':
    main()
