#!/bin/bash
#SBATCH --partition=songmei
#SBATCH --nodelist=feanor
#SBATCH --gres=gpu:1
#SBATCH --cpus-per-task=8
#SBATCH --mem=64G
#SBATCH --time=72:00:00
#SBATCH --job-name=w12_inst
#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

# Same recipe as the slimgpu w12_inst_latent run:
#   K=12 recurrent latent curriculum, no backtrack, no candidate-head BCE.
#   Output prompt = one sampled stage-k instance (CE against that sequence).
# Instance npy files and s12 masks live on feanor /tmp (scratch quota is 20 G).

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

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
LOCAL_LOG=/tmp/sudoku_wave_runs/w12_inst_latent
TARBALL_GANDALF=/tmp/logicpuzzles_env.tar.gz
TARBALL_LOCAL=/tmp/logicpuzzles_env_${SLURM_JOB_ID}.tar.gz
GANDALF_INST=gandalf.berkeley.edu:/tmp/sudoku_superposition
GANDALF_CAND=gandalf.berkeley.edu:/tmp/sudoku_s12

mkdir -p "${RUN_DIR}/logs" "${LOCAL_LOG}" /tmp/sudoku_wave_runs

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 "[$(date)] fetching env tarball"
  rm -rf "${ENV_LOCAL}"
  scp ${SCP_OPTS} "gandalf.berkeley.edu:${TARBALL_GANDALF}" "${TARBALL_LOCAL}"
  tar xzf "${TARBALL_LOCAL}" -C /tmp
  rm -f "${TARBALL_LOCAL}"
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())"

# ---- Stage instance files onto node-local /tmp (scratch cannot hold 8 G) ----
need_inst=0
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}" ] || need_inst=1
done
if [ "${need_inst}" = 1 ]; then
  echo "[$(date)] pulling instances from ${GANDALF_INST}"
  mkdir -p "${INST_DIR}"
  rsync -a --progress -e "ssh ${SCP_OPTS}" \
    "${GANDALF_INST}/" "${INST_DIR}/"
fi
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

need_cand=0
for f in train_cand_masks.npy test_cand_masks.npy; do
  [ -s "${CAND_DIR}/${f}" ] || need_cand=1
done
if [ "${need_cand}" = 1 ]; then
  echo "[$(date)] pulling s12 masks from ${GANDALF_CAND}"
  mkdir -p "${CAND_DIR}"
  rsync -a -e "ssh ${SCP_OPTS}" "${GANDALF_CAND}/" "${CAND_DIR}/" || true
fi
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

# ---- Exact same recipe as slimgpu w12_inst_latent ----
export SUDOKU_RESUME=0
export SUDOKU_START_STAGE=1
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
export SUDOKU_PLATEAU_STEPS=20000
export SUDOKU_PLATEAU_DELTA=0.005
export SUDOKU_PATIENCE=80000
export SUDOKU_MIN_STAGE_STEPS=8000
# Stage 1->2: inset >= 0.85 (chance is ~0.41 on stage-0 sets) AND
# loc_acc >= 0.70. Frontier is min(inset, loc), so both must clear 0.85
# unless loc is the limiter (then 0.85 loc is required). The loc floor
# is applied via min(); SUDOKU_PROMOTE_LOC is the documented intent.
export SUDOKU_PROMOTE_ACC=0.85
export SUDOKU_PROMOTE_LOC=0.70
export SUDOKU_MAX_STEPS="${SUDOKU_MAX_STEPS:-800000}"
export SUDOKU_EVAL_EVERY=2000
export SUDOKU_SAVE_EVERY=10000
export SUDOKU_CKPT_KEEP=3
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}"
export XLA_PYTHON_CLIENT_MEM_FRACTION=0.9

cd "${RUN_DIR}"
(
  while true; do sleep 900
    rsync -a "${LOCAL_LOG}.log" "${RUN_DIR}/logs/w12_inst_latent.log" 2>/dev/null || true
  done
) &
SYNC_PID=$!

echo "[$(date)] starting w12_inst_latent from scratch"
echo "  K=12 recurrent=1 bt=0 aux=0 instance_dir=${INST_DIR}"
CUDA_VISIBLE_DEVICES=0 ${PY} -u -m train.main \
  --workdir="${LOCAL_LOG}" --exp_name="w12_inst_latent" \
  > "${LOCAL_LOG}.log" 2>&1
EC=$?
kill ${SYNC_PID} 2>/dev/null || true
rsync -a "${LOCAL_LOG}.log" "${RUN_DIR}/logs/w12_inst_latent.log" 2>/dev/null || true
echo "[$(date)] w12_inst_latent exit ${EC}"
exit ${EC}