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

def _mahal_eval(npz_path: Path, val_prefix: str, atk_prefix: str) -> float:
    if not npz_path.exists():
        return float('nan')
    z = np.load(npz_path, allow_pickle=True)
    keys = sorted([k.replace(val_prefix, '') for k in z.files if k.startswith(val_prefix) and (not k.endswith('labels'))])
    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)
    val_S = np.nan_to_num(val_S, nan=0.0, posinf=1000000.0, neginf=-1000000.0)
    atk_S = np.nan_to_num(atk_S, nan=0.0, posinf=1000000.0, neginf=-1000000.0)
    if len(val_S) < 50 or len(atk_S) < 50:
        return float('nan')
    y = np.r_[np.zeros(len(val_S)), np.ones(len(atk_S))]
    K = val_S.shape[1]
    try:
        oas = OAS().fit(val_S)
        inv_cov = np.linalg.inv(oas.covariance_ + 1e-09 * np.eye(K))
    except Exception:
        return float('nan')
    mu = val_S.mean(0)

    def m(S):
        d = S - mu
        return np.einsum('ni,ij,nj->n', d, inv_cov, d)
    s = np.r_[m(val_S), m(atk_S)]
    s = np.nan_to_num(s, nan=0.0, posinf=1000000000000.0, neginf=-1000000000000.0)
    try:
        return float(roc_auc_score(y, s))
    except ValueError:
        return float('nan')

def _within_mahal(ds: str, seed: int) -> float:
    md = ROOT / f'route_ac_combo_{ds}_seed{seed}'
    return _mahal_eval(md / 'phase1_scores.npz', 'val_', 'atk_')

def _within_terminal_norm(ds: str, seed: int) -> float:
    f = ROOT / f'route_ac_combo_{ds}_seed{seed}' / 'phase1_summary.json'
    if not f.exists():
        return float('nan')
    return json.loads(f.read_text())['overall'].get('terminal_norm', {}).get('auroc', float('nan'))

def _src_aliases(src: str) -> list[str]:
    aliases = [src]
    if src == 'cicddos2019':
        aliases.append('ddos2019')
    return aliases

def _cross_mahal(src: str, tgt: str, seed: int) -> float:
    candidates = []
    for alias in _src_aliases(src):
        candidates.append(CROSS_DIR / f'route_ac_combo_seed{seed}_{alias}_to_{tgt}.npz')
    if src == 'ciciot2023':
        candidates.append(CROSS_DIR / f'route_ac_combo_seed{seed}_to_{tgt}.npz')
    for c in candidates:
        if c.exists():
            return _mahal_eval(c, 'b_', 'a_')
    return float('nan')

def _cross_terminal_norm(src: str, tgt: str, seed: int) -> float:
    candidates = []
    for alias in _src_aliases(src):
        candidates.append(CROSS_DIR / f'route_ac_combo_seed{seed}_{alias}_to_{tgt}.json')
    if src == 'ciciot2023':
        candidates.append(CROSS_DIR / f'route_ac_combo_seed{seed}_to_{tgt}.json')
    for c in candidates:
        if c.exists():
            d = json.loads(c.read_text())
            return d['overall'].get('terminal_norm', {}).get('auroc', float('nan'))
    return float('nan')

def _ms(vals: list[float]) -> str:
    arr = np.asarray([v for v in vals if not np.isnan(v)], dtype=np.float64)
    if arr.size == 0:
        return '—'
    if arr.size == 1:
        return f'{arr[0]:.4f}'
    return f'{arr.mean():.4f}±{arr.std():.4f}'

def main() -> None:
    rows: list[str] = []
    rows.append('# Full 4×4 Cross Matrix — A+C combo + Mahalanobis-OAS')
    rows.append('')
    rows.append('3-seed mean ± std. Diagonal = within-dataset; off-diagonal = cross.')
    rows.append('Aggregator: Mahalanobis-OAS over 10-d A+C combo score vector,')
    rows.append('fit on **target-dataset benign val only** (no attack labels).')
    rows.append('')
    rows.append('## Mahalanobis-OAS AUROC (4×4)')
    rows.append('')
    rows.append('| Source ↓ \\ Target → | ' + ' | '.join(DATASETS) + ' |')
    rows.append('|---' * (1 + len(DATASETS)) + '|')
    for src in DATASETS:
        cells = [src]
        for tgt in DATASETS:
            if src == tgt:
                vals = [_within_mahal(src, s) for s in SEEDS]
                cells.append(f'_{_ms(vals)}_')
            else:
                vals = [_cross_mahal(src, tgt, s) for s in SEEDS]
                cells.append(_ms(vals))
        rows.append('| ' + ' | '.join(cells) + ' |')
    rows.append('')
    rows.append('(Italic diagonal = within-dataset reference)')
    rows.append('')
    rows.append('## `terminal_norm` AUROC (4×4) — for comparison (selection-bias-free single fixed score)')
    rows.append('')
    rows.append('| Source ↓ \\ Target → | ' + ' | '.join(DATASETS) + ' |')
    rows.append('|---' * (1 + len(DATASETS)) + '|')
    for src in DATASETS:
        cells = [src]
        for tgt in DATASETS:
            if src == tgt:
                vals = [_within_terminal_norm(src, s) for s in SEEDS]
                cells.append(f'_{_ms(vals)}_')
            else:
                vals = [_cross_terminal_norm(src, tgt, s) for s in SEEDS]
                cells.append(_ms(vals))
        rows.append('| ' + ' | '.join(cells) + ' |')
    rows.append('')
    rows.append('## Δ Mahalanobis − terminal_norm (where positive, Mahalanobis is better)')
    rows.append('')
    rows.append('| Source ↓ \\ Target → | ' + ' | '.join(DATASETS) + ' |')
    rows.append('|---' * (1 + len(DATASETS)) + '|')
    for src in DATASETS:
        cells = [src]
        for tgt in DATASETS:
            if src == tgt:
                m = np.mean([v for v in [_within_mahal(src, s) for s in SEEDS] if not np.isnan(v)])
                t = np.mean([v for v in [_within_terminal_norm(src, s) for s in SEEDS] if not np.isnan(v)])
            else:
                m = np.mean([v for v in [_cross_mahal(src, tgt, s) for s in SEEDS] if not np.isnan(v)])
                t = np.mean([v for v in [_cross_terminal_norm(src, tgt, s) for s in SEEDS] if not np.isnan(v)])
            if np.isnan(m) or np.isnan(t):
                cells.append('—')
            else:
                d = m - t
                if abs(d) < 0.005:
                    cells.append(f'{d:+.4f}')
                elif d > 0:
                    cells.append(f'**{d:+.4f}**')
                else:
                    cells.append(f'_{d:+.4f}_')
        rows.append('| ' + ' | '.join(cells) + ' |')
    rows.append('')
    rows.append('## Per-source averaged cross-AUROC (Mahalanobis, off-diagonal mean)')
    rows.append('')
    rows.append('| Source | mean off-diag Mahalanobis | mean off-diag terminal_norm |')
    rows.append('|---|---|---|')
    for src in DATASETS:
        m_offs = []
        t_offs = []
        for tgt in DATASETS:
            if src == tgt:
                continue
            m_vals = [_cross_mahal(src, tgt, s) for s in SEEDS]
            t_vals = [_cross_terminal_norm(src, tgt, s) for s in SEEDS]
            m_offs.extend([v for v in m_vals if not np.isnan(v)])
            t_offs.extend([v for v in t_vals if not np.isnan(v)])
        m_mean = np.mean(m_offs) if m_offs else float('nan')
        t_mean = np.mean(t_offs) if t_offs else float('nan')
        rows.append(f'| {src} | {m_mean:.4f} | {t_mean:.4f} |')
    out = ROOT / 'CROSS_MATRIX.md'
    out.write_text('\n'.join(rows) + '\n')
    print(f'[wrote] {out}')
if __name__ == '__main__':
    main()
