#!/bin/bash
# CUDA Kernel Benchmark: CubeCL vs Hand-crafted on NVIDIA GPUs
#
# Usage (on Linux node with CUDA):
#   ./scripts/run_cuda_kernel_bench.sh
#   ./scripts/run_cuda_kernel_bench.sh --sizes 128,256,512
#   ./scripts/run_cuda_kernel_bench.sh --iterations 200 --output-format json

set -e

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)"

echo ""
echo "=============================================="
echo "  RustyTorch CUDA Kernel Benchmark"
echo "  CubeCL vs Hand-crafted CUDA Kernels"
echo "=============================================="
echo ""

# Check for CUDA
if ! command -v nvcc &> /dev/null; then
    echo "WARNING: nvcc not found in PATH"
    echo "Attempting to load CUDA from common locations..."

    # Try common CUDA locations
    for cuda_path in /usr/local/cuda /opt/cuda /usr/cuda; do
        if [[ -d "$cuda_path" ]]; then
            export PATH="$cuda_path/bin:$PATH"
            export LD_LIBRARY_PATH="$cuda_path/lib64:$LD_LIBRARY_PATH"
            echo "Found CUDA at: $cuda_path"
            break
        fi
    done

    if ! command -v nvcc &> /dev/null; then
        echo "ERROR: CUDA toolkit not found. Please ensure CUDA is installed and nvcc is in PATH."
        exit 1
    fi
fi

# Check for NVIDIA GPU
if ! command -v nvidia-smi &> /dev/null; then
    echo "ERROR: nvidia-smi not found. NVIDIA driver may not be installed."
    exit 1
fi

# Default parameters
OPERATIONS="${OPERATIONS:-add,mul,exp,relu,gelu,silu}"
SIZES="${SIZES:-64,128,256}"
ITERATIONS="${ITERATIONS:-100}"
WARMUP="${WARMUP:-10}"
OUTPUT_FORMAT="${OUTPUT_FORMAT:-table}"
OUTPUT_FILE=""
VERIFY="${VERIFY:-false}"

# Parse command line arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        --operations)
            OPERATIONS="$2"
            shift 2
            ;;
        --sizes)
            SIZES="$2"
            shift 2
            ;;
        --iterations)
            ITERATIONS="$2"
            shift 2
            ;;
        --warmup)
            WARMUP="$2"
            shift 2
            ;;
        --output-format)
            OUTPUT_FORMAT="$2"
            shift 2
            ;;
        --output|-o)
            OUTPUT_FILE="$2"
            shift 2
            ;;
        --verify)
            VERIFY="true"
            shift
            ;;
        --help|-h)
            echo "Usage: $0 [OPTIONS]"
            echo ""
            echo "Options:"
            echo "  --operations OPS    Comma-separated operations (default: add,mul,exp,relu,gelu,silu)"
            echo "  --sizes SIZES       Comma-separated tensor sizes (default: 64,128,256)"
            echo "  --iterations N      Number of benchmark iterations (default: 100)"
            echo "  --warmup N          Number of warmup iterations (default: 10)"
            echo "  --output-format FMT Output format: table, markdown, json, csv (default: table)"
            echo "  --output, -o FILE   Output file path"
            echo "  --verify            Verify outputs match between implementations"
            echo "  --help, -h          Show this help message"
            exit 0
            ;;
        *)
            echo "Unknown option: $1"
            exit 1
            ;;
    esac
done

cd "$PROJECT_ROOT"

echo "System Info:"
echo "  OS: $(uname -s) $(uname -r)"
echo "  CUDA Version: $(nvcc --version | grep release | awk '{print $6}')"
echo ""

echo "GPU Info:"
nvidia-smi --query-gpu=name,memory.total,driver_version --format=csv,noheader | head -1
echo ""

echo "Benchmark Configuration:"
echo "  Operations: $OPERATIONS"
echo "  Sizes: $SIZES"
echo "  Iterations: $ITERATIONS (warmup: $WARMUP)"
echo "  Output format: $OUTPUT_FORMAT"
echo ""

# Build the benchmark binary with CUDA support
echo "Building benchmark binary with CUDA support..."
cargo build --release -p rtx-kernel-bench --features cuda --bin compare-kernels 2>&1 | tail -5

if [[ ! -f "$PROJECT_ROOT/target/release/compare-kernels" ]]; then
    echo "ERROR: Failed to build compare-kernels binary"
    exit 1
fi

echo ""
echo "Running benchmark..."
echo ""

# Build command arguments
CMD_ARGS=(
    --backend cuda
    --operations "$OPERATIONS"
    --sizes "$SIZES"
    --iterations "$ITERATIONS"
    --warmup "$WARMUP"
    --output-format "$OUTPUT_FORMAT"
)

if [[ -n "$OUTPUT_FILE" ]]; then
    # Create output directory if needed
    mkdir -p "$(dirname "$OUTPUT_FILE")"
    CMD_ARGS+=(--output "$OUTPUT_FILE")
fi

if [[ "$VERIFY" == "true" ]]; then
    CMD_ARGS+=(--verify)
fi

# Run the benchmark
"$PROJECT_ROOT/target/release/compare-kernels" "${CMD_ARGS[@]}"

# Save results with timestamp if no output file specified
if [[ -z "$OUTPUT_FILE" && "$OUTPUT_FORMAT" == "markdown" ]]; then
    TIMESTAMP=$(date +%Y%m%d_%H%M%S)
    RESULT_FILE="$PROJECT_ROOT/results/cuda_benchmark_${TIMESTAMP}.md"
    echo ""
    echo "Saving results to: $RESULT_FILE"
    "$PROJECT_ROOT/target/release/compare-kernels" "${CMD_ARGS[@]}" > "$RESULT_FILE"
fi

echo ""
echo "=============================================="
echo "  Benchmark Complete"
echo "=============================================="