#!/usr/bin/env bash
set -euo pipefail

PROFILE="${1:-today}"
WORK_ROOT="${NAYARA_WORK_ROOT:-/teamspace/studios/this_studio/nayara-work}"

case "$PROFILE" in
  preview)
    MAX_TOKENS=50000000
    MAX_MINUTES=30
    DATA_DIR="$WORK_ROOT/packed-preview"
    TOKENIZER="$WORK_ROOT/tokenizer-preview.json"
    ;;
  today)
    MAX_TOKENS=300000000
    MAX_MINUTES=90
    DATA_DIR="$WORK_ROOT/packed-today"
    TOKENIZER="$WORK_ROOT/tokenizer-today.json"
    ;;
  overnight)
    MAX_TOKENS=2000000000
    MAX_MINUTES=540
    DATA_DIR="$WORK_ROOT/packed-overnight"
    TOKENIZER="$WORK_ROOT/tokenizer-overnight.json"
    ;;
  sprint)
    MAX_TOKENS="${NAYARA_SPRINT_TOKENS:-2000000000}"
    MAX_MINUTES="${NAYARA_SPRINT_MINUTES:-205}"
    DATA_DIR="${NAYARA_PACKED_DATA:?Set NAYARA_PACKED_DATA to a pre-tokenized elite corpus}"
    TOKENIZER="${NAYARA_TOKENIZER:?Set NAYARA_TOKENIZER to its tokenizer.json}"
    ;;
  *)
    echo "usage: bash scripts/lightning_train_h100.sh [preview|today|overnight|sprint]" >&2
    exit 2
    ;;
esac

MAX_TOKENS="${NAYARA_MAX_TOKENS:-$MAX_TOKENS}"
MAX_MINUTES="${NAYARA_MAX_MINUTES:-$MAX_MINUTES}"
DATA_DIR="${NAYARA_PACKED_DATA:-$DATA_DIR}"
TOKENIZER="${NAYARA_TOKENIZER:-$TOKENIZER}"

# Never silently cycle a small corpus. Clamp the schedule just below the
# largest no-repeat weighted mixture supported by the packed manifest.
REQUESTED_TOKENS="$MAX_TOKENS"
MAX_TOKENS="$(python - "$DATA_DIR/manifest.json" "$REQUESTED_TOKENS" "${NAYARA_RESUME:-}" <<'PY'
import json
import sys
from pathlib import Path

manifest = json.loads(Path(sys.argv[1]).read_text())
requested = int(sys.argv[2])
resume_tokens = 0
if sys.argv[3]:
    resume_tokens = int(json.loads((Path(sys.argv[3]) / "progress.json").read_text())["tokens"])
capacity = min(
    int((item["tokens"] - resume_tokens * float(item["weight"])) / float(item["weight"]))
    for item in manifest["sources"].values()
    if float(item["weight"]) > 0
)
safe = max(1, int(capacity * 0.98))
print(min(requested, safe))
PY
)"
if [[ "$MAX_TOKENS" -lt "$REQUESTED_TOKENS" ]]; then
  echo "Packed corpus supports $MAX_TOKENS unique mixed tokens; clamped from $REQUESTED_TOKENS."
fi

python - <<'PY'
import torch
assert torch.cuda.is_available(), "Switch the Studio to the H100 machine"
assert torch.cuda.is_bf16_supported(), "BF16 is unavailable"
name = torch.cuda.get_device_name(0)
memory = torch.cuda.get_device_properties(0).total_memory / 2**30
print(f"GPU preflight: {name}, {memory:.1f} GiB, torch={torch.__version__}, CUDA={torch.version.cuda}")
assert memory >= 70, "This profile expects an 80GB H100"
PY

OUTPUT="$WORK_ROOT/checkpoints-${PROFILE}"
MICRO_BATCH="${NAYARA_MICRO_BATCH:-64}"
EXTRA_ARGS=()
if [[ -n "${NAYARA_RESUME:-}" ]]; then
  EXTRA_ARGS+=(--resume "$NAYARA_RESUME")
fi
if [[ -n "${NAYARA_PEAK_LR:-}" ]]; then
  EXTRA_ARGS+=(--peak-lr "$NAYARA_PEAK_LR")
fi
python -m nayara.train \
  --strategy single \
  --config configs/nayara_350m.json \
  --data "$DATA_DIR" \
  --output "$OUTPUT" \
  --micro-batch "$MICRO_BATCH" \
  --grad-accum 1 \
  --max-tokens "$MAX_TOKENS" \
  --max-minutes "$MAX_MINUTES" \
  --checkpoint-steps 2000 \
  --checkpoint-minutes 15 \
  --no-gradient-checkpointing \
  --peak-tflops 989 \
  "${EXTRA_ARGS[@]}" \
  2>&1 | tee "$WORK_ROOT/train-${PROFILE}.log"

FINAL="$(find "$OUTPUT" -maxdepth 1 -type d -name 'step-*' | sort | tail -n 1)"
test -n "$FINAL"
cp "$TOKENIZER" "$FINAL/tokenizer.json"
python -m nayara.quantize "$FINAL"
printf '%s\n' "$FINAL" | tee "$WORK_ROOT/final-${PROFILE}.txt"
echo "Ready: $FINAL"