#!/bin/bash
# Noisy-latent Sudoku trainer. Works from either layout:
#   Sudoku_superposition/          (HF: code/ + data/)
#   sudoku-wavecurriculum/         (GitHub branch; set SUDOKU_DATA_DIR)
#
# Usage:
#   bash run_experiment.sh
#   SUDOKU_INSTANCE_DIR=./datasets_superposition bash run_experiment.sh
#   SUDOKU_DATA_DIR=/path/to/data bash run_experiment.sh
set -euo pipefail

die() { echo "ERROR: $*" >&2; exit 1; }
need_file() { [ -s "$1" ] || die "missing or empty: $1"; }

HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

if [ -d "${HERE}/train" ] && [ -f "${HERE}/train/main.py" ]; then
  CODE="${HERE}"
  ROOT="$(cd "${HERE}/.." && pwd)"
elif [ -d "${HERE}/code/train" ] && [ -f "${HERE}/code/train/main.py" ]; then
  CODE="${HERE}/code"
  ROOT="${HERE}"
elif [ -d "${HERE}/../data" ] && [ -f "${HERE}/train/main.py" ]; then
  CODE="${HERE}"
  ROOT="$(cd "${HERE}/.." && pwd)"
else
  die "cannot find train/main.py next to this script or in ./code/"
fi

# HF download puts arrays in data/. New instance rebuilds go in datasets_superposition/.
DATA="${SUDOKU_DATA_DIR:-}"
if [ -z "${DATA}" ]; then
  if [ -d "${ROOT}/data" ]; then
    DATA="${ROOT}/data"
  elif [ -d "${HERE}/data" ]; then
    DATA="${HERE}/data"
  else
    die "set SUDOKU_DATA_DIR to the folder with the 11 npy/npz files"
  fi
fi
DATA="$(cd "${DATA}" && pwd)"

INST="${SUDOKU_INSTANCE_DIR:-}"
if [ -z "${INST}" ]; then
  if [ -s "${ROOT}/datasets_superposition/train_assignments.npy" ] \
     && [ -s "${ROOT}/datasets_superposition/train_starts.npy" ] \
     && [ -s "${ROOT}/datasets_superposition/train_counts.npy" ]; then
    INST="${ROOT}/datasets_superposition"
    echo "[run] using rebuilt instances in ${INST}"
  else
    INST="${DATA}"
  fi
fi
INST="$(cd "${INST}" && pwd)"

export SUDOKU_TRAIN_PATH="${SUDOKU_TRAIN_PATH:-${DATA}/train_sudoku_puzzles.npy}"
export SUDOKU_TEST_PATH="${SUDOKU_TEST_PATH:-${DATA}/test_sudoku_puzzles.npy}"
export SUDOKU_TRAIN_CAND="${SUDOKU_TRAIN_CAND:-${DATA}/train_cand_masks.npy}"
export SUDOKU_TEST_CAND="${SUDOKU_TEST_CAND:-${DATA}/test_cand_masks.npy}"
export SUDOKU_INSTANCE_DIR="${INST}"
export SUDOKU_GATE_Q="${SUDOKU_GATE_Q:-${DATA}/test_bigq300.npz}"

need_file "${SUDOKU_TRAIN_PATH}"
need_file "${SUDOKU_TEST_PATH}"
need_file "${SUDOKU_TRAIN_CAND}"
need_file "${SUDOKU_TEST_CAND}"
need_file "${SUDOKU_GATE_Q}"
for split in train test; do
  need_file "${INST}/${split}_assignments.npy"
  need_file "${INST}/${split}_starts.npy"
  need_file "${INST}/${split}_counts.npy"
done

export SUDOKU_LATENT_SLOTS="${SUDOKU_LATENT_SLOTS:-12}"
export SUDOKU_AUX_WEIGHT="${SUDOKU_AUX_WEIGHT:-0}"
export SUDOKU_LEAK_WEIGHT="${SUDOKU_LEAK_WEIGHT:-0}"
export SUDOKU_LATENT_NOISE="${SUDOKU_LATENT_NOISE:-1}"
export SUDOKU_MC_SAMPLES="${SUDOKU_MC_SAMPLES:-8}"
export SUDOKU_MC_OBJ="${SUDOKU_MC_OBJ:-mean}"
export SUDOKU_NOISE_LOGSIGMA0="${SUDOKU_NOISE_LOGSIGMA0:--2.5}"
export SUDOKU_START_STAGE="${SUDOKU_START_STAGE:-1}"
export SUDOKU_MAX_STAGE="${SUDOKU_MAX_STAGE:-12}"
export SUDOKU_MINIBATCH="${SUDOKU_MINIBATCH:-64}"
export SUDOKU_EVAL_EVERY="${SUDOKU_EVAL_EVERY:-2000}"
export SUDOKU_MAX_STEPS="${SUDOKU_MAX_STEPS:-2500000}"

PY="${PYTHON:-}"
if [ -z "${PY}" ]; then
  if command -v python >/dev/null 2>&1; then
    PY=python
  elif command -v python3 >/dev/null 2>&1; then
    PY=python3
  else
    die "python not on PATH"
  fi
fi

# Trainer imports these at module load even when wandb is off.
"${PY}" -c "import jax, flax, optax, ml_collections, tensorflow, wandb, absl" \
  || die "install: jax flax optax ml_collections clu absl-py tensorflow wandb"

WORKDIR="${WORKDIR:-${ROOT}/runs/noise}"
EXP_NAME="${SUDOKU_EXP_NAME:-noise}"
mkdir -p "${WORKDIR}"

echo "[run] code=${CODE}"
echo "[run] puzzles/masks/q=${DATA}"
echo "[run] instances=${INST}"
echo "[run] workdir=${WORKDIR}"
echo "[run] noise=${SUDOKU_LATENT_NOISE} M=${SUDOKU_MC_SAMPLES} K=${SUDOKU_LATENT_SLOTS} start=${SUDOKU_START_STAGE}"

cd "${CODE}"
[ -f train/main.py ] || die "train/main.py not in ${CODE}"
exec "${PY}" -m train.main --workdir="${WORKDIR}" --exp_name="${EXP_NAME}"