#!/usr/bin/env bash
# Stage-2 latent backtracking sweep: 8 combinations, one per GPU, from the
# Stage-1 GRPO checkpoint. Compares no-backtrack controls vs backtracking
# (rehearsal of stage-1 consistency with k=1) and warm-up on/off.
set -uo pipefail

PY=/opt/pytorch/bin/python
REPO=/home/ubuntu/curriculum-cot-code
SCRIPT="${REPO}/latent_multi_output_cell_policy/sft_latent_multi_output_train.py"
S1_GRPO=/home/ubuntu/hf_checkpoints/latent_stages/stage01_latent_grpo_i1_20empty_latent_recurrent_hidden
TRAIN="${REPO}/data/sudoku_t3_20empty_value_qwen_text_stage1_train.jsonl"
EVAL="${REPO}/data/sudoku_t3_20empty_value_qwen_text_stage1_eval.jsonl"
ROOT=/home/ubuntu/bt_runs/s2_sweep_$(date +%Y%m%d_%H%M%S)
mkdir -p "${ROOT}/logs"
echo "Sweep root: ${ROOT}"

# Common knobs (kept small enough for an 8-way single-GPU sweep).
COMMON=(
  --model_name Qwen/Qwen2.5-1.5B-Instruct
  --train_jsonl "${TRAIN}" --eval_jsonl "${EVAL}"
  --cache_dir /home/ubuntu/.hf_cache
  --init_adapter_dir "${S1_GRPO}"
  --seed 0 --gpu_id 0
  --stage_i 2 --num_cot_tokens 2 --latent_mode recurrent_hidden
  --total_empties_hint 20
  --per_device_train_batch_size 8 --gradient_accumulation_steps 2
  --num_epochs 8 --learning_rate 5e-5 --enable_gradient_checkpointing
  --logging_steps 20 --eval_steps 400 --save_steps 400
  --eval_rows 20 --max_completion_length 24 --limit_train_rows 4000
  --max_steps 1200
  --eval_value_precision_stop 0 --eval_value_recall_stop 0 --eval_solve_rate_stop 0
  --lora_r 32 --lora_alpha 64 --lora_dropout 0.05
)

# name|extra flags
declare -a COMBOS=(
  "ctrl_nobt|--latent_warmup_steps 0"
  "ctrl_warm|--latent_warmup_steps 150"
  "bt_rr03_warm|--backtrack_enable --remember_rate 0.3 --backtrack_detect_threshold 0 --latent_warmup_steps 150"
  "bt_rr05_warm|--backtrack_enable --remember_rate 0.5 --backtrack_detect_threshold 0 --latent_warmup_steps 150"
  "bt_rr03_adapt|--backtrack_enable --remember_rate 0.3 --backtrack_detect_threshold 0.95 --latent_warmup_steps 150"
  "bt_rr05_adapt|--backtrack_enable --remember_rate 0.5 --backtrack_detect_threshold 0.97 --latent_warmup_steps 150"
  "bt_rr03_nowarm|--backtrack_enable --remember_rate 0.3 --backtrack_detect_threshold 0.95 --latent_warmup_steps 0"
  "bt_rr02_adapt|--backtrack_enable --remember_rate 0.2 --backtrack_detect_threshold 0.97 --latent_warmup_steps 150"
)

launch() {
  local gpu="$1" name="$2" extra="$3"
  local out="${ROOT}/${name}"
  local log="${ROOT}/logs/${name}.log"
  mkdir -p "${out}"
  echo "[gpu ${gpu}] ${name} :: ${extra}"
  CUDA_VISIBLE_DEVICES="${gpu}" nohup "${PY}" "${SCRIPT}" "${COMMON[@]}" \
    --output_dir "${out}" --wandb_mode disabled ${extra} > "${log}" 2>&1 &
  echo "$!" > "${ROOT}/logs/${name}.pid"
}

# Launch combo 0 first and wait for the shared stage-2 data cache + pools to be
# built, so the other 7 jobs hit a warm cache (no concurrent .tmp write race).
IFS='|' read -r name0 extra0 <<< "${COMBOS[0]}"
launch 0 "${name0}" "${extra0}"
echo "Waiting for stage-2 data cache to build (combo 0)..."
for _ in $(seq 1 120); do
  if grep -q "latent sft train step 00001" "${ROOT}/logs/${name0}.log" 2>/dev/null; then
    echo "Cache warm; launching remaining combos."
    break
  fi
  sleep 5
done

for i in $(seq 1 7); do
  IFS='|' read -r name extra <<< "${COMBOS[$i]}"
  launch "${i}" "${name}" "${extra}"
  sleep 1
done

echo "All combos launched under ${ROOT}"
echo "${ROOT}" > /home/ubuntu/bt_runs/LATEST_S2_SWEEP