#!/bin/bash
# Wait for user's merge_shard_artifacts.py to finish, then run transfer eval.
#
# The previous auto_transfer_ddos.sh aborted at STAGE 3 because
# np.savez_compressed on 50+ GB array produces a corrupt zip. User's
# merge_shard_artifacts.py uses uncompressed np.savez and is currently
# running. We just wait for it, validate, then do the 4-cell transfer.
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"

# Append to existing log.
exec > >(tee -a "$LOG") 2>&1

N_VAL=20000
N_ATK=100000
SPLIT_SEED=42

echo ""
echo "=== $(date): [after-merge script] starts ==="
echo "waiting for merge_shard_artifacts.py to finish..."

elapsed=0
while pgrep -f "merge_shard_artifacts" > /dev/null; do
    sleep 30
    elapsed=$((elapsed + 30))
    if (( elapsed % 180 == 0 )); then
        rss=$(pgrep -af "merge_shard_artifacts" | head -1 | awk '{print $1}' | xargs -I{} ps -p {} -o rss= 2>/dev/null || echo "?")
        echo "[heartbeat $(date +%H:%M:%S)] merge still running, waited ${elapsed}s  rss=${rss} kB"
    fi
done
echo "=== $(date): merge exited after ${elapsed}s ==="
sleep 10

PACKETS="$DDOS_DIR/packets.npz"
FLOWS="$DDOS_DIR/flows.parquet"

# ---- validate ----
if [[ ! -f "$PACKETS" || ! -f "$FLOWS" ]]; then
    echo "ERROR: artifacts missing after merge"
    echo "  $PACKETS : $([[ -f $PACKETS ]] && echo OK || echo MISSING)"
    echo "  $FLOWS   : $([[ -f $FLOWS   ]] && echo OK || echo MISSING)"
    exit 1
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: validation failed"
    exit 2
fi

# ---- transfer detect + per_class per cell ----
echo ""
echo "=== $(date): 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"
        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
    [[ -f "$out_dir/scores.npz" ]] || { echo "ERROR: no scores.npz for $cell"; return 3; }

    echo "[per_class] starting"
    uv run python -m eval.per_class --save-dir "$out_dir" 2>&1 | tail -60
    [[ -f "$out_dir/per_class.json" ]] || { echo "ERROR: no per_class.json for $cell"; return 4; }
    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"); fi
done

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

# ---- summary ----
echo ""
echo "=== $(date): 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 = {}
for c in cells:
    jp = dst / c / "per_class.json"
    if not jp.exists():
        print(f"{c:<15s}  (missing)"); continue
    tn = json.loads(jp.read_text())[ch]
    loaded[c] = tn
    line = f"{c:<15s}"
    for k,_ in keys:
        v = tn[k]
        line += f" {'NaN':>14s}" if (isinstance(v,float) and v!=v) else f" {v:>14.4f}"
    line += f" {str(tn['flipped']):>8s}"
    print(line)

if loaded:
    print("\n=== 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}: {len(pair)} seed(s), skip"); continue
        print(f"\n--- σ={sg} ---")
        for k,t in keys:
            vs = [d[k] for _,d in pair if isinstance(d[k],float) and d[k]==d[k]]
            if vs:
                print(f"  {t:<18s}  {np.mean(vs):.4f} ± {np.std(vs,ddof=0):.4f}")

    ref = loaded.get("seed42/s0.6") or next(iter(loaded.values()))
    pc = ref["per_class"]
    print("\n=== per-class (seed42/s0.6 reference) ===")
    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}")

    # Merged-label view.
    print("\n=== per-class after label merge (DrDoS_* → stripped) ===")
    def norm(name):
        return name[len("DrDoS_"):] if name.startswith("DrDoS_") else name
    buckets = {}
    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':>8s}  {'n_total':>8s}  "
              f"{'auroc_wtd':>10s}  {'auprc_wtd':>10s}")
        print("-"*68)
        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_a = sum(r["auroc"] * r["n"] for r in rs) / max(n_tot, 1)
            wtd_ap = sum(r["auprc"] * r["n"] for r in rs) / max(n_tot, 1)
            print(f"  {k:<18s}  {len(rs):>8d}  {n_tot:>8d}  "
                  f"{wtd_a:>10.4f}  {wtd_ap:>10.4f}")
EOF

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