#!/usr/bin/env bash
# 12-stage latent curriculum, instance-wise superposition targets.
#
# Same architecture and curriculum as w12_latent, except:
#   - the output prompt is one sampled stage-k assignment, not the unique
#     solution (input clues stay fixed)
#   - the candidate-head BCE is off (SUDOKU_AUX_WEIGHT=0)
#   - promotion is gated on the in-set rate (emitted digit is a stage-k
#     candidate), not on candidate-set exact match
#
# Usage:
#   bash launch_instance_latent.sh            # GPU 0, full 800k run
#   GPU=6 bash launch_instance_latent.sh      # pick another GPU
#   SUDOKU_MAX_STEPS=200 bash launch_instance_latent.sh   # smoke

set -u

_SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
_CODE_DIR="$(cd "${_SCRIPT_DIR}/.." && pwd)"
export SUDOKU_RUN_DIR="${SUDOKU_RUN_DIR:-${_SCRIPT_DIR}}"
# shellcheck source=env_paths.sh
source "${_SCRIPT_DIR}/env_paths.sh"

cd "$RUN_DIR" || exit 1
if [ -n "${CUDNN_LIB}" ]; then export LD_LIBRARY_PATH="${CUDNN_LIB}:${LD_LIBRARY_PATH:-}"; fi
export XLA_PYTHON_CLIENT_MEM_FRACTION=0.9

GPU="${GPU:-0}"
name="${NAME:-w12_inst_latent}"
mkdir -p logs

export SUDOKU_TRAIN_PATH="${_CODE_DIR}/datasets/train_sudoku_puzzles.npy"
export SUDOKU_TEST_PATH="${_CODE_DIR}/datasets/test_sudoku_puzzles.npy"
# Masks are a metric only (in-set rate / promotion). Not a training target.
export SUDOKU_TRAIN_CAND="${_CODE_DIR}/datasets_multicandidate_s12/train_cand_masks.npy"
export SUDOKU_TEST_CAND="${_CODE_DIR}/datasets_multicandidate_s12/test_cand_masks.npy"
export SUDOKU_INSTANCE_DIR="${_CODE_DIR}/datasets_superposition"

export SUDOKU_RESUME=0
export SUDOKU_LATENT_SLOTS=12
export SUDOKU_RECURRENT=1
export SUDOKU_BACKTRACK=0
export SUDOKU_START_STAGE=1
export SUDOKU_MAX_STAGE=12
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="${SUDOKU_PLATEAU_STEPS:-20000}"
export SUDOKU_PLATEAU_DELTA=0.005
export SUDOKU_PATIENCE="${SUDOKU_PATIENCE:-80000}"
export SUDOKU_MIN_STAGE_STEPS="${SUDOKU_MIN_STAGE_STEPS:-8000}"
export SUDOKU_PROMOTE_ACC="${SUDOKU_PROMOTE_ACC:-0.90}"
export SUDOKU_MAX_STEPS="${SUDOKU_MAX_STEPS:-800000}"
export SUDOKU_EVAL_EVERY="${SUDOKU_EVAL_EVERY:-2000}"
export SUDOKU_SAVE_EVERY="${SUDOKU_SAVE_EVERY:-10000}"
export SUDOKU_CKPT_KEEP=100

export SUDOKU_LR=0.0002
export SUDOKU_DROPOUT=0.2
export SUDOKU_WD=0.005

echo "launching ${name} on GPU ${GPU}"
echo "  K=12 recurrent=1 aux=0 instance_dir=${SUDOKU_INSTANCE_DIR}"
echo "  max_steps=${SUDOKU_MAX_STEPS} plateau=${SUDOKU_PLATEAU_STEPS} patience=${SUDOKU_PATIENCE}"
CUDA_VISIBLE_DEVICES="${GPU}" \
  nohup "$PY" -u -m train.main \
    --workdir="./logs/${name}" \
    --exp_name="${name}" \
    > "logs/${name}.log" 2>&1 &
echo "  pid $!  log=${RUN_DIR}/logs/${name}.log"