#!/bin/bash
#SBATCH --partition=songmei
#SBATCH --nodelist=feanor
#SBATCH --gres=gpu:2
#SBATCH --cpus-per-task=16
#SBATCH --mem=96G
#SBATCH --time=72:00:00
#SBATCH --requeue
#SBATCH --job-name=w12_gateq
#SBATCH --output=/scratch/users/gatmiry/llm-reasoning-logic-puzzles/sudoku-code/wavecurriculum_ceq/logs/slurm_%x_%j.out
#SBATCH --error=/scratch/users/gatmiry/llm-reasoning-logic-puzzles/sudoku-code/wavecurriculum_ceq/logs/slurm_%x_%j.err

# Walk stages 4..12 from the verified stage-4 checkpoint, on a CE(q||p) gate
# that is finally measured correctly. Same criterion as sbatch_ceq_multigpu.sh;
# three bugs in how it was READ are fixed, and every one of them was inflating
# the statistic, which is why every stage past 4 stalled.
#
#   prefix  data.py rewrote value tokens with an instance only when train=True,
#           so eval scored p on a solution-valued sequence: every earlier empty
#           cell handed the model the true digit, the current cell fell out by
#           elimination, and the gate then charged that against a q which
#           conditions on nothing of the sort. Instance-valued sequences now
#           come through the batch and p is averaged over SUDOKU_EVAL_READS of
#           them, which is both what training feeds and a marginal comparable
#           to q. Worth ~0.07 nats of the stage-4 hinge.
#   depth   evaluater.py pinned k_budget = K, so one full-depth prediction was
#           compared against every stage's q. The gate pass now activates
#           exactly `stage` slots and rebuilds the latents at that depth.
#   q       the stored pool stops at coverage (`while uncovered`, plus
#           require_new dropping redundant draws), so it is ~5 grids per puzzle
#           -- 1.4 by stage 11 -- and it deliberately spends those draws on
#           rare candidates, which flattens the histogram. That is fatal for a
#           gate whose headroom is log|S| - H(q): re-sampling the same
#           generator 200x per cell moves q(true) from 0.40 to 0.63 at stage 4
#           and H(q) from 0.96 to 0.50, i.e. real headroom 0.60 nats, not 0.15.
#           SUDOKU_GATE_Q points at gen_bigq.py output and the eval split is
#           restricted to the puzzles it covers.
#
# Measured on the stage-4..11 checkpoints of the ungated w12_full12 run, the
# hinge goes from 0.07-0.46 (fail everywhere) to 0.000-0.002 (pass everywhere).
# Verified in the training loop itself: stage 4 reads hinge 0.0026 against a
# 0.05 bar and promotes.
#
# Submit with:
#   sbatch --gres=gpu:2 --job-name=w12_gateq \
#     --export=ALL,CEQ_TAG=w12_gateq sbatch_gateq.sh

set -u
hostname
nvidia-smi -L
echo "[$(date)] ${CEQ_TAG:-w12_gateq} starting"

TAG="${CEQ_TAG:-w12_gateq}"
SCRATCH_ROOT=/scratch/users/gatmiry/llm-reasoning-logic-puzzles
RUN_DIR=${SCRATCH_ROOT}/sudoku-code/wavecurriculum_ceq
ENV_LOCAL=/tmp/logicpuzzles
CAND_DIR=/tmp/sudoku_s12
INST_DIR=/tmp/sudoku_superposition
LOCAL_LOG=/tmp/sudoku_wave_runs/${TAG}
WARM_DIR=/tmp/sudoku_warmstart_${TAG}
# The stage-4 checkpoint this run starts from, tarred up from the diagonal
# probe so the run begins at exactly the weights the gate was verified on.
S4_TAR=${RUN_DIR}/d4_ckpt.tgz
LOGIN_STAGE=gandalf.berkeley.edu:/tmp/sudoku_hf_${TAG}

mkdir -p "${RUN_DIR}/logs" "${LOCAL_LOG}" "${LOCAL_LOG}/_hf_tars" "${WARM_DIR}"
echo "[$(date)] /tmp free:"
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 XLA_PYTHON_CLIENT_PREALLOCATE=false ${ENV_LOCAL}/bin/python -u -c \
     "import jax; assert jax.default_backend()=='gpu'" 2>/dev/null; then
  echo "[$(date)] reusing ${ENV_LOCAL}"
else
  echo "ERROR: ${ENV_LOCAL} missing or not GPU-capable" >&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:-}

NDEV=$(XLA_PYTHON_CLIENT_PREALLOCATE=false ${PY} -u -c \
  "import jax; print(jax.local_device_count())")
export SUDOKU_MINIBATCH="${SUDOKU_MINIBATCH:-64}"
if [ -z "${NDEV}" ] || [ "${NDEV}" -lt 1 ]; then
  echo "ERROR: could not determine device count" >&2
  exit 1
fi
if [ $(( SUDOKU_MINIBATCH % NDEV )) -ne 0 ]; then
  echo "ERROR: minibatch ${SUDOKU_MINIBATCH} not divisible by ${NDEV} devices" >&2
  exit 1
fi
echo "[$(date)] ${NDEV} GPU(s), global batch ${SUDOKU_MINIBATCH}"

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
[ -s "${RUN_DIR}/test_bigq300.npz" ] || {
  echo "ERROR: missing ${RUN_DIR}/test_bigq300.npz (gen_bigq.py output)" >&2
  exit 1; }

# ---- Resume, in priority order ----------------------------------------------
#   1. this job's own node-local workdir (survives a requeue onto the same node)
#   2. the rolling latest.tar staged on the login node
#   3. the stage-4 checkpoint tarball (genuine cold start only)
CKPT_LOC=""
FRESH=1
if ls -1d "${LOCAL_LOG}"/checkpoint_* >/dev/null 2>&1; then
  CKPT_LOC="${LOCAL_LOG}"
  FRESH=0
  echo "[$(date)] resuming in place from step $(ls -1d ${LOCAL_LOG}/checkpoint_* | grep -v orbax-checkpoint-tmp | sed 's/.*checkpoint_//' | sort -n | tail -1)"
elif scp ${SCP_OPTS} "${LOGIN_STAGE}/latest.tar" "${WARM_DIR}/latest.tar" 2>/dev/null \
     && tar -C "${WARM_DIR}" -xf "${WARM_DIR}/latest.tar"; then
  rm -f "${WARM_DIR}/latest.tar"
  CKPT_LOC="${WARM_DIR}"
  FRESH=0
  echo "[$(date)] resuming from staged step $(ls -1d ${WARM_DIR}/checkpoint_* | sed 's/.*checkpoint_//' | sort -n | tail -1)"
elif [ "${CEQ_SCRATCH:-0}" = "1" ]; then
  # The data-curriculum baseline cannot inherit d4_ckpt.tgz: those weights were
  # trained with a depth ladder, which is the thing the baseline removes.
  CKPT_LOC="${WARM_DIR}"
  echo "[$(date)] random init (CEQ_SCRATCH=1)"
else
  [ -s "${S4_TAR}" ] || { echo "ERROR: no ${S4_TAR}" >&2; exit 1; }
  rm -rf "${WARM_DIR:?}"/checkpoint_*
  tar -C "${WARM_DIR}" -xzf "${S4_TAR}" || { echo "ERROR: bad ${S4_TAR}" >&2; exit 1; }
  CKPT_LOC="${WARM_DIR}"
  echo "[$(date)] cold start from stage-4 checkpoint $(ls -1d ${WARM_DIR}/checkpoint_* | sed 's/.*checkpoint_//' | sort -n | tail -1)"
fi
du -sh "${CKPT_LOC}" 2>/dev/null || true

export SUDOKU_RESUME=1
BASE_STAGE="${CEQ_START_STAGE:-4}"
CEQ_LOG_PROBE="${RUN_DIR}/logs/${TAG}.log"
if [ "${FRESH}" = "0" ] && [ -f "${CEQ_LOG_PROBE}" ]; then
  RECOVER_STAGE=$(grep -oE "promote to stage [0-9]+" "${CEQ_LOG_PROBE}" \
                  | grep -oE "[0-9]+$" | sort -n | tail -1)
  if [ -n "${RECOVER_STAGE}" ] && [ "${RECOVER_STAGE}" -gt "${BASE_STAGE}" ]; then
    echo "[$(date)] recovered curriculum stage ${RECOVER_STAGE} from log (base ${BASE_STAGE})"
    BASE_STAGE="${RECOVER_STAGE}"
  fi
fi
export SUDOKU_START_STAGE="${BASE_STAGE}"
export SUDOKU_MAX_STAGE=12
export SUDOKU_LATENT_SLOTS=12
export SUDOKU_RECURRENT=1
export SUDOKU_CAND_SLOT_MODE=depth
# num_passes = min(pps*stage, K). pps=1 is the latent curriculum: stage t reads
# out of slot t. pps=K pins every stage at full depth, so the stage only selects
# the data and there is a single readout -- the data-curriculum baseline.
export SUDOKU_PASSES_PER_STAGE="${CEQ_PPS:-1}"
export SUDOKU_AUX_WEIGHT=0.0
export SUDOKU_LEVEL_BALANCED=0
export SUDOKU_DATA_CURRICULUM=none
export SUDOKU_BACKTRACK="${CEQ_BACKTRACK:-0}"
export SUDOKU_BACKTRACK_MODE=adaptive
export SUDOKU_BACKTRACK_MARGIN=0.02

export SUDOKU_INSTANCE_EPOCHS=5
export SUDOKU_INSTANCE_PUZZLES="${CEQ_PUZZLES:-100000}"
export SUDOKU_MIN_STAGE_STEPS="${CEQ_MIN_STAGE_STEPS:-50000}"

# ---- The gate ----
export SUDOKU_PROMOTE_CEQ=0.90
export SUDOKU_PROMOTE_CEQ_MODE="${CEQ_MODE:-hinge}"
export SUDOKU_PROMOTE_CEQ_MARGIN=0.0
export SUDOKU_PROMOTE_HINGE="${CEQ_HINGE:-0.0}"
export SUDOKU_PROMOTE_HINGE_SLACK="${CEQ_HINGE_SLACK:-0.05}"
# Large-sample q, and the instance prefixes p is averaged over. These two are
# the whole point of this script; without them the gate reads 20-200x too high.
export SUDOKU_GATE_Q="${RUN_DIR}/test_bigq300.npz"
export SUDOKU_EVAL_READS="${CEQ_EVAL_READS:-4}"

export SUDOKU_PROMOTE_ACC=0.95
export SUDOKU_PROMOTE_SPREAD=0.88
export SUDOKU_PROMOTE_HBOUND=0.0
export SUDOKU_PROMOTE_LOC_WAVE=0.0
export SUDOKU_PLATEAU_STEPS=0
export SUDOKU_PATIENCE=0

export SUDOKU_MAX_STEPS="${CEQ_MAX_STEPS:-2000000}"
export SUDOKU_EVAL_EVERY="${CEQ_EVAL_EVERY:-2000}"
export SUDOKU_SAVE_EVERY=10000
export SUDOKU_CKPT_KEEP=2
export SUDOKU_LR="${CEQ_LR:-0.0002}"
export SUDOKU_DROPOUT="${CEQ_DROPOUT:-0.2}"
export SUDOKU_WD="${CEQ_WD:-0.005}"
export SUDOKU_LOC_WEIGHT="${CEQ_LOC_WEIGHT:-1.0}"
export SUDOKU_VALUE_WEIGHT="${CEQ_VALUE_WEIGHT:-1.0}"
export SUDOKU_LEAK_WEIGHT="${CEQ_LEAK_WEIGHT:-0.0}"
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="${CEQ_MEM_FRAC:-0.8}"

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

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
  exit 1
fi

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/${TAG} \
  --login-stage "${LOGIN_STAGE}" \
  --token-file "${HF_TOKEN_FILE}" \
  --interval 60 \
  > "${SYNC_LOG}" 2>&1 &
SYNC_PID=$!

echo "[$(date)] starting ${TAG}"
echo "  stages ${SUDOKU_START_STAGE}..12, ${NDEV} GPU(s), min ${SUDOKU_MIN_STAGE_STEPS} steps/stage"
echo "  gate: per-puzzle hinge V <= hinge_floor + ${SUDOKU_PROMOTE_HINGE_SLACK}"
echo "  gate q: ${SUDOKU_GATE_Q} (large-sample), ${SUDOKU_EVAL_READS} instance prefixes"
echo "  hf=Avra98/Sudoku_superposition/runs/${TAG}"
${PY} -u -m train.main \
  --workdir="${LOCAL_LOG}" --exp_name="${TAG}" \
  --ckpt_loc="${CKPT_LOC}" \
  >> "${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/${TAG} \
  --login-stage "${LOGIN_STAGE}" \
  --token-file "${HF_TOKEN_FILE}" \
  >> "${SYNC_LOG}" 2>&1 || true
echo "[$(date)] ${TAG} exit ${EC}"
exit ${EC}