from __future__ import annotations
import argparse
import json
import sys
from pathlib import Path
import numpy as np
import torch
from sklearn.metrics import average_precision_score, roc_auc_score
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from data import _preprocess_packet_batch
from detect import _load_model
from packet_store import PacketShardStore

@torch.no_grad()
def _score_indices(*, store: PacketShardStore, indices: np.ndarray, model, device: torch.device, preprocess: str, mean: np.ndarray, std: np.ndarray, clip_lo: np.ndarray | None, clip_hi: np.ndarray | None, split_tag: str, split_seed: int, batch: int, materialize_batch: int, n_steps: int) -> dict[str, np.ndarray]:
    out = {'terminal_norm': [], 'arc_length': [], 'kinetic_energy': [], 'velocity_score': []}
    total = len(indices)
    report_every = max(1, total // 4)
    next_report = 0
    for start in range(0, total, materialize_batch):
        idx = indices[start:start + materialize_batch]
        (x_np, lens_np) = store.read_packets(idx, T=model.cfg.T)
        x_np = _preprocess_packet_batch(x_np, lens_np, preprocess=preprocess, mean=mean, std=std, clip_lo=clip_lo, clip_hi=clip_hi, split_tag=split_tag, split_seed=split_seed, flow_ids=idx)
        for pos in range(0, len(idx), batch):
            bx = torch.from_numpy(x_np[pos:pos + batch]).float().to(device)
            bl = torch.from_numpy(lens_np[pos:pos + batch]).long().to(device)
            m = model.trajectory_metrics(bx, lens=bl, cond=None, n_steps=n_steps)
            for key in ('terminal_norm', 'arc_length', 'kinetic_energy'):
                out[key].append(m[key].cpu().numpy())
            vs = model.velocity_score(bx, lens=bl, cond=None, t_eval=(0.5, 0.75, 1.0))
            out['velocity_score'].append(vs.cpu().numpy())
        done = min(start + len(idx), total)
        if done >= next_report or done == total:
            print(f'[{split_tag}] {done:,}/{total:,}', flush=True)
            next_report = done + report_every
    return {key: np.concatenate(parts) for (key, parts) in out.items()}

def run(args: argparse.Namespace) -> None:
    device = torch.device('cuda' if args.device == 'auto' and torch.cuda.is_available() else 'cpu' if args.device == 'auto' else args.device)
    save_dir = Path(args.save_dir)
    ckpt = torch.load(save_dir / 'model.pt', map_location='cpu', weights_only=False)
    preprocess = str(ckpt.get('preprocess', 'zscore'))
    mean = np.asarray(ckpt['packet_mean'], dtype=np.float32)
    std = np.asarray(ckpt['packet_std'], dtype=np.float32)
    clip_lo = np.asarray(ckpt['clip_lo'], dtype=np.float32) if 'clip_lo' in ckpt else None
    clip_hi = np.asarray(ckpt['clip_hi'], dtype=np.float32) if 'clip_hi' in ckpt else None
    model = _load_model(save_dir, device)
    store = PacketShardStore.open(Path(args.target_store))
    flows = store.read_flows(columns=['flow_id', 'label'])
    labels = flows['label'].to_numpy().astype(str)
    lens = store.manifest['packet_length'].to_numpy(dtype=np.int32)
    keep = lens >= int(args.min_len)
    benign_idx = flows.loc[keep & (labels == args.benign_label), 'flow_id'].to_numpy(dtype=np.int64)
    attack_df = flows.loc[keep & (labels != args.benign_label), ['flow_id', 'label']]
    attack_idx_all = attack_df['flow_id'].to_numpy(dtype=np.int64)
    attack_labels_all = attack_df['label'].to_numpy().astype(str)
    if len(benign_idx) < args.n_benign:
        raise ValueError(f'target has only {len(benign_idx)} benign rows, need {args.n_benign}')
    if len(attack_idx_all) < args.n_attack:
        raise ValueError(f'target has only {len(attack_idx_all)} attack rows, need {args.n_attack}')
    print(f'[target] store={args.target_store} benign_pool={len(benign_idx):,} attack_pool={len(attack_idx_all):,} T={model.cfg.T} preprocess={preprocess}', flush=True)
    results: dict[str, object] = {'save_dir': str(save_dir), 'target_store': str(args.target_store), 'n_benign': int(args.n_benign), 'n_attack': int(args.n_attack), 'seeds': [], 'mean': {}, 'std': {}}
    metrics = ('terminal_norm', 'arc_length', 'kinetic_energy', 'velocity_score')
    per_metric_values = {f'{metric}_auroc': [] for metric in metrics}
    per_metric_values.update({f'{metric}_auprc': [] for metric in metrics})
    for seed in args.seeds:
        rng = np.random.default_rng(int(seed))
        b_idx = np.sort(rng.choice(benign_idx, args.n_benign, replace=False))
        a_pos = rng.choice(len(attack_idx_all), args.n_attack, replace=False)
        a_pos.sort()
        a_idx = attack_idx_all[a_pos]
        a_labels = attack_labels_all[a_pos]
        print(f'[seed={seed}] scoring benign={len(b_idx):,} attack={len(a_idx):,}', flush=True)
        b_scores = _score_indices(store=store, indices=b_idx, model=model, device=device, preprocess=preprocess, mean=mean, std=std, clip_lo=clip_lo, clip_hi=clip_hi, split_tag='val', split_seed=int(seed), batch=args.batch, materialize_batch=args.materialize_batch, n_steps=args.n_steps)
        a_scores = _score_indices(store=store, indices=a_idx, model=model, device=device, preprocess=preprocess, mean=mean, std=std, clip_lo=clip_lo, clip_hi=clip_hi, split_tag='attack', split_seed=int(seed), batch=args.batch, materialize_batch=args.materialize_batch, n_steps=args.n_steps)
        seed_result: dict[str, object] = {'seed': int(seed), 'attack_label_counts': {str(k): int(v) for (k, v) in zip(*np.unique(a_labels, return_counts=True))}, 'metrics': {}}
        for metric in metrics:
            y = np.r_[np.zeros(len(b_scores[metric])), np.ones(len(a_scores[metric]))]
            s = np.r_[b_scores[metric], a_scores[metric]]
            s = np.nan_to_num(s, nan=0.0, posinf=1000000.0, neginf=-1000000.0)
            auroc = float(roc_auc_score(y, s))
            auprc = float(average_precision_score(y, s))
            seed_result['metrics'][metric] = {'auroc': auroc, 'auprc': auprc}
            per_metric_values[f'{metric}_auroc'].append(auroc)
            per_metric_values[f'{metric}_auprc'].append(auprc)
            print(f'[seed={seed}] {metric:<16s} AUROC={auroc:.4f} AUPRC={auprc:.4f}', flush=True)
        results['seeds'].append(seed_result)
    for (key, values) in per_metric_values.items():
        arr = np.asarray(values, dtype=np.float64)
        results['mean'][key] = float(arr.mean())
        results['std'][key] = float(arr.std(ddof=0))
    Path(args.output).parent.mkdir(parents=True, exist_ok=True)
    Path(args.output).write_text(json.dumps(results, indent=2, sort_keys=True) + '\n')
    print(f'[saved] {args.output}', flush=True)
    for metric in metrics:
        print(f"[mean] {metric:<16s} AUROC={results['mean'][metric + '_auroc']:.4f}±{results['std'][metric + '_auroc']:.4f} AUPRC={results['mean'][metric + '_auprc']:.4f}±{results['std'][metric + '_auprc']:.4f}", flush=True)

def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--save-dir', type=Path, required=True)
    parser.add_argument('--target-store', type=Path, required=True)
    parser.add_argument('--output', type=Path, required=True)
    parser.add_argument('--n-benign', type=int, default=10000)
    parser.add_argument('--n-attack', type=int, default=10000)
    parser.add_argument('--seeds', type=int, nargs='+', default=[0, 1, 2, 3, 4])
    parser.add_argument('--benign-label', type=str, default='normal')
    parser.add_argument('--min-len', type=int, default=2)
    parser.add_argument('--n-steps', type=int, default=16)
    parser.add_argument('--batch', type=int, default=4096)
    parser.add_argument('--materialize-batch', type=int, default=32768)
    parser.add_argument('--device', type=str, default='auto')
    run(parser.parse_args())
if __name__ == '__main__':
    main()
