#!/bin/bash
# One-time setup of a fresh Trn2 (Neuron DLAMI, SDK 2.33+) for the DeepInfer Neuron port.
#
#   git clone -b neuron https://github.com/whn09/DeepInfer.git && cd DeepInfer
#   bash tests/neuron/setup_trn2.sh            # WORK=~/di by default
#
# Creates $WORK/venv (layered on the DLAMI's read-only torch-neuronx venv), builds the C++ extension,
# downloads the official reference implementation + tokenizer, generates the tiny checkpoint and the
# CPU golden. Idempotent: finished steps are skipped. Host CPU only -- never touches a NeuronCore.
set -euo pipefail

REPO="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
WORK="${WORK:-$HOME/di}"
OPT_VENV="${OPT_VENV:-/opt/aws_neuronx_venv_torch_neuronx_2_14_0}"
HF_REPO="deepseek-ai/DeepSeek-V4.1-Flash-1017"
mkdir -p "$WORK"
log() { echo "[setup $(date +%H:%M:%S)] $*"; }

# 1. venv layered on the read-only DLAMI torch-neuronx venv
if [ ! -x "$WORK/venv/bin/python" ]; then
  log "venv -> $WORK/venv (layered on $OPT_VENV)"
  PY="$(readlink -f "$OPT_VENV/bin/python")"
  "$PY" -m venv "$WORK/venv"
  SP="$("$WORK/venv/bin/python" -c 'import site; print(site.getsitepackages()[0])')"
  echo "import site; site.addsitedir('$OPT_VENV/lib/python$("$PY" -c 'import sys; print(f"{sys.version_info[0]}.{sys.version_info[1]}")')/site-packages')" > "$SP/zz_opt_neuron.pth"
  # neuronx-cc and the neuron tools live in the DLAMI venv
  echo "export PATH=\"\$PATH:$OPT_VENV/bin\"" >> "$WORK/venv/bin/activate"
fi
# shellcheck disable=SC1091
source "$WORK/venv/bin/activate"
python -c "import torch, torch_neuronx; print('torch', torch.__version__, 'torch-neuronx', torch_neuronx.__version__)"
which neuronx-cc >/dev/null || { echo "neuronx-cc not on PATH"; exit 1; }

# 2. Python deps: requirements.txt minus torch / tilelang / tile-kernels (pure-torch shims are vendored)
log "pip deps"
pip install -q einops "huggingface_hub<1" ninja orjson pillow posix_ipc pybind11 pyzmq requests scipy sympy tabulate \
  tokenizers tqdm "transformers==4.50.0" xgrammar pytest psutil safetensors

# 3. C++ extension (host-only, -DWITH_NEURON)
if ! ls "$REPO"/deepinfer/_so/neuron-*/_C*.so >/dev/null 2>&1; then
  log "build C++ extension"
  (cd "$REPO" && DEEPINFER_ACCELERATOR=neuron MAX_JOBS=16 bash ./scripts/build.sh build > "$WORK/build.log" 2>&1) \
    || { tail -30 "$WORK/build.log"; exit 1; }
fi

# 4. official reference implementation + tokenizer + real config (small files only, no weights)
REF="$WORK/ref_dsv41"
if [ ! -f "$REF/inference/model.py" ]; then
  log "download reference implementation from $HF_REPO"
  mkdir -p "$REF/inference"
  for f in inference/model.py inference/kernel.py inference/convert.py inference/engram.py inference/generate.py \
           inference/vision.py inference/image_processor.py inference/config.json tokenizer.json tokenizer_config.json config.json; do
    curl -fsSL "https://huggingface.co/$HF_REPO/resolve/main/$f" -o "$REF/$f"
  done
fi

# 5. tiny checkpoint (~1.6 GB) and the CPU golden from the official reference
if [ ! -f "$WORK/tiny_ckpt/config.json" ]; then
  log "tiny checkpoint -> $WORK/tiny_ckpt"
  TORCH_DEVICE_BACKEND_AUTOLOAD=0 python "$REPO/examples/neuron/make_tiny_checkpoint.py" \
    --out "$WORK/tiny_ckpt" --tokenizer-dir "$REF" --real-config "$REF/config.json"
fi
if [ ! -f "$WORK/golden/long.pt" ]; then
  log "golden (official reference on CPU) -> $WORK/golden"
  TORCH_DEVICE_BACKEND_AUTOLOAD=0 python -I "$REPO/examples/neuron/reference_golden.py" \
    --hf-ckpt "$WORK/tiny_ckpt" --ref-dir "$REF/inference" --converted "$WORK/tiny_ckpt_ref_mp1" \
    --out "$WORK/golden" --threads 32
fi
log "done. Next: bash tests/neuron/run_all.sh"