#!/bin/bash
# Metal Kernel Benchmark: CubeCL vs Hand-crafted on Apple Silicon
#
# Usage:
#   ./scripts/run_metal_kernel_bench.sh
#   ./scripts/run_metal_kernel_bench.sh --sizes 128,256,512
#   ./scripts/run_metal_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 Metal Kernel Benchmark"
echo "  CubeCL vs Hand-crafted MSL Kernels"
echo "=============================================="
echo ""

# Check we're on macOS
if [[ "$(uname)" != "Darwin" ]]; then
    echo "ERROR: This script must be run on macOS (requires Metal)"
    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 "  macOS Version: $(sw_vers -productVersion)"
echo "  Chip: $(sysctl -n machdep.cpu.brand_string 2>/dev/null || system_profiler SPHardwareDataType | grep 'Chip' | awk -F': ' '{print $2}')"
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 Metal support
echo "Building benchmark binary..."
cargo build --release -p rtx-kernel-bench --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 metal
    --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/metal_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 "=============================================="