#!/bin/bash
# Run all missing cross-direction evals for A+C combo.
# Targets are routed to packets-npz or full_store as appropriate.

set -e
ROOT=/home/chy/mambafortrafficmodeling
EVAL=${ROOT}/Mixed_CFM/eval_cross.py
CROSS_DIR=${ROOT}/artifacts/route_comparison/cross
mkdir -p ${CROSS_DIR}

# Target paths
TGT_iscxtor2016_npz=${ROOT}/datasets/iscxtor2016/processed/packets.npz
TGT_iscxtor2016_flows=${ROOT}/datasets/iscxtor2016/processed/flows.parquet
TGT_iscxtor2016_features=${ROOT}/datasets/iscxtor2016/processed/flow_features.parquet
TGT_iscxtor2016_label=nontor
TGT_iscxtor2016_natk=1888

TGT_cicids2017_store=${ROOT}/datasets/cicids2017/processed/full_store
TGT_cicids2017_flows=${ROOT}/datasets/cicids2017/processed/flows.parquet
TGT_cicids2017_features=${ROOT}/datasets/cicids2017/processed/flow_features.parquet
TGT_cicids2017_label=normal

TGT_cicddos2019_store=${ROOT}/datasets/cicddos2019/processed/full_store
TGT_cicddos2019_flows=${ROOT}/datasets/cicddos2019/processed/flows.parquet
TGT_cicddos2019_features=${ROOT}/datasets/cicddos2019/processed/flow_features.parquet
TGT_cicddos2019_label=normal

TGT_ciciot2023_store=${ROOT}/datasets/ciciot2023/processed/full_store
TGT_ciciot2023_flows=${ROOT}/datasets/ciciot2023/processed/full_store/flows.parquet
TGT_ciciot2023_features=${ROOT}/datasets/ciciot2023/processed/flow_features.parquet
TGT_ciciot2023_label=normal

run_one() {
  local gpu=$1 src=$2 tgt=$3 seed=$4
  local md=${ROOT}/artifacts/route_comparison/route_ac_combo_${src}_seed${seed}
  local out=${CROSS_DIR}/route_ac_combo_seed${seed}_${src}_to_${tgt}.json
  if [ -f "${out}" ]; then echo "[skip] ${src}→${tgt} seed${seed}"; return; fi
  if [ ! -f "${md}/model.pt" ]; then echo "[missing] ${md}/model.pt"; return; fi

  # Resolve target args
  local tgt_args
  if [ "${tgt}" = "iscxtor2016" ]; then
    tgt_args="--target-packets-npz ${TGT_iscxtor2016_npz} --target-flows ${TGT_iscxtor2016_flows} --target-flow-features ${TGT_iscxtor2016_features} --benign-label nontor --n-attack 1888"
  elif [ "${tgt}" = "cicids2017" ]; then
    tgt_args="--target-store ${TGT_cicids2017_store} --target-flows ${TGT_cicids2017_flows} --target-flow-features ${TGT_cicids2017_features} --benign-label normal --n-attack 10000"
  elif [ "${tgt}" = "cicddos2019" ]; then
    tgt_args="--target-store ${TGT_cicddos2019_store} --target-flows ${TGT_cicddos2019_flows} --target-flow-features ${TGT_cicddos2019_features} --benign-label normal --n-attack 10000"
  elif [ "${tgt}" = "ciciot2023" ]; then
    tgt_args="--target-store ${TGT_ciciot2023_store} --target-flows ${TGT_ciciot2023_flows} --target-flow-features ${TGT_ciciot2023_features} --benign-label normal --n-attack 10000"
  fi

  echo "[gpu${gpu}] ${src} → ${tgt} seed${seed}"
  cd ${ROOT}/Mixed_CFM
  CUDA_VISIBLE_DEVICES=${gpu} stdbuf -oL uv run --no-sync python -u ${EVAL} \
    --model-dir ${md} \
    ${tgt_args} \
    --out ${out} \
    --n-benign 10000 --seed 42 --T 64 --batch-size 256 --n-steps 16 \
    > ${CROSS_DIR}/route_ac_combo_seed${seed}_${src}_to_${tgt}.log 2>&1
}

# 8 missing directions × 3 seeds = 24 evals
# Split across 2 GPUs to balance load
{
for dir in "ciciot2023:iscxtor2016" "cicids2017:iscxtor2016" "cicddos2019:iscxtor2016" "iscxtor2016:cicids2017"; do
  src=${dir%:*}; tgt=${dir#*:}
  for seed in 42 43 44; do
    run_one 0 ${src} ${tgt} ${seed}
  done
done
echo "[gpu0 done]"
} > /tmp/cross_matrix_gpu0.log 2>&1 &
G0=$!

{
for dir in "cicids2017:ciciot2023" "cicddos2019:ciciot2023" "iscxtor2016:cicddos2019" "iscxtor2016:ciciot2023"; do
  src=${dir%:*}; tgt=${dir#*:}
  for seed in 42 43 44; do
    run_one 1 ${src} ${tgt} ${seed}
  done
done
echo "[gpu1 done]"
} > /tmp/cross_matrix_gpu1.log 2>&1 &
G1=$!

wait $G0
wait $G1
echo "[all done]"
