#!/bin/bash
#
# RustyTorch++ Metal & WASM Benchmark Runner
#
# This script runs comprehensive benchmarks comparing RustyTorch++ against
# PyTorch MPS for Apple Silicon, and WASM inference comparisons.
#
# Usage:
#   ./scripts/run_metal_benchmarks.sh [--quick] [--rust-only] [--python-only] [--wasm]
#
# Prerequisites:
#   - Apple Silicon Mac (M1/M2/M3/M4)
#   - PyTorch with MPS support: pip install torch
#   - Rust toolchain with wasm-pack: cargo install wasm-pack
#   - Node.js 18+ for WASM benchmarks
#

set -e

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Configuration
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(dirname "$SCRIPT_DIR")"
BENCHMARKS_DIR="$PROJECT_ROOT/benchmarks"
REPORTS_DIR="$BENCHMARKS_DIR/reports"
TIMESTAMP=$(date +"%Y%m%d_%H%M%S")

# Default options
QUICK_MODE=false
RUST_ONLY=false
PYTHON_ONLY=false
RUN_WASM=false
ITERATIONS=1000
WARMUP=50

# Parse arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        --quick)
            QUICK_MODE=true
            ITERATIONS=100
            WARMUP=10
            shift
            ;;
        --rust-only)
            RUST_ONLY=true
            shift
            ;;
        --python-only)
            PYTHON_ONLY=true
            shift
            ;;
        --wasm)
            RUN_WASM=true
            shift
            ;;
        --help)
            echo "Usage: $0 [--quick] [--rust-only] [--python-only] [--wasm]"
            echo ""
            echo "Options:"
            echo "  --quick       Run with fewer iterations (quick validation)"
            echo "  --rust-only   Only run Rust benchmarks"
            echo "  --python-only Only run Python/PyTorch benchmarks"
            echo "  --wasm        Also run WASM browser benchmarks"
            echo ""
            exit 0
            ;;
        *)
            echo "Unknown option: $1"
            exit 1
            ;;
    esac
done

# Print banner
echo -e "${BLUE}"
echo "╔══════════════════════════════════════════════════════════════════╗"
echo "║     RustyTorch++ Metal & WASM Benchmark Suite                    ║"
echo "║     Comparing against PyTorch MPS on Apple Silicon               ║"
echo "╚══════════════════════════════════════════════════════════════════╝"
echo -e "${NC}"

# Check prerequisites
echo -e "${YELLOW}Checking prerequisites...${NC}"

# Check for Apple Silicon
if [[ "$(uname -m)" != "arm64" ]]; then
    echo -e "${RED}Error: This benchmark requires Apple Silicon (arm64)${NC}"
    exit 1
fi
echo -e "  ${GREEN}✓${NC} Apple Silicon detected"

# Check for Rust
if ! command -v cargo &> /dev/null; then
    echo -e "${RED}Error: Rust/Cargo not found${NC}"
    exit 1
fi
echo -e "  ${GREEN}✓${NC} Rust installed: $(rustc --version)"

# Check for Python
if ! command -v python3 &> /dev/null; then
    echo -e "${RED}Error: Python3 not found${NC}"
    exit 1
fi
echo -e "  ${GREEN}✓${NC} Python installed: $(python3 --version)"

# Check for PyTorch MPS
if ! python3 -c "import torch; assert torch.backends.mps.is_available()" 2>/dev/null; then
    echo -e "${YELLOW}Warning: PyTorch MPS not available, skipping PyTorch benchmarks${NC}"
    PYTHON_ONLY=false
    if [[ "$RUST_ONLY" == false ]]; then
        RUST_ONLY=true
    fi
else
    echo -e "  ${GREEN}✓${NC} PyTorch MPS available"
fi

# Create reports directory
mkdir -p "$REPORTS_DIR"

echo ""
echo -e "${BLUE}Configuration:${NC}"
echo "  Iterations: $ITERATIONS"
echo "  Warmup: $WARMUP"
echo "  Quick mode: $QUICK_MODE"
echo "  Reports dir: $REPORTS_DIR"
echo ""

# ============================================================================
# Run Rust Benchmarks
# ============================================================================
if [[ "$PYTHON_ONLY" == false ]]; then
    echo -e "${BLUE}════════════════════════════════════════════════════════════════════${NC}"
    echo -e "${BLUE}Running Rust Metal Benchmarks${NC}"
    echo -e "${BLUE}════════════════════════════════════════════════════════════════════${NC}"
    echo ""

    cd "$PROJECT_ROOT"

    # Flash Attention benchmark
    echo -e "${YELLOW}[1/3] Flash Attention Metal benchmark...${NC}"
    if cargo bench -p rtx-flash-metal-attention --bench metal_vs_pytorch 2>/dev/null; then
        echo -e "${GREEN}  ✓ Flash Attention benchmark complete${NC}"
    else
        echo -e "${YELLOW}  ⚠ Flash Attention benchmark skipped (not available)${NC}"
    fi

    # MoE benchmark
    echo ""
    echo -e "${YELLOW}[2/3] MoE Metal benchmark...${NC}"
    if cargo bench -p rtx-transformers --bench metal_moe_bench 2>/dev/null; then
        echo -e "${GREEN}  ✓ MoE benchmark complete${NC}"
    else
        echo -e "${YELLOW}  ⚠ MoE benchmark skipped (not available)${NC}"
    fi

    # Mamba benchmark
    echo ""
    echo -e "${YELLOW}[3/3] Mamba Metal benchmark...${NC}"
    if cargo bench -p rtx-transformers --bench metal_mamba_bench 2>/dev/null; then
        echo -e "${GREEN}  ✓ Mamba benchmark complete${NC}"
    else
        echo -e "${YELLOW}  ⚠ Mamba benchmark skipped (not available)${NC}"
    fi

    echo ""
fi

# ============================================================================
# Run Python/PyTorch Benchmarks
# ============================================================================
if [[ "$RUST_ONLY" == false ]]; then
    echo -e "${BLUE}════════════════════════════════════════════════════════════════════${NC}"
    echo -e "${BLUE}Running PyTorch MPS Benchmarks${NC}"
    echo -e "${BLUE}════════════════════════════════════════════════════════════════════${NC}"
    echo ""

    cd "$PROJECT_ROOT"

    # Flash Attention
    echo -e "${YELLOW}[1/3] PyTorch Flash Attention MPS benchmark...${NC}"
    python3 benchmarks/metal/bench_flash_attention.py \
        --json \
        --iterations "$ITERATIONS" \
        --warmup "$WARMUP" \
        > "$REPORTS_DIR/pytorch_flash_attention_$TIMESTAMP.json"
    echo -e "${GREEN}  ✓ Results saved to pytorch_flash_attention_$TIMESTAMP.json${NC}"

    # MoE
    echo ""
    echo -e "${YELLOW}[2/3] PyTorch MoE MPS benchmark...${NC}"
    python3 benchmarks/metal/bench_moe_mps.py \
        --json \
        --iterations "$((ITERATIONS / 2))" \
        --warmup "$WARMUP" \
        > "$REPORTS_DIR/pytorch_moe_$TIMESTAMP.json"
    echo -e "${GREEN}  ✓ Results saved to pytorch_moe_$TIMESTAMP.json${NC}"

    # Mamba
    echo ""
    echo -e "${YELLOW}[3/3] PyTorch Mamba MPS benchmark...${NC}"
    python3 benchmarks/metal/bench_mamba_mps.py \
        --json \
        --iterations "$((ITERATIONS / 5))" \
        --warmup "$WARMUP" \
        > "$REPORTS_DIR/pytorch_mamba_$TIMESTAMP.json"
    echo -e "${GREEN}  ✓ Results saved to pytorch_mamba_$TIMESTAMP.json${NC}"

    echo ""
fi

# ============================================================================
# Run WASM Benchmarks
# ============================================================================
if [[ "$RUN_WASM" == true ]]; then
    echo -e "${BLUE}════════════════════════════════════════════════════════════════════${NC}"
    echo -e "${BLUE}Building WASM Module${NC}"
    echo -e "${BLUE}════════════════════════════════════════════════════════════════════${NC}"
    echo ""

    cd "$PROJECT_ROOT"

    # Build WASM module
    echo -e "${YELLOW}Building rtx-wasm-inference...${NC}"
    if wasm-pack build crates/production/rtx-wasm-inference --target web 2>/dev/null; then
        echo -e "${GREEN}  ✓ WASM module built successfully${NC}"

        # Copy to benchmarks directory
        cp -r crates/production/rtx-wasm-inference/pkg benchmarks/wasm/

        echo ""
        echo -e "${YELLOW}To run browser benchmarks:${NC}"
        echo "  cd benchmarks/wasm && npm install && npm run serve"
        echo "  Open http://localhost:8080/comparison.html"
    else
        echo -e "${YELLOW}  ⚠ WASM build failed (check wasm-pack installation)${NC}"
    fi

    echo ""
fi

# ============================================================================
# Summary
# ============================================================================
echo -e "${BLUE}════════════════════════════════════════════════════════════════════${NC}"
echo -e "${BLUE}Benchmark Summary${NC}"
echo -e "${BLUE}════════════════════════════════════════════════════════════════════${NC}"
echo ""

echo "Results saved to: $REPORTS_DIR"
echo ""
echo "Files generated:"
ls -la "$REPORTS_DIR"/*.json 2>/dev/null || echo "  (no JSON files yet)"

echo ""
echo -e "${GREEN}Benchmark run complete!${NC}"
echo ""
echo "Next steps:"
echo "  1. View PyTorch results: cat $REPORTS_DIR/pytorch_*.json | jq ."
echo "  2. Compare with Rust: cargo bench -p rtx-flash-metal-attention"
echo "  3. For WASM browser tests: Run with --wasm flag"
echo ""