#!/bin/bash
# Score saved checkpoints on the CE(q||p) <= log|S| promotion criterion.
#
# Exports the same SUDOKU_* configuration the training jobs used, so the model
# shape and the eval split match the checkpoint. Paths point at the local copies
# of the datasets rather than the cluster's /tmp staging.
#
# Usage: run_ceq_eval.sh <ckpt_dir> <eval_stage> <label> [gpu] [eval_epochs]
set -euo pipefail

CKPT_DIR="$1"
STAGE="$2"
LABEL="$3"
GPU="${4:-3}"
EPOCHS="${5:-40}"

ROOT=/egr/research-slim/ghoshavr/llm-reasoning-logic-puzzles/sudoku-code
PY=/egr/research-slim/ghoshavr/conda-envs/logicpuzzles/bin/python

# Model / curriculum shape: must match the checkpoint exactly.
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_MAX_STAGE=12
export SUDOKU_START_STAGE="${STAGE}"
export SUDOKU_MINIBATCH=64
export SUDOKU_LR=0.0002
export SUDOKU_DROPOUT=0.2
export SUDOKU_WD=0.005
export SUDOKU_RESUME=0

# Data: local copies. The instance pool supplies q; the candidate masks supply S.
export SUDOKU_TRAIN_PATH="${ROOT}/datasets/train_sudoku_puzzles.npy"
export SUDOKU_TEST_PATH="${ROOT}/datasets/test_sudoku_puzzles.npy"
export SUDOKU_TRAIN_CAND="${ROOT}/datasets_multicandidate_s12/train_cand_masks.npy"
export SUDOKU_TEST_CAND="${ROOT}/datasets_multicandidate_s12/test_cand_masks.npy"
export SUDOKU_INSTANCE_DIR="${ROOT}/datasets_superposition"
# Eval reads no schedule, but instance_epochs is what puts the run in
# instance mode, which is what makes the value metrics defined.
export SUDOKU_INSTANCE_EPOCHS=5
export SUDOKU_INSTANCE_PUZZLES=100000

export SUDOKU_VERBOSE_EVAL=0
export CUDA_VISIBLE_DEVICES="${GPU}"
export XLA_PYTHON_CLIENT_PREALLOCATE=false
export XLA_PYTHON_CLIENT_MEM_FRACTION=.25
export TF_CPP_MIN_LOG_LEVEL=3
export HF_HOME=/tmp/hfcache

# train.* is imported as a top-level package, matching how the jobs run it.
cd "${ROOT}/wavecurriculum_run"
exec ${PY} -u eval_ceq_ckpt.py \
  --ckpt_dir "${CKPT_DIR}" --stage "${STAGE}" --label "${LABEL}" \
  --eval_epochs "${EPOCHS}" \
  --out_json "/tmp/ceq_results/${LABEL}.json"