from __future__ import annotations
import argparse
import json
import os
import sys
import time
from pathlib import Path
from typing import Any
import numpy as np
import pandas as pd
import torch
import yaml
from sklearn.metrics import average_precision_score, roc_auc_score
REPO = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(REPO / 'Packet_CFM'))
sys.path.insert(0, str(REPO / 'Unified_CFM'))
sys.path.insert(0, str(REPO / 'baselines/Kitsune-py'))
from data import _apply_mixed_dequant, _zscore, load_unified_data
from packet_store import PacketShardStore
from KitNET.KitNET import KitNET
WITHIN_DIRS = {'iscxtor_within': ('phase25_multiseed_2026_04_25/iscxtor2016_lambda0p3_seed{seed}', {'n_val': 10000, 'n_atk': None}), 'cicids_within': ('phase25_sigma06_multiseed_2026_04_25/cicids2017_lambda0p3_sigma0p6_seed{seed}', {'n_val': 10000, 'n_atk': 30000}), 'cicddos_within': ('phase25_multiseed_2026_04_25/cicddos2019_lambda0p3_seed{seed}', {'n_val': 10000, 'n_atk': 20000})}
CROSS_DIRS = {'forward_cross': {'model_template': 'phase25_sigma06_multiseed_2026_04_25/cicids2017_lambda0p3_sigma0p6_seed{seed}', 'target_store': 'datasets/cicddos2019/processed/full_store', 'target_flows': 'datasets/cicddos2019/processed/flows.parquet', 'target_flow_features': 'datasets/cicddos2019/processed/flow_features.parquet', 'n_benign': 10000, 'n_attack': 10000}, 'reverse_cross': {'model_template': 'phase25_multiseed_2026_04_25/cicddos2019_lambda0p3_seed{seed}', 'target_store': 'datasets/cicids2017/processed/full_store', 'target_flows': 'datasets/cicids2017/processed/flows.parquet', 'target_flow_features': 'datasets/cicids2017/processed/flow_features.parquet', 'n_benign': 10000, 'n_attack': 10000}}

def _safe_metric(fn, y, s) -> float:
    s = np.nan_to_num(s, nan=0.0, posinf=1000000000000.0, neginf=-1000000000000.0)
    try:
        return float(fn(y, s))
    except ValueError:
        return float('nan')

def _load_within(model_dir: Path, n_val, n_atk, n_train_cap, seed):
    cfg = yaml.safe_load((model_dir / 'config.yaml').read_text())
    data = load_unified_data(packets_npz=Path(cfg['packets_npz']) if cfg.get('packets_npz') else None, source_store=Path(cfg['source_store']) if cfg.get('source_store') else None, flows_parquet=Path(cfg['flows_parquet']), flow_features_path=Path(cfg['flow_features_path']) if cfg.get('flow_features_path') else None, flow_feature_columns=cfg.get('flow_feature_columns'), flow_features_align=str(cfg.get('flow_features_align', 'auto')), T=int(cfg['T']), split_seed=int(cfg.get('data_seed', cfg.get('seed', 42))), train_ratio=float(cfg.get('train_ratio', 0.8)), benign_label=str(cfg.get('benign_label', 'normal')), min_len=int(cfg.get('min_len', 2)), packet_preprocess=str(cfg.get('packet_preprocess', 'mixed_dequant')), attack_cap=int(cfg['attack_cap']) if cfg.get('attack_cap') else n_atk, val_cap=int(cfg['val_cap']) if cfg.get('val_cap') else n_val)
    rng = np.random.default_rng(seed)
    (train_packets, train_len) = (data.train_packets, data.train_len)
    if len(train_packets) > n_train_cap:
        idx = np.sort(rng.choice(len(train_packets), size=n_train_cap, replace=False))
        (train_packets, train_len) = (train_packets[idx], train_len[idx])
    (val_packets, val_len) = (data.val_packets, data.val_len)
    (atk_packets, atk_len, atk_labels) = (data.attack_packets, data.attack_len, data.attack_labels)
    if n_val is not None and len(val_packets) > n_val:
        idx = np.sort(rng.choice(len(val_packets), size=n_val, replace=False))
        (val_packets, val_len) = (val_packets[idx], val_len[idx])
    if n_atk is not None and len(atk_packets) > n_atk:
        idx = np.sort(rng.choice(len(atk_packets), size=n_atk, replace=False))
        (atk_packets, atk_len, atk_labels) = (atk_packets[idx], atk_len[idx], atk_labels[idx])
    return {'train_packets': train_packets, 'train_len': train_len, 'val_packets': val_packets, 'val_len': val_len, 'atk_packets': atk_packets, 'atk_len': atk_len, 'atk_labels': atk_labels}

def _load_cross(spec, ckpt, seed, n_train_cap, T):
    packet_mean = np.asarray(ckpt['packet_mean'], dtype=np.float32)
    packet_std = np.asarray(ckpt['packet_std'], dtype=np.float32)
    packet_preprocess = str(ckpt.get('packet_preprocess', 'mixed_dequant'))
    src_cfg_path = REPO / 'artifacts' / spec['model_template'].format(seed=seed) / 'config.yaml'
    src_cfg = yaml.safe_load(src_cfg_path.read_text())
    src_data = load_unified_data(packets_npz=Path(src_cfg['packets_npz']) if src_cfg.get('packets_npz') else None, source_store=Path(src_cfg['source_store']) if src_cfg.get('source_store') else None, flows_parquet=Path(src_cfg['flows_parquet']), flow_features_path=Path(src_cfg['flow_features_path']) if src_cfg.get('flow_features_path') else None, flow_feature_columns=src_cfg.get('flow_feature_columns'), flow_features_align=str(src_cfg.get('flow_features_align', 'auto')), T=int(src_cfg['T']), split_seed=int(src_cfg.get('data_seed', src_cfg.get('seed', 42))), train_ratio=float(src_cfg.get('train_ratio', 0.8)), benign_label=str(src_cfg.get('benign_label', 'normal')), min_len=int(src_cfg.get('min_len', 2)), packet_preprocess=packet_preprocess, attack_cap=None, val_cap=None)
    rng = np.random.default_rng(seed + 1000)
    (train_packets, train_len) = (src_data.train_packets, src_data.train_len)
    if len(train_packets) > n_train_cap:
        idx = np.sort(rng.choice(len(train_packets), size=n_train_cap, replace=False))
        (train_packets, train_len) = (train_packets[idx], train_len[idx])
    target_store = REPO / spec['target_store']
    target_flows = REPO / spec['target_flows']
    (n_benign, n_attack) = (int(spec['n_benign']), int(spec['n_attack']))
    flows = pd.read_parquet(target_flows, columns=['flow_id', 'label'])
    labels = flows['label'].astype(str).to_numpy()
    rng2 = np.random.default_rng(seed)
    benign_idx = np.flatnonzero(labels == 'normal')
    attack_idx = np.flatnonzero(labels != 'normal')
    b_sel = np.sort(rng2.choice(benign_idx, size=n_benign, replace=False))
    atk_classes = sorted(set(labels[attack_idx]))
    per_class = max(1, n_attack // len(atk_classes))
    chunks = []
    for cls in atk_classes:
        pool = attack_idx[labels[attack_idx] == cls]
        k = min(per_class, len(pool))
        if k:
            chunks.append(rng2.choice(pool, size=k, replace=False))
    a_sel = np.sort(np.concatenate(chunks))
    if len(a_sel) > n_attack:
        a_sel = np.sort(rng2.choice(a_sel, size=n_attack, replace=False))
    store = PacketShardStore.open(target_store)

    def _materialize(idx):
        (tok, ll) = store.read_packets(idx, T=T)
        ll = np.minimum(ll, T).astype(np.int32)
        return (tok.astype(np.float32), ll)
    (b_tok, b_len) = _materialize(b_sel)
    (a_tok, a_len) = _materialize(a_sel)
    if packet_preprocess == 'mixed_dequant':
        val_packets = _apply_mixed_dequant(b_tok, b_len, packet_mean, packet_std, split_tag='val', seed=seed)
        atk_packets = _apply_mixed_dequant(a_tok, a_len, packet_mean, packet_std, split_tag='attack', seed=seed)
    else:
        val_packets = _zscore(b_tok, packet_mean, packet_std)
        atk_packets = _zscore(a_tok, packet_mean, packet_std)
        msk_b = np.arange(T)[None, :] < b_len[:, None]
        msk_a = np.arange(T)[None, :] < a_len[:, None]
        val_packets = (val_packets * msk_b[:, :, None]).astype(np.float32)
        atk_packets = (atk_packets * msk_a[:, :, None]).astype(np.float32)
    return {'train_packets': train_packets, 'train_len': train_len, 'val_packets': val_packets, 'val_len': b_len, 'atk_packets': atk_packets, 'atk_len': a_len, 'atk_labels': labels[a_sel]}

def _flatten_packets(packets: np.ndarray, lens: np.ndarray) -> np.ndarray:
    out_chunks = []
    for i in range(len(packets)):
        L = int(lens[i])
        if L > 0:
            out_chunks.append(packets[i, :L])
    if not out_chunks:
        return np.empty((0, packets.shape[-1]), dtype=np.float32)
    return np.concatenate(out_chunks, axis=0).astype(np.float32)

def _train_kitnet(kit: KitNET, train_flat: np.ndarray) -> dict[str, float]:
    t0 = time.time()
    last_rmse = 0.0
    for i in range(len(train_flat)):
        last_rmse = kit.process(train_flat[i])
        if (i + 1) % 50000 == 0:
            print(f'  [train] processed {i + 1:,}/{len(train_flat):,}  last_rmse={last_rmse:.4f}', flush=True)
    return {'t_train_sec': round(time.time() - t0, 2), 'n_trained_packets': len(train_flat)}

def _score_flows(kit: KitNET, packets: np.ndarray, lens: np.ndarray) -> dict[str, np.ndarray]:
    N = len(packets)
    means = np.zeros(N, dtype=np.float32)
    maxes = np.zeros(N, dtype=np.float32)
    medians = np.zeros(N, dtype=np.float32)
    p90s = np.zeros(N, dtype=np.float32)
    for i in range(N):
        L = int(lens[i])
        if L == 0:
            continue
        rmses = np.zeros(L, dtype=np.float32)
        for t in range(L):
            rmses[t] = kit.execute(packets[i, t])
        means[i] = rmses.mean()
        maxes[i] = rmses.max()
        medians[i] = np.median(rmses)
        p90s[i] = np.percentile(rmses, 90)
    return {'mean': means, 'max': maxes, 'median': medians, 'p90': p90s}

def _per_class(val_score: np.ndarray, atk_score: np.ndarray, atk_labels: np.ndarray):
    out = {}
    for cls in sorted(set(atk_labels)):
        m = atk_labels == cls
        n_c = int(m.sum())
        v_c = atk_score[m]
        y = np.r_[np.zeros(len(val_score)), np.ones(len(v_c))]
        s = np.r_[val_score, v_c]
        out[cls] = {'_n': float(n_c), 'auroc': _safe_metric(roc_auc_score, y, s)}
    return out

def main():
    p = argparse.ArgumentParser()
    p.add_argument('--protocol', required=True, choices=list(WITHIN_DIRS) + list(CROSS_DIRS))
    p.add_argument('--seed', type=int, required=True, choices=[42, 43, 44])
    p.add_argument('--out-dir', type=Path, required=True)
    p.add_argument('--n-train-cap', type=int, default=2000, help='Cap source-benign train flows (each contributes ~T packets).')
    p.add_argument('--fm-grace', type=int, default=2000, help='Kitsune feature-mapper grace period (packets).')
    p.add_argument('--ad-grace', type=int, default=20000, help='Kitsune anomaly-detector grace period (packets).')
    p.add_argument('--max-ae-size', type=int, default=10)
    p.add_argument('--lr', type=float, default=0.1)
    p.add_argument('--hidden-ratio', type=float, default=0.75)
    p.add_argument('--T', type=int, default=64)
    args = p.parse_args()
    args.out_dir.mkdir(parents=True, exist_ok=True)
    is_within = args.protocol in WITHIN_DIRS
    if is_within:
        (template, caps) = WITHIN_DIRS[args.protocol]
        model_dir = REPO / 'artifacts' / template.format(seed=args.seed)
    else:
        spec = CROSS_DIRS[args.protocol]
        model_dir = REPO / 'artifacts' / spec['model_template'].format(seed=args.seed)
    print(f'[run] kitsune protocol={args.protocol} seed={args.seed}')
    print(f'[run] using packet stats from {model_dir}/model.pt')
    ckpt = torch.load(model_dir / 'model.pt', map_location='cpu', weights_only=False)
    if is_within:
        arrays = _load_within(model_dir, n_val=caps['n_val'], n_atk=caps['n_atk'], n_train_cap=args.n_train_cap, seed=args.seed)
    else:
        arrays = _load_cross(spec, ckpt, args.seed, args.n_train_cap, args.T)
    n_train = len(arrays['train_packets'])
    n_val = len(arrays['val_packets'])
    n_atk = len(arrays['atk_packets'])
    D = arrays['train_packets'].shape[-1]
    print(f'[data] train_flows={n_train:,}  val={n_val:,}  attack={n_atk:,}  D={D}')
    train_flat = _flatten_packets(arrays['train_packets'], arrays['train_len'])
    print(f'[data] train_flat packets={len(train_flat):,}  FM_grace={args.fm_grace} AD_grace={args.ad_grace}')
    if len(train_flat) < args.fm_grace + args.ad_grace:
        raise ValueError(f'Need at least FM+AD={args.fm_grace + args.ad_grace} packets, have {len(train_flat)} (try increasing --n-train-cap).')
    kit = KitNET(n=D, max_autoencoder_size=args.max_ae_size, FM_grace_period=args.fm_grace, AD_grace_period=args.ad_grace, learning_rate=args.lr, hidden_ratio=args.hidden_ratio)
    train_meta = _train_kitnet(kit, train_flat)
    print(f'[train] {train_meta}')
    t0 = time.time()
    val_aggs = _score_flows(kit, arrays['val_packets'], arrays['val_len'])
    print(f'[score] benign in {time.time() - t0:.1f}s')
    t0 = time.time()
    atk_aggs = _score_flows(kit, arrays['atk_packets'], arrays['atk_len'])
    print(f'[score] attack in {time.time() - t0:.1f}s')
    overall = {}
    per_class_by_agg = {}
    for agg in ('mean', 'max', 'median', 'p90'):
        v = val_aggs[agg]
        a = atk_aggs[agg]
        y = np.r_[np.zeros(len(v)), np.ones(len(a))]
        s = np.r_[v, a]
        overall[agg] = {'auroc': _safe_metric(roc_auc_score, y, s), 'auprc': _safe_metric(average_precision_score, y, s)}
        per_class_by_agg[agg] = _per_class(v, a, np.asarray(arrays['atk_labels']).astype(str))
    out = {'method': 'kitsune_path_b', 'protocol': args.protocol, 'seed': args.seed, 'model_dir': str(model_dir), 'n_train_flows': n_train, 'n_train_packets': int(len(train_flat)), 'n_val': n_val, 'n_atk': n_atk, 'D': int(D), 'fm_grace': args.fm_grace, 'ad_grace': args.ad_grace, 'max_ae_size': args.max_ae_size, 't_train_sec': train_meta['t_train_sec'], 'overall_by_agg': overall, 'per_class_by_agg': per_class_by_agg}
    out_json = args.out_dir / f'{args.protocol}_seed{args.seed}.json'
    out_json.write_text(json.dumps(out, indent=2))
    npz_path = out_json.with_suffix('.npz')
    save = {'a_labels': np.asarray(arrays['atk_labels']).astype(str)}
    for agg in ('mean', 'max', 'median', 'p90'):
        save[f'b_{agg}'] = val_aggs[agg].astype(np.float32)
        save[f'a_{agg}'] = atk_aggs[agg].astype(np.float32)
    np.savez_compressed(npz_path, **save)
    print(f'[saved] {out_json}')
    print(f'[saved] {npz_path}')
    best = max(overall, key=lambda k: overall[k]['auroc'])
    print(f"[best agg={best}] AUROC={overall[best]['auroc']:.4f}  AUPRC={overall[best]['auprc']:.4f}")
    print()
    print('=== overall AUROC by aggregator ===')
    for k in sorted(overall, key=lambda kk: -overall[kk]['auroc']):
        print(f"  {k:<8s} AUROC={overall[k]['auroc']:.4f}  AUPRC={overall[k]['auprc']:.4f}")
if __name__ == '__main__':
    main()
