#!/bin/bash
#
# PINN Benchmark: Python (PyTorch) vs Rust (RustyTorch++)
#
# This script runs benchmarks for both implementations and generates a comparison report.
#
# Usage:
#   ./comparison/run_benchmark.sh [--quick] [--cuda] [--venv] [--cpu-only] [--gpu-only]
#
# Options:
#   --quick     Run quick benchmarks with fewer iterations
#   --cuda      Run GPU benchmarks (both Python and Rust)
#   --venv      Create and use a Python virtual environment
#   --cpu-only  Run only CPU benchmarks (skip GPU)
#   --gpu-only  Run only GPU benchmarks (skip CPU)

set -e

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

# Use python3 by default, allow override via PYTHON env var
PYTHON="${PYTHON:-python3}"

echo "=== PINN Benchmark: Python vs Rust ==="
echo "Date: $(date)"
echo "Project: $PROJECT_DIR"
echo ""

# Parse arguments
QUICK_MODE=false
CUDA_MODE=false
VENV_MODE=false
CPU_ONLY=false
GPU_ONLY=false
for arg in "$@"; do
    case $arg in
        --quick)
            QUICK_MODE=true
            ;;
        --cuda)
            CUDA_MODE=true
            ;;
        --venv)
            VENV_MODE=true
            ;;
        --cpu-only)
            CPU_ONLY=true
            ;;
        --gpu-only)
            GPU_ONLY=true
            ;;
    esac
done

# Determine what to run
RUN_CPU=true
RUN_GPU=false

if [ "$GPU_ONLY" = true ]; then
    RUN_CPU=false
    RUN_GPU=true
elif [ "$CPU_ONLY" = true ]; then
    RUN_CPU=true
    RUN_GPU=false
elif [ "$CUDA_MODE" = true ]; then
    RUN_CPU=true
    RUN_GPU=true
fi

# Set point sizes based on quick mode
if [ "$QUICK_MODE" = true ]; then
    POINT_SIZES="200,1000"
else
    POINT_SIZES="200,1000,10000"
fi

# Create results directory
mkdir -p "$SCRIPT_DIR/results"

cd "$PROJECT_DIR"

# Setup virtual environment if requested
if [ "$VENV_MODE" = true ]; then
    echo ">>> Setting up Python virtual environment..."
    if [ ! -d ".venv" ]; then
        $PYTHON -m venv .venv
        echo "    Created .venv"
    fi
    source .venv/bin/activate
    echo "    Activated .venv"

    echo ">>> Installing Python dependencies..."
    pip install -q -r python/requirements.txt
    echo "    Dependencies installed"
    echo ""

    # Use the venv python
    PYTHON="python"
fi

# Check if Python is available
if ! command -v $PYTHON &> /dev/null; then
    echo "Error: $PYTHON not found. Please install Python 3 or use --venv flag."
    exit 1
fi

echo "Configuration:"
echo "  Quick mode: $QUICK_MODE"
echo "  Run CPU benchmarks: $RUN_CPU"
echo "  Run GPU benchmarks: $RUN_GPU"
echo "  Point sizes: $POINT_SIZES"
echo ""

# =============================================================================
# PYTHON BENCHMARKS
# =============================================================================

if [ "$RUN_CPU" = true ]; then
    echo ">>> Running Python (PyTorch) benchmarks (CPU)..."
    echo "    Using: $PYTHON"
    $PYTHON python/benchmark_runner.py \
        --output comparison/results/python_cpu_results.json \
        --point-sizes "$POINT_SIZES" \
        --device cpu
    echo ""
fi

if [ "$RUN_GPU" = true ]; then
    echo ">>> Running Python (PyTorch) benchmarks (GPU)..."
    echo "    Using: $PYTHON"
    $PYTHON python/benchmark_runner.py \
        --output comparison/results/python_gpu_results.json \
        --point-sizes "$POINT_SIZES" \
        --device cuda
    echo ""
fi

# =============================================================================
# RUST BENCHMARKS
# =============================================================================

if [ "$RUN_CPU" = true ]; then
    echo ">>> Running Rust (RustyTorch++) benchmarks (CPU)..."

    if [ "$QUICK_MODE" = true ]; then
        cargo bench --bench pinn_benchmark -- --noplot --save-baseline rust_cpu 2>&1 | tee comparison/results/rust_cpu_output.txt
    else
        cargo bench --bench pinn_benchmark -- --save-baseline rust_cpu 2>&1 | tee comparison/results/rust_cpu_output.txt
    fi

    # Copy criterion JSON if available
    if [ -d "target/criterion" ]; then
        find target/criterion -name "estimates.json" -exec cp {} comparison/results/rust_cpu_estimates.json \; 2>/dev/null || true
    fi
    echo ""
fi

if [ "$RUN_GPU" = true ]; then
    echo ">>> Running Rust (RustyTorch++) benchmarks (CUDA)..."

    # Check if CUDA feature compiles
    if cargo check --features cuda 2>/dev/null; then
        if [ "$QUICK_MODE" = true ]; then
            cargo bench --bench pinn_benchmark --features cuda -- --noplot --save-baseline rust_cuda 2>&1 | tee comparison/results/rust_cuda_output.txt
        else
            cargo bench --bench pinn_benchmark --features cuda -- --save-baseline rust_cuda 2>&1 | tee comparison/results/rust_cuda_output.txt
        fi

        if [ -d "target/criterion" ]; then
            find target/criterion -name "estimates.json" -exec cp {} comparison/results/rust_cuda_estimates.json \; 2>/dev/null || true
        fi
    else
        echo "Warning: CUDA feature not available, skipping Rust CUDA benchmarks"
    fi
    echo ""
fi

# =============================================================================
# GENERATE REPORT
# =============================================================================

echo ">>> Generating comparison report..."

# Build report generator arguments
REPORT_ARGS="--output $SCRIPT_DIR/results/comparison_report.md"

if [ -f "$SCRIPT_DIR/results/python_cpu_results.json" ]; then
    REPORT_ARGS="$REPORT_ARGS --python-cpu-results $SCRIPT_DIR/results/python_cpu_results.json"
fi

if [ -f "$SCRIPT_DIR/results/python_gpu_results.json" ]; then
    REPORT_ARGS="$REPORT_ARGS --python-gpu-results $SCRIPT_DIR/results/python_gpu_results.json"
fi

if [ -f "$SCRIPT_DIR/results/rust_cpu_output.txt" ]; then
    REPORT_ARGS="$REPORT_ARGS --rust-cpu-output $SCRIPT_DIR/results/rust_cpu_output.txt"
fi

if [ -f "$SCRIPT_DIR/results/rust_cuda_output.txt" ]; then
    REPORT_ARGS="$REPORT_ARGS --rust-cuda-output $SCRIPT_DIR/results/rust_cuda_output.txt"
fi

$PYTHON "$SCRIPT_DIR/generate_report.py" $REPORT_ARGS

echo ""
echo "=== Benchmark Complete ==="
echo ""
echo "Results directory: $SCRIPT_DIR/results/"
ls -la "$SCRIPT_DIR/results/"
echo ""

# Display summary
if [ -f "$SCRIPT_DIR/results/comparison_report.md" ]; then
    echo "--- Report Summary ---"
    head -80 "$SCRIPT_DIR/results/comparison_report.md"
fi