#!/bin/bash
# Setup and run the layer test on B200 — no Docker, no vLLM
# Compiles the kernel raw and runs the comparison test

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO_DIR="$(dirname "$SCRIPT_DIR")"
VENV_DIR="$REPO_DIR/tests/.venv"

echo "=== NVFP4 Layer Test Setup ==="
echo "Repo: $REPO_DIR"
echo ""

# 1. Create venv
if [ ! -d "$VENV_DIR" ]; then
    echo "[1/4] Creating venv..."
    python3 -m venv "$VENV_DIR"
else
    echo "[1/4] Venv already exists, skipping"
fi

source "$VENV_DIR/bin/activate"

# 2. Install dependencies
echo "[2/4] Installing Python deps..."
pip install --upgrade pip -q
pip install -r "$SCRIPT_DIR/requirements.txt" -q

# 3. Build and install the kernel from source
echo "[3/4] Building kernel from source (this takes a few minutes)..."
cd "$REPO_DIR"
pip install -e . --no-build-isolation -q

# Build the CUTLASS NVFP4 GEMM C++ extension
echo "       Building CUTLASS C++ extension..."
cd "$REPO_DIR/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm"
python setup.py install

# 4. Run the test
echo "[4/4] Running layer test..."
cd "$SCRIPT_DIR"
python3 layertest.py

echo ""
echo "=== Done ==="