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
os.environ.setdefault('JAX_PLATFORMS', 'cpu')
import optax
from pzflow import Flow
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'))
from data import _apply_mixed_dequant, _zscore, load_unified_data
from model import UnifiedCFMConfig, UnifiedTokenCFM
from packet_store import PacketShardStore
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}), 'ciciot_within': ('runs/unified_cfm_ciciot2023_shafir5_2026_04_29', {'n_val': 10000, 'n_atk': 30000})}
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 _load_within(model_dir: Path, n_val: int | None, n_atk: int | None, n_train_cap: int, seed: int) -> dict[str, Any]:
    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_flow = data.train_flow
    if len(train_flow) > n_train_cap:
        idx = np.sort(rng.choice(len(train_flow), size=n_train_cap, replace=False))
        train_flow = train_flow[idx]
    val_flow = data.val_flow
    (atk_flow, atk_labels) = (data.attack_flow, data.attack_labels)
    if n_val is not None and len(val_flow) > n_val:
        idx = np.sort(rng.choice(len(val_flow), size=n_val, replace=False))
        val_flow = val_flow[idx]
    if n_atk is not None and len(atk_flow) > n_atk:
        idx = np.sort(rng.choice(len(atk_flow), size=n_atk, replace=False))
        atk_flow = atk_flow[idx]
        atk_labels = atk_labels[idx]
    return {'train_flow': train_flow, 'val_flow': val_flow, 'atk_flow': atk_flow, 'atk_labels': atk_labels}

def _load_cross(spec: dict[str, Any], ckpt_dict: dict[str, Any], seed: int, T: int, n_train_cap: int) -> dict[str, Any]:
    flow_mean = np.asarray(ckpt_dict['flow_mean'], dtype=np.float32)
    flow_std = np.asarray(ckpt_dict['flow_std'], dtype=np.float32)
    flow_names = [str(n) for n in ckpt_dict['flow_feature_names']]
    target_store = REPO / spec['target_store']
    target_flows = REPO / spec['target_flows']
    target_flow_features = REPO / spec['target_flow_features']
    (n_benign, n_attack) = (int(spec['n_benign']), int(spec['n_attack']))
    flows = pd.read_parquet(target_flows, columns=['flow_id', 'label'])
    ff = pd.read_parquet(target_flow_features)
    if not np.array_equal(flows['flow_id'].to_numpy(dtype=np.uint64), ff['flow_id'].to_numpy(dtype=np.uint64)):
        raise ValueError('target flows and flow_features not row-aligned')
    labels = flows['label'].astype(str).to_numpy()
    rng = np.random.default_rng(seed)
    benign_idx = np.flatnonzero(labels == 'normal')
    attack_idx = np.flatnonzero(labels != 'normal')
    b_sel = np.sort(rng.choice(benign_idx, size=n_benign, replace=False))
    atk_classes = sorted(set(labels[attack_idx]))
    per_class = max(1, n_attack // len(atk_classes))
    a_sel_chunks = []
    for cls in atk_classes:
        pool = attack_idx[labels[attack_idx] == cls]
        k = min(per_class, len(pool))
        if k:
            a_sel_chunks.append(rng.choice(pool, size=k, replace=False))
    a_sel = np.sort(np.concatenate(a_sel_chunks))
    if len(a_sel) > n_attack:
        a_sel = np.sort(rng.choice(a_sel, size=n_attack, replace=False))

    def _flow_only(idx):
        f = ff.iloc[idx][flow_names].to_numpy(dtype=np.float64)
        f = np.nan_to_num(f, nan=0.0, posinf=0.0, neginf=0.0).astype(np.float32)
        return ((f - flow_mean) / np.maximum(flow_std, 1e-06)).astype(np.float32)
    val_flow = _flow_only(b_sel)
    atk_flow = _flow_only(a_sel)
    atk_labels = labels[a_sel]
    src_flows = pd.read_parquet(REPO / ckpt_dict_paths(ckpt_dict)['flows'], columns=['flow_id', 'label'])
    src_ff = pd.read_parquet(REPO / ckpt_dict_paths(ckpt_dict)['flow_features'])
    if not np.array_equal(src_flows['flow_id'].to_numpy(dtype=np.uint64), src_ff['flow_id'].to_numpy(dtype=np.uint64)):
        raise ValueError('source flows and flow_features not row-aligned')
    src_labels = src_flows['label'].astype(str).to_numpy()
    src_benign_idx = np.flatnonzero(src_labels == 'normal')
    rng2 = np.random.default_rng(seed + 1000)
    if len(src_benign_idx) > n_train_cap:
        src_benign_idx = np.sort(rng2.choice(src_benign_idx, size=n_train_cap, replace=False))
    src_train = src_ff.iloc[src_benign_idx][flow_names].to_numpy(dtype=np.float64)
    src_train = np.nan_to_num(src_train, nan=0.0, posinf=0.0, neginf=0.0).astype(np.float32)
    train_flow = ((src_train - flow_mean) / np.maximum(flow_std, 1e-06)).astype(np.float32)
    return {'train_flow': train_flow, 'val_flow': val_flow, 'atk_flow': atk_flow, 'atk_labels': atk_labels, 'flow_names': flow_names}

def ckpt_dict_paths(ckpt: dict[str, Any]) -> dict[str, str]:
    raise NotImplementedError('paths must be passed via main()')

def _train_and_score(train_flow: np.ndarray, val_flow: np.ndarray, atk_flow: np.ndarray, *, epochs: int, lr: float, optimizer: str, verbose: bool):
    cols = [f'x{i}' for i in range(train_flow.shape[1])]
    df_train = pd.DataFrame(train_flow.astype(np.float32), columns=cols)
    df_val = pd.DataFrame(val_flow.astype(np.float32), columns=cols)
    df_atk = pd.DataFrame(atk_flow.astype(np.float32), columns=cols)
    if optimizer == 'sgd':
        opt = optax.sgd(learning_rate=lr)
    elif optimizer == 'adam':
        opt = optax.adam(learning_rate=lr)
    else:
        raise ValueError(f'unknown optimizer {optimizer!r}')
    flow = Flow(df_train.columns.tolist())
    t0 = time.time()
    losses = flow.train(df_train, optimizer=opt, epochs=epochs, verbose=verbose)
    t_train = time.time() - t0
    t0 = time.time()
    lp_val = np.asarray(flow.log_prob(df_val))
    lp_atk = np.asarray(flow.log_prob(df_atk))
    t_score = time.time() - t0
    return {'score_val': (-lp_val).astype(np.float32), 'score_atk': (-lp_atk).astype(np.float32), 'losses': np.asarray(losses, dtype=np.float64), 't_train': t_train, 't_score': t_score}

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 _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() -> None:
    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=10000, help='Cap benign train (default 10k mirrors Shafir).')
    p.add_argument('--epochs', type=int, default=100)
    p.add_argument('--lr', type=float, default=0.001)
    p.add_argument('--optimizer', choices=['sgd', 'adam'], default='sgd')
    p.add_argument('--T', type=int, default=64)
    p.add_argument('--verbose', action='store_true')
    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] shafir_nf protocol={args.protocol} seed={args.seed}')
    print(f'[run] using normalization stats from {model_dir}/model.pt (source ckpt)')
    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:
        cfg = yaml.safe_load((model_dir / 'config.yaml').read_text())
        flows_parquet = Path(cfg['flows_parquet'])
        flow_features_path = Path(cfg['flow_features_path'])
        flow_mean = np.asarray(ckpt['flow_mean'], dtype=np.float32)
        flow_std = np.asarray(ckpt['flow_std'], dtype=np.float32)
        flow_names = [str(n) for n in ckpt['flow_feature_names']]
        src_flows = pd.read_parquet(flows_parquet, columns=['flow_id', 'label'])
        src_ff = pd.read_parquet(flow_features_path)
        if not np.array_equal(src_flows['flow_id'].to_numpy(dtype=np.uint64), src_ff['flow_id'].to_numpy(dtype=np.uint64)):
            raise ValueError('source flows and flow_features not row-aligned')
        src_labels = src_flows['label'].astype(str).to_numpy()
        src_benign_idx = np.flatnonzero(src_labels == 'normal')
        rng2 = np.random.default_rng(args.seed + 1000)
        if len(src_benign_idx) > args.n_train_cap:
            src_benign_idx = np.sort(rng2.choice(src_benign_idx, size=args.n_train_cap, replace=False))
        src_train = src_ff.iloc[src_benign_idx][flow_names].to_numpy(dtype=np.float64)
        src_train = np.nan_to_num(src_train, nan=0.0, posinf=0.0, neginf=0.0).astype(np.float32)
        train_flow = ((src_train - flow_mean) / np.maximum(flow_std, 1e-06)).astype(np.float32)
        target_store = REPO / spec['target_store']
        target_flows = REPO / spec['target_flows']
        target_flow_features = REPO / spec['target_flow_features']
        (n_benign, n_attack) = (int(spec['n_benign']), int(spec['n_attack']))
        flows = pd.read_parquet(target_flows, columns=['flow_id', 'label'])
        ff = pd.read_parquet(target_flow_features)
        labels = flows['label'].astype(str).to_numpy()
        rng = np.random.default_rng(args.seed)
        b_sel = np.sort(rng.choice(np.flatnonzero(labels == 'normal'), size=n_benign, replace=False))
        atk_idx = np.flatnonzero(labels != 'normal')
        atk_classes = sorted(set(labels[atk_idx]))
        per_class_n = max(1, n_attack // len(atk_classes))
        chunks = []
        for cls in atk_classes:
            pool = atk_idx[labels[atk_idx] == cls]
            k = min(per_class_n, len(pool))
            if k:
                chunks.append(rng.choice(pool, size=k, replace=False))
        a_sel = np.sort(np.concatenate(chunks))
        if len(a_sel) > n_attack:
            a_sel = np.sort(rng.choice(a_sel, size=n_attack, replace=False))

        def _flow_only(idx):
            f = ff.iloc[idx][flow_names].to_numpy(dtype=np.float64)
            f = np.nan_to_num(f, nan=0.0, posinf=0.0, neginf=0.0).astype(np.float32)
            return ((f - flow_mean) / np.maximum(flow_std, 1e-06)).astype(np.float32)
        val_flow = _flow_only(b_sel)
        atk_flow = _flow_only(a_sel)
        atk_labels = labels[a_sel]
        arrays = {'train_flow': train_flow, 'val_flow': val_flow, 'atk_flow': atk_flow, 'atk_labels': atk_labels}
    print(f"[data] train={len(arrays['train_flow']):,}  val={len(arrays['val_flow']):,}  attack={len(arrays['atk_flow']):,}  D={arrays['train_flow'].shape[1]}")
    res = _train_and_score(arrays['train_flow'], arrays['val_flow'], arrays['atk_flow'], epochs=args.epochs, lr=args.lr, optimizer=args.optimizer, verbose=args.verbose)
    (val_score, atk_score) = (res['score_val'], res['score_atk'])
    y = np.r_[np.zeros(len(val_score)), np.ones(len(atk_score))]
    s = np.r_[val_score, atk_score]
    overall = {'neg_log_prob': {'auroc': _safe_metric(roc_auc_score, y, s), 'auprc': _safe_metric(average_precision_score, y, s)}}
    per_cls = _per_class(val_score, atk_score, np.asarray(arrays['atk_labels']).astype(str))
    out = {'method': 'shafir_nf', 'protocol': args.protocol, 'seed': args.seed, 'model_dir': str(model_dir), 'n_train': int(len(arrays['train_flow'])), 'n_val': int(len(arrays['val_flow'])), 'n_atk': int(len(arrays['atk_flow'])), 'epochs': args.epochs, 'lr': args.lr, 'optimizer': args.optimizer, 't_train_sec': round(res['t_train'], 2), 't_score_sec': round(res['t_score'], 2), 'loss_first_last': [float(res['losses'][0]), float(res['losses'][-1])], 'overall': overall, 'per_class': per_cls}
    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')
    np.savez_compressed(npz_path, b_neg_log_prob=val_score, a_neg_log_prob=atk_score, a_labels=np.asarray(arrays['atk_labels']).astype(str), losses=res['losses'])
    print(f'[saved] {out_json}')
    print(f'[saved] {npz_path}')
    print(f"[result] AUROC={overall['neg_log_prob']['auroc']:.4f}  AUPRC={overall['neg_log_prob']['auprc']:.4f}  train={res['t_train']:.1f}s  score={res['t_score']:.1f}s")
if __name__ == '__main__':
    main()
