#!/usr/bin/env bash
# Purpose: prepare a from-scratch pretraining run: build the nanochat virtual
#          environment for this platform track, download a chosen number of corpus
#          shards, train a byte-pair tokeniser on them, and evaluate its compression
#          against the GPT-2 and GPT-4 vocabularies
# Platform: all (CUDA on Tracks S and N, ROCm on Track X, MPS on Track M, CPU anywhere)
# Minimum memory: 8 GB
# Assumes: nanochat has been cloned to $NANOCHAT (default ~/nanochat); uv, git and
#          curl are on PATH; there is about 150 MB of free disk per corpus shard
#          under $NANOCHAT_BASE_DIR (default ~/.cache/nanochat); the machine can
#          reach pypi.org, download.pytorch.org and huggingface.co
set -euo pipefail

# ---------------------------------------------------------------- settings ---
TRACK="${TRACK:-}"                 # spark | strix | mac | nvidia | cpu
NANOCHAT="${NANOCHAT:-$HOME/nanochat}"
SHARDS="${SHARDS:-8}"              # corpus shards to download; ~250M characters each
VOCAB_SIZE="${VOCAB_SIZE:-32768}"  # nanochat's default, 2^15
MAX_CHARS="${MAX_CHARS:-2000000000}"   # characters the tokeniser is trained on
ROCM_INDEX="${ROCM_INDEX:-https://download.pytorch.org/whl/rocm6.4}"
# Optional: replace torch after the sync with a build from another wheel index.
# Track X always does this because nanochat's pyproject.toml knows only a CPU index
# and a CUDA index. Any track can do it, which is how Track S moves to a different
# CUDA build without editing the project's files.
TORCH_INDEX="${TORCH_INDEX:-}"
TORCH_VERSION="${TORCH_VERSION:-2.9.1}"   # the version nanochat pins
LOG_DIR="${LOG_DIR:-$PWD/part-12-logs}"

export NANOCHAT_BASE_DIR="${NANOCHAT_BASE_DIR:-$HOME/.cache/nanochat}"

usage() {
    cat <<'USAGE'
Usage: TRACK=<spark|strix|mac|nvidia|cpu> bash prepare-data.sh

Environment variables (all optional except TRACK):
  NANOCHAT           path to the cloned repository        (default ~/nanochat)
  SHARDS             corpus shards to download            (default 8)
  VOCAB_SIZE         tokeniser vocabulary size            (default 32768)
  MAX_CHARS          characters to train the tokeniser on (default 2000000000)
  NANOCHAT_BASE_DIR  where shards and checkpoints live    (default ~/.cache/nanochat)
  ROCM_INDEX         PyTorch ROCm wheel index, Track X   (default rocm6.4)
  TORCH_INDEX        replace torch from this wheel index after the sync
  TORCH_VERSION      version to pull from TORCH_INDEX     (default 2.9.1)
  LOG_DIR            where this script writes its logs    (default ./part-12-logs)
USAGE
}

case "$TRACK" in
    spark|strix|mac|nvidia|cpu) ;;
    *) usage; echo; echo "ERROR: set TRACK to one of spark, strix, mac, nvidia, cpu." >&2; exit 2 ;;
esac

# ------------------------------------------------------------- preflight ----
need() {
    command -v "$1" >/dev/null 2>&1 || { echo "ERROR: '$1' is not on PATH. $2" >&2; exit 1; }
}
need uv "Install it from https://docs.astral.sh/uv/ as in Part 1."
need git "Install your distribution's git package."

if [ ! -d "$NANOCHAT/.git" ]; then
    echo "ERROR: no git repository at $NANOCHAT." >&2
    echo "       git clone https://github.com/karpathy/nanochat.git \"$NANOCHAT\"" >&2
    exit 1
fi

mkdir -p "$LOG_DIR" "$NANOCHAT_BASE_DIR"
cd "$NANOCHAT"

COMMIT="$(git rev-parse --short HEAD)"
echo "==> nanochat at $NANOCHAT, commit $COMMIT"
echo "==> artefacts in $NANOCHAT_BASE_DIR, logs in $LOG_DIR"

# --------------------------------------------------------- the environment --
# nanochat pins torch in pyproject.toml and selects the wheel index with an extra:
# "gpu" points at the CUDA index and "cpu" at the CPU index, which is also the one
# that carries the Metal-capable macOS wheels. Track X installs the CPU set first
# and then replaces torch with a ROCm build from AMD's index.
echo "==> creating the virtual environment"
[ -d .venv ] || uv venv

case "$TRACK" in
    spark|nvidia) uv sync --extra gpu ;;
    mac|cpu)      uv sync --extra cpu ;;
    strix)
        # The CPU extra brings every dependency except an accelerated torch.
        uv sync --extra cpu
        TORCH_INDEX="${TORCH_INDEX:-$ROCM_INDEX}"
        ;;
esac

if [ -n "$TORCH_INDEX" ]; then
    echo "==> replacing torch $TORCH_VERSION from $TORCH_INDEX"
    # --index-url makes that index authoritative for torch itself and the wheels it
    # ships alongside; --extra-index-url keeps PyPI available for the ordinary
    # dependencies (filelock, sympy, networkx and friends).
    uv pip install \
        --index-url "$TORCH_INDEX" \
        --extra-index-url https://pypi.org/simple \
        "torch==$TORCH_VERSION"
fi

# shellcheck disable=SC1091
source .venv/bin/activate

# ------------------------------------------------------------ verification --
# Print what we actually got. A wrong wheel here is the single most common reason
# a run is mysteriously slow, and it is silent unless you look.
echo "==> environment"
python - <<'PY'
import platform
import torch

print(f"    python  {platform.python_version()} on {platform.system()} {platform.machine()}")
print(f"    torch   {torch.__version__}")
if torch.cuda.is_available():
    print(f"    device  cuda -> {torch.cuda.get_device_name(0)}")
    print(f"    hip     {getattr(torch.version, 'hip', None)}")
    print(f"    cuda    {getattr(torch.version, 'cuda', None)}")
elif getattr(torch.backends, "mps", None) is not None and torch.backends.mps.is_available():
    print("    device  mps -> Apple silicon GPU")
else:
    print("    device  cpu -> no accelerator visible; the run will be slow")
PY

# ------------------------------------------------------------------ corpus --
# Shards are parquet files with a "text" column. The download always adds the
# final shard as well, which nanochat holds out as the validation split.
echo "==> downloading $SHARDS corpus shard(s) plus the validation shard"
python -m nanochat.dataset -n "$SHARDS" 2>&1 | tee "$LOG_DIR/dataset.log"

SHARD_DIR="$NANOCHAT_BASE_DIR/base_data_climbmix"
if [ -d "$SHARD_DIR" ]; then
    echo "==> on disk: $(du -sh "$SHARD_DIR" | cut -f1) in $(find "$SHARD_DIR" -name '*.parquet' | wc -l) file(s)"
fi

# --------------------------------------------------------------- tokeniser --
echo "==> training the tokeniser: vocabulary $VOCAB_SIZE, from $MAX_CHARS characters"
python -m scripts.tok_train \
    --max-chars="$MAX_CHARS" \
    --vocab-size="$VOCAB_SIZE" 2>&1 | tee "$LOG_DIR/tok-train.log"

echo "==> evaluating the tokeniser against the GPT-2 and GPT-4 vocabularies"
python -m scripts.tok_eval 2>&1 | tee "$LOG_DIR/tok-eval.log"

# ----------------------------------------------------------------- summary --
cat <<SUMMARY

==> done
    commit         $COMMIT
    track          $TRACK
    shards         $SHARDS (+1 validation shard)
    vocabulary     $VOCAB_SIZE
    tokeniser      $NANOCHAT_BASE_DIR/tokenizer
    logs           $LOG_DIR

    Record in the lab notebook: the commit, the torch version and device the
    verification step printed, the shard count, and the bytes-per-token ratio
    your tokeniser achieved on the corpus rows of the tok_eval table, next to
    what GPT-2 and GPT-4 achieved on the same text.

    Next: TRACK=$TRACK bash train-small.sh
SUMMARY