#!/bin/bash
# Unattended cross-dataset transfer eval (v2, with per-shard merge).
#
# Pipeline:
#   STAGE 1 : wait for 01-12 re-extraction to finish
#   STAGE 2 : merge packets.{01-12,03-11}.npz + flows.{01-12,03-11}.parquet
#             → unified packets.npz + flows.parquet
#   STAGE 3 : validate unified artifacts
#   STAGE 4 : detect + per_class across 4 cells (seed × σ from CICIDS2017 sweep)
#   STAGE 5 : summary table + merged-label view
#
# Log: runs/transfer_ddos/run.log
set -uo pipefail

ROOT=/home/chy/mambafortrafficmodeling
cd "$ROOT"

SRC_SWEEP="runs/n10k_refactor_20260422_220351"
DST="runs/transfer_ddos"
LOG="$DST/run.log"
DDOS_DIR="datasets/cicddos2019/processed"
mkdir -p "$DST"
exec > >(tee -a "$LOG") 2>&1

N_VAL=20000
N_ATK=100000
SPLIT_SEED=42

echo "=== $(date): script starts (v2 with merge) ==="
echo "source sweep : $SRC_SWEEP"
echo "destination  : $DST"
echo "scoring      : n_val=$N_VAL  n_atk=$N_ATK  split_seed=$SPLIT_SEED"

# =====================================================================
# STAGE 1: wait for extraction_cicddos2019 (01-12 shard) to finish
# =====================================================================
echo ""
echo "=== $(date): STAGE 1 — waiting for 01-12 re-extraction ==="
elapsed=0
while pgrep -f "scripts/extract_cicddos2019" > /dev/null; do
    sleep 60
    elapsed=$((elapsed + 60))
    if (( elapsed % 600 == 0 )); then
        # Heartbeat every 10 min
        rss=$(pgrep -af "scripts/extract_cicddos2019" | head -1 | awk '{print $1}' | xargs -I{} ps -p {} -o rss= 2>/dev/null || echo "?")
        echo "[heartbeat $(date +%H:%M:%S)] 01-12 extraction running, waited ${elapsed}s  rss(parent)=${rss} kB"
    fi
done
echo "=== $(date): extraction process exited after ${elapsed}s wait ==="
sleep 15

# =====================================================================
# STAGE 2: merge per-shard artifacts
# =====================================================================
echo ""
echo "=== $(date): STAGE 2 — merge shards ==="

SHARDS_PACKETS="$DDOS_DIR/packets.01-12.npz $DDOS_DIR/packets.03-11.npz"
SHARDS_FLOWS="$DDOS_DIR/flows.01-12.parquet $DDOS_DIR/flows.03-11.parquet"

missing=0
for f in $SHARDS_PACKETS $SHARDS_FLOWS; do
    if [[ ! -f "$f" ]]; then
        echo "ERROR: shard artifact missing: $f"
        missing=1
    fi
done
if (( missing )); then
    echo "--- tail of 01-12 extraction log ---"
    tail -40 runs/extract_logs/extract_ddos_0112.log 2>&1 || true
    exit 1
fi

echo "all 4 shard artifacts present; running merge"
if ! uv run python scripts/merge_cicddos_shards.py 2>&1; then
    echo "ERROR: merge failed"
    exit 2
fi

# =====================================================================
# STAGE 3: validate unified artifacts
# =====================================================================
echo ""
echo "=== $(date): STAGE 3 — validate unified artifacts ==="
PACKETS="$DDOS_DIR/packets.npz"
FLOWS="$DDOS_DIR/flows.parquet"
if [[ ! -f "$PACKETS" || ! -f "$FLOWS" ]]; then
    echo "ERROR: merge output missing"
    echo "  $PACKETS : $([[ -f $PACKETS ]] && echo OK || echo MISSING)"
    echo "  $FLOWS   : $([[ -f $FLOWS   ]] && echo OK || echo MISSING)"
    exit 3
fi
ls -lh "$PACKETS" "$FLOWS"

uv run python - <<'EOF'
import numpy as np, pandas as pd, sys
try:
    p = np.load('datasets/cicddos2019/processed/packets.npz')
    f = pd.read_parquet('datasets/cicddos2019/processed/flows.parquet')
    assert set(p.files) == {'packet_tokens', 'packet_lengths', 'flow_id'}, p.files
    assert set(f.columns) == {'flow_id', 'label'}, f.columns
    assert p['flow_id'].shape[0] == len(f)
    assert np.array_equal(p['flow_id'], f['flow_id'].to_numpy())
    n = len(f)
    n_benign = int((f['label'] == 'normal').sum())
    print(f'[validate] OK: N={n:,}  benign={n_benign:,}  attack={n-n_benign:,}')
    print(f'[validate] packet_tokens shape/dtype: {p["packet_tokens"].shape} {p["packet_tokens"].dtype}')
    print('[validate] label value_counts (top 20):')
    print(f['label'].value_counts().head(20).to_string())
except Exception as e:
    print(f'[validate] FAILED: {type(e).__name__}: {e}')
    sys.exit(2)
EOF
if [[ $? -ne 0 ]]; then
    echo "ERROR: unified artifact validation failed"
    exit 4
fi

# =====================================================================
# STAGE 4: transfer detect + per_class
# =====================================================================
echo ""
echo "=== $(date): STAGE 4 — transfer detect + per_class (4 cells) ==="
CELLS=( "seed42/s0" "seed42/s0.6" "seed43/s0" "seed43/s0.6" )

run_cell() {
    local cell=$1
    local src_dir="$SRC_SWEEP/$cell"
    local out_dir="$DST/$cell"

    if [[ ! -f "$src_dir/model.pt" ]]; then
        echo "SKIP $cell : no model.pt at $src_dir"
        return 1
    fi

    mkdir -p "$out_dir"
    cp "$src_dir/model.pt" "$out_dir/model.pt"
    [[ -f "$src_dir/history.json" ]] && cp "$src_dir/history.json" "$out_dir/"

    echo ""
    echo "----- [$cell] $(date) -----"
    echo "[detect] starting"
    if ! uv run python -m detect \
        --save-dir "$out_dir" \
        --packets-npz "$PACKETS" \
        --flows-parquet "$FLOWS" \
        --n-val "$N_VAL" --n-atk "$N_ATK" \
        --seed "$SPLIT_SEED" \
        2>&1 | tail -25
    then
        echo "ERROR: detect failed for $cell"
        return 2
    fi

    if [[ ! -f "$out_dir/scores.npz" ]]; then
        echo "ERROR: detect produced no scores.npz for $cell"
        return 3
    fi

    echo "[per_class] starting"
    if ! uv run python -m eval.per_class --save-dir "$out_dir" 2>&1 | tail -80
    then
        echo "ERROR: per_class failed for $cell"
        return 4
    fi

    if [[ ! -f "$out_dir/per_class.json" ]]; then
        echo "ERROR: per_class.json missing for $cell"
        return 5
    fi

    echo "[$cell] OK"
    return 0
}

OK_CELLS=()
FAIL_CELLS=()
for cell in "${CELLS[@]}"; do
    if run_cell "$cell"; then
        OK_CELLS+=("$cell")
    else
        FAIL_CELLS+=("$cell")
        echo "[$cell] continuing despite failure"
    fi
done

echo ""
echo "=== per-cell status ==="
echo "OK   (${#OK_CELLS[@]}): ${OK_CELLS[*]:-none}"
echo "FAIL (${#FAIL_CELLS[@]}): ${FAIL_CELLS[*]:-none}"

# =====================================================================
# STAGE 5: summary
# =====================================================================
echo ""
echo "=== $(date): STAGE 5 — summary ==="

uv run python - "$DST" <<'EOF'
import json, sys
from pathlib import Path
import numpy as np

dst = Path(sys.argv[1])
cells = ["seed42/s0", "seed42/s0.6", "seed43/s0", "seed43/s0.6"]

ch = "terminal_norm"
keys = [("overall_auroc", "overall AUROC"),
        ("overall_auprc", "overall AUPRC"),
        ("macro_auroc",   "macro AUROC"),
        ("macro_auprc",   "macro AUPRC"),
        ("tpr_at_1fpr",   "TPR@1%FPR"),
        ("fpr_at_95tpr",  "FPR@95%TPR")]

header = f"{'cell':<15s}" + "".join(f" {t:>14s}" for _, t in keys) + f" {'flipped':>8s}"
print(header)
print("-" * len(header))

loaded: dict[str, dict] = {}
for c in cells:
    jp = dst / c / "per_class.json"
    if not jp.exists():
        print(f"{c:<15s}  (missing per_class.json)")
        continue
    try:
        tn = json.loads(jp.read_text())[ch]
    except Exception as e:
        print(f"{c:<15s}  (parse error: {e})")
        continue
    loaded[c] = tn
    row = [tn[k] for k, _ in keys]
    line = f"{c:<15s}"
    for v in row:
        if isinstance(v, float) and (v != v):
            line += f" {'NaN':>14s}"
        else:
            line += f" {v:>14.4f}"
    line += f" {str(tn['flipped']):>8s}"
    print(line)

if not loaded:
    print("\n(no cells loaded — nothing to aggregate)")
    sys.exit(0)

print("")
print("=== mean ± std across seeds (per σ) ===")
for sg in ["0", "0.6"]:
    pair = [(c, d) for c, d in loaded.items() if c.endswith(f"/s{sg}")]
    if len(pair) < 2:
        print(f"σ={sg}: only {len(pair)} seed(s), skip aggregate")
        continue
    print(f"\n--- σ={sg} ({len(pair)} seeds) ---")
    for k, t in keys:
        vals = [d[k] for _, d in pair if isinstance(d[k], float) and d[k] == d[k]]
        if not vals:
            continue
        m = float(np.mean(vals)); s = float(np.std(vals, ddof=0))
        print(f"  {t:<18s}  {m:.4f} ± {s:.4f}")

ref = loaded.get("seed42/s0.6") or next(iter(loaded.values()))
print("")
print("=== per-class AUROC (seed42/s0.6 reference) ===")
pc = ref["per_class"]
print(f"{'class':<30s}  {'n':>8s}  {'auroc':>8s}  {'auprc':>8s}  {'tpr@1%':>8s}")
print("-" * 70)
for r in pc:
    fmt = lambda v: "—" if (isinstance(v, float) and v != v) else f"{v:.4f}"
    print(f"  {r['class']:<28s}  {r['n']:>8d}  "
          f"{fmt(r['auroc']):>8s}  {fmt(r['auprc']):>8s}  {fmt(r['tpr_at_1fpr']):>8s}")

print("")
print("=== per-class after label merge (DrDoS_* → stripped) ===")
def norm(name):
    if name.startswith("DrDoS_"):
        return name[len("DrDoS_"):]
    return name
buckets: dict[str, list] = {}
for r in pc:
    if isinstance(r["auroc"], float) and r["auroc"] == r["auroc"]:
        buckets.setdefault(norm(r["class"]), []).append(r)
if buckets:
    print(f"{'merged class':<20s}  {'shards':>6s}  {'n_total':>8s}  "
          f"{'auroc_wtd':>10s}  {'auroc_mean':>10s}")
    print("-" * 64)
    for k, rs in sorted(buckets.items(), key=lambda x: -sum(r["n"] for r in x[1])):
        n_tot = sum(r["n"] for r in rs)
        wtd = sum(r["auroc"] * r["n"] for r in rs) / max(n_tot, 1)
        mean = sum(r["auroc"] for r in rs) / len(rs)
        print(f"  {k:<18s}  {len(rs):>6d}  {n_tot:>8d}  "
              f"{wtd:>10.4f}  {mean:>10.4f}")
EOF

echo ""
echo "=== $(date): script done ==="
