#!/usr/bin/env bash
# Provisions a vast.ai instance and runs the full pipeline: generate -> verify ->
# train -> push to the Hub. Written to be re-runnable: a reclaimed instance can be
# replaced and this script started again from scratch.
#
# Expects, in the environment:
#   HF_TOKEN     write token, for checkpoint pushes
#   HUB_REPO     e.g. yourname/tiny-cube-solver
# Optional: PAIRS (default 10000000), STEPS, HIDDEN, LAYERS
set -euo pipefail

PAIRS="${PAIRS:-10000000}"
STEPS="${STEPS:-40000}"
HIDDEN="${HIDDEN:-384}"
LAYERS="${LAYERS:-8}"
CORES="$(nproc)"

echo "=== $CORES cores, $(nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null || echo 'no GPU') ==="

# Build against SYSTEM python, not the image's conda python. kociemba ships a C
# extension and a pure-Python fallback, and it falls back *silently* when the
# extension cannot build -- which is what happens if the venv's interpreter has no
# matching dev headers. python3-dev is therefore not optional here.
apt-get update -qq
# libffi-dev is the one that actually matters and is the least obvious: kociemba's
# native extension is built through cffi, which needs ffi.h. Without it the
# compile dies with "fatal error: ffi.h: No such file or directory", kociemba
# installs its pure-Python fallback instead, and pip still exits 0.
apt-get install -y -qq build-essential python3-dev python3-venv libffi-dev git
# Prefer python3.11. kociemba 1.2.1 builds its C extension cleanly on 3.11 (the
# version this repo is developed against) but silently produces no extension on
# the 3.10 that ships as system python in the Ubuntu 22.04 based images -- pip
# still reports success, so the only symptom is the fallback. The PyTorch images
# carry a conda 3.11, so prefer that and fall back to system python.
PYBIN=""
for cand in /opt/conda/bin/python3.11 /opt/conda/bin/python "$(command -v python3.11 || true)" /usr/bin/python3; do
  [ -x "$cand" ] || continue
  if "$cand" -c 'import sys; sys.exit(0 if sys.version_info[:2] >= (3,11) else 1)' 2>/dev/null; then
    PYBIN="$cand"; break
  fi
done
[ -z "$PYBIN" ] && PYBIN=/usr/bin/python3
echo "building venv from $PYBIN ($($PYBIN -V 2>&1))"
"$PYBIN" -m venv /opt/venv
/opt/venv/bin/pip install -q --upgrade pip
# kociemba's setup.py predates modern setuptools' removal of install_layout.
/opt/venv/bin/pip install -q "setuptools<60" wheel
# Full output, not a tail: when the C extension fails to build, kociemba still
# installs successfully (as pure Python), so the compiler error is the only
# evidence of what went wrong and it must not be truncated away.
/opt/venv/bin/pip install --no-build-isolation kociemba 2>&1 | tail -40

# Last resort before giving up: build from the upstream git source, whose build
# metadata is newer than the 2017 PyPI sdist.
if ! /opt/venv/bin/python -c "import kociemba.ckociembawrapper" 2>/dev/null; then
  echo "=== sdist build produced no extension; trying upstream git ==="
  /opt/venv/bin/pip install --force-reinstall --no-build-isolation \
    "git+https://github.com/muodov/kociemba.git" 2>&1 | tail -15 || true
fi

if ! /opt/venv/bin/python -c "import kociemba.ckociembawrapper" 2>/dev/null; then
  echo "=== NATIVE BUILD FAILED -- diagnostics ==="
  echo "python: $(/opt/venv/bin/python -V 2>&1) at $(/opt/venv/bin/python -c 'import sys;print(sys.executable)')"
  echo "gcc: $(gcc --version 2>&1 | head -1)"
  echo "includes: $(/opt/venv/bin/python -c 'import sysconfig;print(sysconfig.get_paths()["include"])')"
  ls -l "$(/opt/venv/bin/python -c 'import sysconfig;print(sysconfig.get_paths()["include"])')/Python.h" 2>&1 || echo "  Python.h MISSING"
  dpkg -l | grep -E "python3.*dev|build-essential" || echo "  no dev packages found"
  echo "=== retrying verbosely ==="
  /opt/venv/bin/pip install --force-reinstall --no-build-isolation -v kociemba 2>&1 | grep -iE "error|gcc|Python.h|fatal" | head -30
fi
/opt/venv/bin/pip install -q torch --index-url https://download.pytorch.org/whl/cu128
/opt/venv/bin/pip install -q transformers huggingface_hub

# Search evaluation only: skip data generation and training entirely, download a
# trained value model and score the state search on the canonical scrambles. The
# search needs a GPU -- each solve evaluates beam*18 children per step -- so it
# cannot be done from the session container.
if [ -n "${EVAL_VALUE_ONLY:-}" ]; then
  cd "${CODE_DIR:-/workspace/code}"
  echo "=== evaluating value search from ${EVAL_VALUE_ONLY} ==="
  /opt/venv/bin/python eval_value_search.py --value "$EVAL_VALUE_ONLY" \
    --beam "${SEARCH_BEAM:-64}" --limit-per-depth "${SEARCH_N:-25}" \
    ${SEARCH_DEPTHS:+--depths "$SEARCH_DEPTHS"} ${SEARCH_LAM:+--lam "$SEARCH_LAM"}
  echo "=== done ==="
  exit 0
fi

# Hard gate. The pure-Python fallback is ~50x slower, so generating this dataset
# with it would take days rather than hours -- and the only symptom is a warning
# buried in the log while the meter runs. Fail here instead.
echo "=== verifying native kociemba ==="
/opt/venv/bin/python - <<'PYCHECK'
import sys, time
import kociemba
from kociemba import ckociembawrapper  # absent entirely when the C build failed
SOLVED = "".join("URFDLB"[i // 9] for i in range(54))
kociemba.solve("DRLUUBFBRBLURRLRUBLRDDFDLFUFUFFDBRDUBRUFLLFDDBFLUBLRBD")
t = time.time()
for _ in range(20):
    kociemba.solve("DRLUUBFBRBLURRLRUBLRDDFDLFUFUFFDBRDUBRUFLLFDDBFLUBLRBD")
ms = (time.time() - t) / 20 * 1000
print(f"kociemba: {ms:.1f} ms/solve (native)")
if ms > 200:
    sys.exit(f"ABORT: {ms:.0f} ms/solve is fallback-speed; refusing to generate")
PYCHECK

cd "${CODE_DIR:-/workspace/code}"
mkdir -p data checkpoints

# Smoke-test the GPU *before* generating anything. A torch build without kernels
# for this GPU fails only when the first tensor op runs, which used to be training
# -- so an unusable GPU was discovered an hour and a full dataset later. This
# runs the real model for one forward and backward pass, so it catches missing
# kernels rather than merely that a device is visible.
echo "=== verifying GPU can run the model ==="
/opt/venv/bin/python - <<PYCHECK
import sys, torch
if not torch.cuda.is_available():
    sys.exit("ABORT: no CUDA device visible")
name = torch.cuda.get_device_name(0)
cap = torch.cuda.get_device_capability(0)
vram = torch.cuda.get_device_properties(0).total_memory / 1e9
print(f"gpu: {name} sm_{cap[0]}{cap[1]} {vram:.0f}GB | torch {torch.__version__} | arch {torch.cuda.get_arch_list()}")
try:
    from transformers import LlamaConfig, LlamaForCausalLM
    # The REAL config and batch size. A scaled-down probe is worse than useless
    # here: it passes on a GPU that cannot hold the actual model, so the OOM
    # surfaces only after data generation has already been paid for.
    cfg = LlamaConfig(vocab_size=28, hidden_size=$HIDDEN, intermediate_size=$HIDDEN * 4,
                      num_hidden_layers=$LAYERS, num_attention_heads=8,
                      num_key_value_heads=8, max_position_embeddings=83,
                      tie_word_embeddings=True)
    m = LlamaForCausalLM(cfg).cuda()
    opt = torch.optim.AdamW(m.parameters(), lr=1e-4)
    ids = torch.randint(0, 28, ($BATCH, 83), device="cuda")
    with torch.autocast("cuda", dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16):
        loss = m(input_ids=ids, labels=ids).loss
    loss.backward()
    opt.step()   # optimizer state is where the memory actually peaks
    torch.cuda.synchronize()
    peak = torch.cuda.max_memory_allocated() / 1e9
    print(f"gpu smoke test: full step at batch $BATCH OK, peak {peak:.1f}GB of {vram:.0f}GB")
except torch.OutOfMemoryError:
    sys.exit(f"ABORT: batch $BATCH of a $LAYERS-layer/$HIDDEN-wide model does not fit in {vram:.0f}GB")
except Exception as e:
    sys.exit(f"ABORT: GPU cannot run the model ({type(e).__name__}: {e})")
PYCHECK

cd "${CODE_DIR:-/workspace/code}"
mkdir -p data checkpoints

# Generation is the wall-clock bottleneck, not training -- use every core.
echo "=== generating $PAIRS pairs on $CORES cores ==="
# --workers 0 detects the cgroup quota. --augment takes several states from each
# solution path, so one Kociemba solve yields several pairs: an ~8x cut in solves,
# which is the actual bottleneck. It also samples distance-to-solved far more
# evenly than uniform states do (those are ~95% distance-18-to-20), which happens
# to match a benchmark whose buckets run from depth 3 upward.
/opt/venv/bin/python gen_data.py --count "$PAIRS" --workers 0 --augment 8 --out data/train.jsonl

# A separate seed, and canonical labels only: the holdout must reflect the real
# task, not the augmented training mix, or the reported solve rate is inflated.
echo "=== generating canonical holdout ==="
# Holdout stays unaugmented and uniform: the hard case. Scoring against the
# augmented mix would flatter the model, since those states sit partway along a
# solution path and are nearer to solved.
# --shards 32 fixes the shard count so the holdout is byte-identical on every
# machine. Without it the per-shard seeds follow the core count, so two boxes
# comparing two models evaluate them against merely similar holdouts rather than
# the same one -- silently downgrading a paired comparison to an unpaired one,
# exactly where the pairing matters most.
/opt/venv/bin/python gen_data.py --count 20000 --workers 0 --shards 32 --seed 999999 \
  --augment 0 --out data/val.jsonl

# DAgger phase: if ROLLOUT_FROM names a checkpoint, roll that policy out, label
# the states it actually reaches, and append them to the training set.
#
# The model is otherwise trained only on states from the sampler and never on the
# ones its own attempts produce -- which is where it measurably fails, walking
# from 10 moves out to 18 and staying there. Behaviour cloning alone compounds
# error quadratically in horizon; training on the policy's own state
# distribution is the standard fix.
if [ -n "${ROLLOUT_FROM:-}" ]; then
  echo "=== rolling out ${ROLLOUT_FROM} for ${ROLLOUT_PAIRS:-2000000} states ==="
  /opt/venv/bin/python gen_rollout_data.py --checkpoint "$ROLLOUT_FROM" \
    --count "${ROLLOUT_PAIRS:-2000000}" --batch-size "${ROLLOUT_BATCH:-2048}" --rounds 4 \
    --out data/rollout.jsonl
  /opt/venv/bin/python verify_data.py data/rollout.jsonl --limit 20000
  cat data/rollout.jsonl >> data/train.jsonl
  echo "=== training set is now $(wc -l < data/train.jsonl) pairs ==="
fi

echo "=== verifying (must report 0 FAILED) ==="
/opt/venv/bin/python verify_data.py data/train.jsonl --limit 100000
/opt/venv/bin/python verify_data.py data/val.jsonl

# TRAIN_VALUE switches the objective from "emit a solution" to "estimate distance
# to solved". The policy shape has been shown four times over not to reach past
# ~8 moves; this is the heuristic that lets search steer. Same data either way --
# the label is just the solution length instead of the solution.

if [ -n "${TRAIN_VALUE:-}" ]; then
  echo "=== training value function ==="
  /opt/venv/bin/python train_value.py \
    --data data/train.jsonl --val-data data/val.jsonl \
    --hidden "$HIDDEN" --layers "$LAYERS" --heads 8 \
    --batch-size "${BATCH:-1024}" --lr 6e-4 --max-steps "$STEPS" \
    --eval-every 1000 --eval-n 8192 \
    --out checkpoints/value --hub-repo "${HUB_REPO:-}"
  echo "=== done ==="
  exit 0
fi

echo "=== training ==="
/opt/venv/bin/python train.py \
  --data data/train.jsonl --val-data data/val.jsonl \
  --hidden "$HIDDEN" --layers "$LAYERS" --heads 8 \
  --batch-size "${BATCH:-1024}" --lr 6e-4 --max-steps "$STEPS" \
  --eval-every 1000 --eval-n 512 \
  --out checkpoints/cube --hub-repo "${HUB_REPO:-}" ${RESUME:+--resume "$RESUME"}

echo "=== done ==="