#!/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}" "${LOCAL_LOG}/_hf_tars"
# Only this run's own directory may be cleared. The siblings under
# /tmp/sudoku_wave_runs are live checkpoints of the other jobs on this node.
rm -rf "${LOCAL_LOG:?}"/checkpoint_*
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 "[$(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
# Instance arm never promotes on plateau or patience. Both gates must clear.
export SUDOKU_PLATEAU_STEPS=0
export SUDOKU_PATIENCE=0
export SUDOKU_MIN_STAGE_STEPS=8000
# Gate 1: P(digit in raw candidate set S) on |S|>=2 cells. Chance ~0.41.
export SUDOKU_PROMOTE_ACC=0.85
# Gate 2: H(p restricted to S) / log|S|. 1.0 = uniform over S, 0 = collapsed.
# S is the wave-solver candidate set, not the filtered-instance support.
export SUDOKU_PROMOTE_SPREAD=0.85
export SUDOKU_PROMOTE_LOC_WAVE=0.0
export SUDOKU_MAX_STEPS="${SUDOKU_MAX_STEPS:-800000}"
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}"
# Leave a little H200 headroom so eval + sidecar tar do not OOM the node.
export XLA_PYTHON_CLIENT_MEM_FRACTION=0.85

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_inst_latent.log"
SYNC_LOG="${RUN_DIR}/logs/hf_sync_w12_inst_latent.log"

# Need room for 2 rolling ckpts (~1 GB) plus a tar in flight.
tmp_avail_kb=$(df -Pk /tmp | awk 'NR==2{print $4}')
echo "[$(date)] /tmp avail ${tmp_avail_kb} KB"
if [ "${tmp_avail_kb}" -lt 3000000 ]; then
  echo "ERROR: /tmp has less than 3G free; refusing to start" >&2
  df -h /tmp
  du -sh /tmp/* 2>/dev/null | sort -h | tail -20 >&2 || true
  exit 1
fi

cd "${RUN_DIR}"
# huggingface_hub lives on scratch (feanor /tmp is too full to pip-install).
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_inst_latent \
  --token-file "${HF_TOKEN_FILE}" \
  --interval 60 \
  > "${SYNC_LOG}" 2>&1 &
SYNC_PID=$!

echo "[$(date)] starting w12_inst_latent from scratch"
echo "  K=12 recurrent=1 bt=0 aux=0 instance_dir=${INST_DIR}"
echo "  workdir=${LOCAL_LOG}  log=${TRAIN_LOG}"
echo "  hf=Avra98/Sudoku_superposition/runs/w12_inst_latent"
# Log goes to scratch immediately so a /tmp-full crash still leaves a traceback.
CUDA_VISIBLE_DEVICES=0 ${PY} -u -m train.main \
  --workdir="${LOCAL_LOG}" --exp_name="w12_inst_latent" \
  > "${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_inst_latent \
  --token-file "${HF_TOKEN_FILE}" \
  >> "${SYNC_LOG}" 2>&1 || true
echo "[$(date)] w12_inst_latent exit ${EC}"
exit ${EC}