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

export RIEMANN_ROOT="${RIEMANN_ROOT:-/workspace/RiemannModel}"
export HF_HUB_ENABLE_HF_TRANSFER=1
export TOKENIZERS_PARALLELISM=false
export NCCL_DEBUG=WARN
export PYTHONUNBUFFERED=1
PROJECT="${PROJECT:-/workspace/riemann}"
STAGING_REPO="${STAGING_REPO:?set STAGING_REPO to the private Hub dataset holding data/, the stage 1 checkpoint and controller}"
LOGDIR="$RIEMANN_ROOT/logs"
mkdir -p "$RIEMANN_ROOT"/{data,checkpoints,ntk_controllers,hf_cache,staging} "$LOGDIR"
cd "$PROJECT"

log() { echo "$(date '+%F %T') [bootstrap] $*"; }
gate() { echo; echo "================ GATE: $* ================"; echo; }

gate "preflight"
log "host: $(hostname)  kernel: $(uname -r)"
df -h "$RIEMANN_ROOT" | tail -1
nproc; free -g | head -2
if command -v rocm-smi >/dev/null 2>&1; then rocm-smi --showproductname | grep -i "card series" | head -8 || true; fi
if command -v nvidia-smi >/dev/null 2>&1; then nvidia-smi -L || true; fi

gate "python environment"
VENV="/workspace/venv"
if [ ! -x "$VENV/bin/python" ]; then
  log "installing python3-venv/pip and creating venv"
  export DEBIAN_FRONTEND=noninteractive
  apt-get update -qq && apt-get install -y -qq python3-venv python3-pip git >/dev/null
  python3 -m venv "$VENV"
fi
export PATH="$VENV/bin:$PATH"
PY="$VENV/bin/python"
$PY -m pip install -q --upgrade pip
if [ "$($PY -c 'import torch;print(torch.cuda.is_available())' 2>/dev/null)" != "True" ]; then
  log "installing torch GPU build"
  if command -v rocm-smi >/dev/null 2>&1 || command -v amd-smi >/dev/null 2>&1; then
    $PY -m pip install -q torch --index-url "${TORCH_INDEX:-https://download.pytorch.org/whl/rocm6.4}"
  else
    $PY -m pip install -q torch --index-url https://download.pytorch.org/whl/cu124
  fi
fi
$PY -m pip install -q "transformers>=4.56,<5" "safetensors>=0.4" "huggingface_hub>=0.24" hf_transfer "numpy>=1.26" "pandas>=2.0" "pyarrow>=15" "datasets>=2.20" "accelerate>=0.33" "git+https://github.com/leochlon/ntkmirror.git"
$PY - <<'EOF'
import torch, transformers
print("torch", torch.__version__, "| hip:", getattr(torch.version, "hip", None), "| cuda:", torch.version.cuda)
n = torch.cuda.device_count()
print("gpu count:", n)
for i in range(n):
    p = torch.cuda.get_device_properties(i)
    print(f"  [{i}] {p.name}  {p.total_memory/2**30:.0f} GiB")
assert n == 8, f"expected 8 GPUs, found {n}"
print("transformers", transformers.__version__)
import ntkmirror; print("ntkmirror ok")
EOF

gate "pull staging data"
if [ ! -f "$RIEMANN_ROOT/data/pretrain_corpus.txt" ]; then
  hf download "$STAGING_REPO" --repo-type dataset --local-dir "$RIEMANN_ROOT/staging"
  log "decompressing corpus"
  gunzip -c "$RIEMANN_ROOT/staging/pretrain_corpus.txt.gz" > "$RIEMANN_ROOT/data/pretrain_corpus.txt"
  gunzip -c "$RIEMANN_ROOT/staging/ntk_finetune.jsonl.gz" > "$RIEMANN_ROOT/data/ntk_finetune.jsonl"
  cp -r "$RIEMANN_ROOT/staging/checkpoints/stage1_279m" "$RIEMANN_ROOT/checkpoints/"
  cp "$RIEMANN_ROOT/staging/ntk_controllers/"* "$RIEMANN_ROOT/ntk_controllers/"
fi
wc -l "$RIEMANN_ROOT/data/pretrain_corpus.txt" "$RIEMANN_ROOT/data/ntk_finetune.jsonl"
ls -la "$RIEMANN_ROOT/checkpoints/stage1_279m" "$RIEMANN_ROOT/ntk_controllers"

if [ "${SKIP_SMOKE:-0}" != "1" ]; then
gate "FSDP smoke test (tiny model, 8 ranks, ~2 min)"
$PY cloud/make_tiny_stage.py "$RIEMANN_ROOT/_smoke/checkpoints"
torchrun --nproc_per_node 8 scripts/pretrain_cuda.py --stage stage2_538m \
  --checkpoint-root "$RIEMANN_ROOT/_smoke/checkpoints" \
  --max-steps 6 --batch-size 2 --grad-accum 2 --seq-len 128 --log-every 2 --save-every 3 --num-workers 0 \
  2>&1 | tee "$LOGDIR/smoke.log" | grep -E "INFO|Error|Traceback" | tail -20
$PY -c "import sys; sys.path.insert(0,'.'); from riemann.modeling_riemann import RiemannForCausalLM; RiemannForCausalLM.from_pretrained('$RIEMANN_ROOT/_smoke/checkpoints/stage2_538m'); print('smoke checkpoint reloads OK')"
rm -rf "$RIEMANN_ROOT/_smoke"
fi

run_rung() {
  local from="$1" to="$2" epochs="$3"
  local marker="$RIEMANN_ROOT/checkpoints/$to/.rung_done"
  if [ -f "$marker" ]; then
    log "rung $to already complete, skipping"
    return 0
  fi
  gate "rung $from -> $to (epochs=$epochs)"
  local t0; t0=$(date +%s)
  if [ ! -f "$RIEMANN_ROOT/checkpoints/$to/config.json" ]; then
    $PY scripts/grow_and_continue.py --from-stage "$from" --to-stage "$to" 2>&1 | tee "$LOGDIR/grow_$to.log" | grep -E "INFO|Error" | tail -5
  fi
  local extra=""
  if [ -f "$RIEMANN_ROOT/pretrain_extra_$to.txt" ]; then
    extra="$(cat "$RIEMANN_ROOT/pretrain_extra_$to.txt")"
  elif [ -f "$RIEMANN_ROOT/pretrain_extra.txt" ]; then
    extra="$(cat "$RIEMANN_ROOT/pretrain_extra.txt")"
  fi
  log "pretrain extra flags for $to: ${extra:-<none>}"
  torchrun --nproc_per_node 8 scripts/pretrain_cuda.py --stage "$to" --epochs "$epochs" \
    --num-workers 6 --log-every 20 --save-every 250 $extra \
    2>&1 | tee "$LOGDIR/pretrain_$to.log" | grep -E "INFO \[rank 0\]|Error|Traceback"
  $PY scripts/fit_ntk.py --stage "$to" --init-from "$from" --gates 8192 --batch-size 8 --max-length 512 \
    2>&1 | tee "$LOGDIR/ntk_$to.log" | grep -E "INFO|Error|Traceback"
  touch "$marker"
  local t1; t1=$(date +%s)
  log "rung $to finished in $(( (t1 - t0) / 60 )) min"
}

gate "stage 2 alone (throughput check)"
run_rung stage1_279m stage2_538m 1
log "stage 2 done. Inspect $LOGDIR/pretrain_stage2_538m.log tok/s before continuing."
if [ "${STOP_AFTER_STAGE2:-0}" = "1" ]; then
  log "STOP_AFTER_STAGE2=1, exiting for review"
  exit 0
fi

run_rung stage2_538m stage3_1.19b 1
run_rung stage3_1.19b stage4_2.65b 1
run_rung stage4_2.65b stage5_6.66b 1
run_rung stage5_6.66b stage6_12.85b 2

gate "ALL DONE"
log "Final checkpoint: $RIEMANN_ROOT/checkpoints/stage6_12.85b. Controllers: $RIEMANN_ROOT/ntk_controllers. Archive these before destroying the node."