#!/bin/bash
#SBATCH --partition=songmei
#SBATCH --nodelist=feanor
#SBATCH --gres=gpu:1
#SBATCH --cpus-per-task=8
#SBATCH --mem=64G
#SBATCH --time=48:00:00
#SBATCH --job-name=w12_s2up
#SBATCH --output=/scratch/users/gatmiry/llm-reasoning-logic-puzzles/sudoku-code/wavecurriculum_run/logs/slurm_%x_%j.out
#SBATCH --error=/scratch/users/gatmiry/llm-reasoning-logic-puzzles/sudoku-code/wavecurriculum_run/logs/slurm_%x_%j.err

# Multi-stage continuation. Starts from the stage-1 run's checkpoint and walks
# stage 2, 3, 4, ... Each stage keeps the SAME sampling structure as stage 1:
# every (puzzle, instance) pair of that stage is enumerated and walked 5 times,
# one instance of a given puzzle per batch. On promotion the pair list is
# rebuilt for the new stage, since each stage has its own instances.
#
# A stage is only left when all three hold:
#   1. every one of its (puzzle, instance) pairs has been seen 5 times
#   2. mass*   >= SUDOKU_PROMOTE_ACC     (support inside the candidate set)
#   3. spread* >= SUDOKU_PROMOTE_SPREAD  (uniform over the candidate set)
# with SUDOKU_MAX_STAGE_STEPS as a cap so an unreachable gate cannot stall the
# whole run.
#
# Runs on a SECOND GPU. It must never touch the stage-1 job's files: no
# rm -rf of /tmp/sudoku_wave_runs, and the stage-1 checkpoint is copied out
# before being restored.

set -u
hostname
nvidia-smi -L
echo "[$(date)] w12_s2plus"

SCRATCH_ROOT=/scratch/users/gatmiry/llm-reasoning-logic-puzzles
RUN_DIR=${SCRATCH_ROOT}/sudoku-code/wavecurriculum_run
ENV_LOCAL=/tmp/logicpuzzles
CAND_DIR=/tmp/sudoku_s12
INST_DIR=/tmp/sudoku_superposition
STAGE1_DIR=/tmp/sudoku_wave_runs/w12_s1_super      # the running job -- READ ONLY
LOCAL_LOG=/tmp/sudoku_wave_runs/w12_s2plus
SEED_CKPT=/tmp/sudoku_wave_runs/w12_s2plus_seed
LOGIN_STAGE=gandalf.berkeley.edu:/tmp/sudoku_hf_s2plus

mkdir -p "${RUN_DIR}/logs" "${LOCAL_LOG}" "${SEED_CKPT}"
echo "[$(date)] /tmp:"
df -h /tmp | tail -1

SCP_OPTS="-o IdentitiesOnly=yes -o StrictHostKeyChecking=accept-new"
[ -f "${HOME}/.ssh/id_ed25519_berkeley" ] && SCP_OPTS="${SCP_OPTS} -i ${HOME}/.ssh/id_ed25519_berkeley"

if ${ENV_LOCAL}/bin/python -u -c "import jax; assert jax.default_backend()=='gpu' or 'cuda' in str(jax.devices()[0]).lower()" 2>/dev/null; then
  echo "[$(date)] reusing ${ENV_LOCAL}"
else
  echo "ERROR: ${ENV_LOCAL} missing; the stage-1 job should have created it" >&2
  exit 1
fi

export PY=${ENV_LOCAL}/bin/python
export LD_LIBRARY_PATH=\
${ENV_LOCAL}/lib/python3.9/site-packages/nvidia/cudnn/lib:\
${ENV_LOCAL}/lib/python3.9/site-packages/nvidia/cublas/lib:\
${ENV_LOCAL}/lib/python3.9/site-packages/nvidia/cuda_runtime/lib:\
${ENV_LOCAL}/lib/python3.9/site-packages/nvidia/cuda_nvrtc/lib:\
${ENV_LOCAL}/lib/python3.9/site-packages/nvidia/nccl/lib:\
${LD_LIBRARY_PATH:-}

${PY} -u -c "import jax; print(jax.__version__, jax.devices(), jax.default_backend())"

for f in train_assignments.npy train_starts.npy train_counts.npy \
         test_assignments.npy test_starts.npy test_counts.npy; do
  [ -s "${INST_DIR}/${f}" ] || { echo "missing ${INST_DIR}/${f}" >&2; exit 1; }
done
for f in train_cand_masks.npy test_cand_masks.npy; do
  [ -s "${CAND_DIR}/${f}" ] || { echo "missing ${CAND_DIR}/${f}" >&2; exit 1; }
done

# ---- Copy the stage-1 checkpoint out before restoring ----
# The stage-1 job rolls its checkpoints (keep=2) while running, so restore from
# a private copy rather than reading a directory that is being rewritten.
latest=$(ls -d ${STAGE1_DIR}/checkpoint_* 2>/dev/null | sort -t_ -k2 -n | tail -1)
if [ -z "${latest}" ]; then
  echo "ERROR: no checkpoint found in ${STAGE1_DIR}" >&2
  ls -la "${STAGE1_DIR}" >&2 || true
  exit 1
fi
echo "[$(date)] seeding from ${latest}"
cp -r "${latest}" "${SEED_CKPT}/" 2>/dev/null || true
ls -la "${SEED_CKPT}"

# ---- Stage 2 onward, same instance schedule as stage 1 ----
# ckpt_loc is a CLI flag, not an env var; SUDOKU_RESUME only turns restore on.
export SUDOKU_RESUME=1
export SUDOKU_START_STAGE=2
export SUDOKU_MAX_STAGE=12
export SUDOKU_LATENT_SLOTS=12
export SUDOKU_RECURRENT=1
export SUDOKU_BACKTRACK=0
export SUDOKU_CAND_SLOT_MODE=depth
export SUDOKU_PASSES_PER_STAGE=1
export SUDOKU_AUX_WEIGHT=0.0
export SUDOKU_LEVEL_BALANCED=0
export SUDOKU_DATA_CURRICULUM=none
# Same structure as stage 1: every (puzzle, instance) pair of the CURRENT stage
# seen 5 times. 200k puzzles keeps one stage near 2 h so several stages land
# overnight; stage 2 has ~5.87 instances/puzzle = 1.17M pairs = 92k steps.
export SUDOKU_INSTANCE_EPOCHS=5
export SUDOKU_INSTANCE_PUZZLES="${SUDOKU_INSTANCE_PUZZLES:-200000}"
# Quality gates, both required on top of the 5-pass data gate.
export SUDOKU_PROMOTE_ACC=0.85
export SUDOKU_PROMOTE_SPREAD=0.85
export SUDOKU_PROMOTE_LOC_WAVE=0.0
export SUDOKU_PLATEAU_STEPS=0
export SUDOKU_PATIENCE=0
export SUDOKU_MIN_STAGE_STEPS=8000
# Cap per stage: above the ~92k steps that 5 passes need at stage 2, so the
# data gate always gets to finish, but bounded so one stage cannot eat the run.
export SUDOKU_MAX_STAGE_STEPS="${SUDOKU_MAX_STAGE_STEPS:-130000}"
# Resume starts the loop at the restored step (~250k+), so leave plenty of room.
export SUDOKU_MAX_STEPS="${SUDOKU_MAX_STEPS:-1400000}"
export SUDOKU_EVAL_EVERY=2000
export SUDOKU_SAVE_EVERY=10000
export SUDOKU_CKPT_KEEP=2
export SUDOKU_LR=0.0002
export SUDOKU_DROPOUT=0.2
export SUDOKU_WD=0.005
export SUDOKU_TRAIN_PATH="${SCRATCH_ROOT}/sudoku-code/datasets/train_sudoku_puzzles.npy"
export SUDOKU_TEST_PATH="${SCRATCH_ROOT}/sudoku-code/datasets/test_sudoku_puzzles.npy"
export SUDOKU_TRAIN_CAND="${CAND_DIR}/train_cand_masks.npy"
export SUDOKU_TEST_CAND="${CAND_DIR}/test_cand_masks.npy"
export SUDOKU_INSTANCE_DIR="${INST_DIR}"
# Two jobs share the node's GPUs, so take a smaller slice of this one.
export XLA_PYTHON_CLIENT_MEM_FRACTION=0.80

HF_TOKEN_FILE="${HF_TOKEN_FILE:-/scratch/users/gatmiry/.hf_token}"
if [ -z "${HF_TOKEN:-}" ] && [ -s "${HF_TOKEN_FILE}" ]; then
  HF_TOKEN="$(cat "${HF_TOKEN_FILE}")"
  export HF_TOKEN
fi
export HUGGING_FACE_HUB_TOKEN="${HF_TOKEN:-}"
HF_PKGS="${HF_PKGS:-/scratch/users/gatmiry/hf_pkgs}"
TRAIN_LOG="${RUN_DIR}/logs/w12_s2plus.log"
SYNC_LOG="${RUN_DIR}/logs/hf_sync_w12_s2plus.log"

cd "${RUN_DIR}"
PYTHONPATH="${HF_PKGS}${PYTHONPATH:+:${PYTHONPATH}}" ${PY} -u "${RUN_DIR}/hf_sync.py" \
  --workdir "${LOCAL_LOG}" \
  --log "${TRAIN_LOG}" \
  --scratch-log "${TRAIN_LOG}" \
  --repo Avra98/Sudoku_superposition \
  --prefix runs/w12_s2plus \
  --login-stage "${LOGIN_STAGE}" \
  --token-file "${HF_TOKEN_FILE}" \
  --interval 60 \
  > "${SYNC_LOG}" 2>&1 &
SYNC_PID=$!

echo "[$(date)] starting w12_s2plus"
echo "  resume from ${SEED_CKPT}, curriculum starts at stage 2"
echo "  ${SUDOKU_INSTANCE_PUZZLES} puzzles, every instance of each stage seen 5x"
echo "  gates: mass*>=${SUDOKU_PROMOTE_ACC} spread*>=${SUDOKU_PROMOTE_SPREAD}"
echo "  hf=Avra98/Sudoku_superposition/runs/w12_s2plus"
CUDA_VISIBLE_DEVICES=0 ${PY} -u -m train.main \
  --workdir="${LOCAL_LOG}" --exp_name="w12_s2plus" \
  --ckpt_loc="${SEED_CKPT}" \
  > "${TRAIN_LOG}" 2>&1
EC=$?
kill ${SYNC_PID} 2>/dev/null || true
wait ${SYNC_PID} 2>/dev/null || true
PYTHONPATH="${HF_PKGS}${PYTHONPATH:+:${PYTHONPATH}}" ${PY} -u "${RUN_DIR}/hf_sync.py" --once \
  --workdir "${LOCAL_LOG}" \
  --log "${TRAIN_LOG}" \
  --scratch-log "${TRAIN_LOG}" \
  --repo Avra98/Sudoku_superposition \
  --prefix runs/w12_s2plus \
  --login-stage "${LOGIN_STAGE}" \
  --token-file "${HF_TOKEN_FILE}" \
  >> "${SYNC_LOG}" 2>&1 || true
echo "[$(date)] w12_s2plus exit ${EC}"
exit ${EC}