#!/usr/bin/env bash
# 四个 baseline 臂（linear / xattn / prompt / adaln）的单卡端到端 smoke：
#   每臂：单进程 train.py（不走 DeepSpeed）--limit 8 --max_steps 2 → 导出 step-2.safetensors
#         → 用该 ckpt 推理一条 right:1.0:21（20 步）并实测位移 → check_ckpt（CPU）strict 装回
#   最后把每臂的 n_new_params / step_time / 峰值显存 / 实测位移汇总到 outputs/smoke_baselines/summary.json。
# 用法：bash scripts/smoke_baselines.sh                 # 四臂依次在 GPUS（默认 4 5 6 7）上轮流跑
#      ARMS="linear adaln" bash scripts/smoke_baselines.sh
#      PARALLEL=1 bash scripts/smoke_baselines.sh      # 四臂各占一张卡并行（要有四张空闲卡）
#      SUMMARY_ONLY=1 bash scripts/smoke_baselines.sh  # 只重算汇总（各臂产物已在）
# 单卡无 DeepSpeed 时 AdamW 直接更新 bf16 参数，峰值显存 ≈ 47 GB（权重 + 梯度 + 两份 Adam 状态各 10 GB + 激活）。
set -euo pipefail
ROOT=/opt/dlami/nvme/zhiyangdeng/ActionRoPE
PY=$ROOT/.venv/bin/python
export DIFFSYNTH_SKIP_DOWNLOAD=True
export PYTHONPATH="$ROOT"
ARMS="${ARMS:-linear xattn prompt adaln}"
GPUS=(${GPUS:-4 5 6 7})
PARALLEL="${PARALLEL:-0}"
INFER_STEPS="${INFER_STEPS:-20}"
CLIP="${CLIP:-clip_Eybx_200000958_000273}"    # val_eybx，tidal_flats，sidecar valid
OUT=$ROOT/outputs/smoke_baselines
mkdir -p "$OUT" "$ROOT/outputs/samples"
LOG=$OUT/smoke_baselines.log
exec > >(tee -a "$LOG") 2>&1
T_ALL=$(date +%s)

run_arm() {  # run_arm <arm> <gpu>
  local ARM=$1 GPU=$2 T0 RUN CKPT
  T0=$(date +%s)
  RUN=$ROOT/outputs/smoke_${ARM}_1gpu
  mkdir -p "$RUN"
  echo "===== [$ARM] GPU $GPU：单卡 2 步训练 → $RUN ($(date '+%F %T')) ====="
  CUDA_VISIBLE_DEVICES=$GPU "$PY" "$ROOT/actionrope/train.py" \
    --arm "$ARM" --limit 8 --max_steps 2 --save_every 2 --val_every 0 --warmup_steps 0 \
    --num_workers 2 --seed 0 --output "$RUN" 2>&1 | tee "$RUN/train.log"
  CKPT=$RUN/step-2.safetensors
  test -s "$CKPT"
  echo "===== [$ARM] GPU $GPU：step-2 推理 right:1.0:21（$INFER_STEPS 步） ====="
  CUDA_VISIBLE_DEVICES=$GPU "$PY" -m actionrope.infer --ckpt "$CKPT" --first_frame "$CLIP" --scene tidal_flats \
    --actions "right:1.0:21" --steps "$INFER_STEPS" --seed 0 \
    --out "$ROOT/outputs/samples/smoke_${ARM}_right_x1.mp4" 2>&1 | tee "$RUN/infer.log"
  echo "===== [$ARM]：check_ckpt（CPU，臂按 ckpt 自动识别） ====="
  "$PY" "$ROOT/tests/check_ckpt.py" "$CKPT" --json_out "$RUN/check_ckpt.json" > "$RUN/check_ckpt.log" 2>&1
  "$PY" -c "import json; r=json.load(open('$RUN/check_ckpt.json')); assert r['strict_load_ok'] and r['arm']=='$ARM', r; print('check_ckpt', {k: r[k] for k in ('arm','n_tensors','n_arm_tensors','n_new_params','n_changed_tensors','arm_nonzero_tensors') if k in r})"
  echo "----- [$ARM] 通过，耗时 $(( $(date +%s) - T0 )) s -----"
}

i=0
PIDS=()
[ "${SUMMARY_ONLY:-0}" = 1 ] && ARMS_RUN="" || ARMS_RUN="$ARMS"
for ARM in $ARMS_RUN; do
  GPU=${GPUS[$(( i % ${#GPUS[@]} ))]}
  if [ "$PARALLEL" = 1 ]; then
    run_arm "$ARM" "$GPU" > "$OUT/${ARM}.log" 2>&1 &
    PIDS+=($!)
  else
    run_arm "$ARM" "$GPU"
  fi
  i=$((i + 1))
done
if [ "$PARALLEL" = 1 ]; then
  FAIL=0
  for p in "${PIDS[@]}"; do wait "$p" || FAIL=1; done
  for ARM in $ARMS; do cat "$OUT/${ARM}.log"; done
  [ "$FAIL" = 0 ] || { echo "有臂失败"; exit 1; }
fi

# 汇总：每臂 n_new_params / 2 步的 step_time / 峰值显存 / 推理实测位移
"$PY" - "$OUT/summary.json" $ARMS <<'EOF'
import csv, json, sys
ROOT = "/opt/dlami/nvme/zhiyangdeng/ActionRoPE"
out_path, arms = sys.argv[1], sys.argv[2:]
summary = {}
for arm in arms:
    run = f"{ROOT}/outputs/smoke_{arm}_1gpu"
    cfg = json.load(open(f"{run}/config.json"))
    rows = list(csv.DictReader(open(f"{run}/train_log.csv")))
    inf = json.load(open(f"{ROOT}/outputs/samples/smoke_{arm}_right_x1.json"))
    chk = json.load(open(f"{run}/check_ckpt.json"))
    summary[arm] = {
        "n_new_params": cfg["model"]["n_new_params"], "arm_class": cfg["model"]["arm_class"], "arm_kwargs": cfg["model"]["arm_kwargs"],
        "text_mode": cfg["args"]["text_mode"], "train_used": cfg["dataset"]["train_used"],
        "step_time_s": [round(float(r["step_time"]), 2) for r in rows], "loss": [round(float(r["loss"]), 4) for r in rows],
        "grad_norm": [round(float(r["grad_norm"]), 3) for r in rows],
        "peak_mem_gb": round(max(float(r["mem_gb"]) for r in rows), 1),
        "infer_steps": inf["steps"], "infer_sample_sec": round(inf["sample_sec"], 1), "infer_peak_mem_gb": round(inf["peak_mem_gb"], 1),
        "infer_measured_bg_shift_80": inf["measured"]["sift"], "infer_expected_bg_shift_80": inf["expected_bg_shift_80"],
        "infer_error_px": inf["error_px"], "infer_sift_source": inf["measured"]["sift_source"],
        "ckpt_n_tensors": chk["n_tensors"], "ckpt_n_arm_tensors": chk["n_arm_tensors"],
        "ckpt_dit_changed_tensors": chk["n_changed_tensors"], "ckpt_arm_nonzero_tensors": chk.get("arm_nonzero_tensors"),
        "ckpt_size_gb": round(chk["size_gb"], 2),
    }
json.dump(summary, open(out_path, "w"), indent=2, ensure_ascii=False)
print(json.dumps(summary, indent=1, ensure_ascii=False))
EOF
echo "===== smoke_baselines 全部通过（$ARMS），总耗时 $(( $(date +%s) - T_ALL )) s ($(date '+%F %T')) ====="