#!/bin/bash
# =============================================================================
# Unified PyTorch vs RustyTorch++ PINN Benchmark
# =============================================================================
#
# This script runs comprehensive PINN benchmarks comparing PyTorch and RustyTorch++
# on both CPU and GPU, generating a unified HTML report with performance charts.
#
# Usage:
#   ./scripts/run_pinn_comparison.sh [OPTIONS]
#
# Options:
#   --quick       Run quick benchmarks (fewer iterations/point sizes)
#   --cpu-only    Run only CPU benchmarks (skip GPU)
#   --gpu-only    Run only GPU benchmarks (skip CPU)
#   --no-venv     Use system Python instead of virtual environment
#   -h, --help    Show this help message
#
# Output:
#   benchmark_reports/<hostname>-rust-py-pinn-benchmark-<MM-DD-YYYY>/
#     - system_info.json
#     - pytorch_cpu_results.json
#     - pytorch_gpu_results.json
#     - rust_cpu_output.txt
#     - rust_gpu_output.txt
#     - comparison_report.html
#     - comparison_report.md
#
# =============================================================================

set -e

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

# =============================================================================
# Progress Indicator Functions
# =============================================================================

# Track start time
START_SECONDS=$SECONDS
TOTAL_STEPS=5
CURRENT_STEP=0

# Format time as Xm Ys
format_time() {
    local secs=$1
    if [[ $secs -lt 60 ]]; then
        printf "%ds" $secs
    else
        printf "%dm %ds" $((secs/60)) $((secs%60))
    fi
}

# Display progress bar with current step and elapsed time
show_progress() {
    local step=$1
    local step_name=$2
    local width=40
    local percent=$((step * 100 / TOTAL_STEPS))
    local filled=$((width * step / TOTAL_STEPS))
    local empty=$((width - filled))
    local elapsed=$((SECONDS - START_SECONDS))

    # Build progress bar
    local bar=""
    for ((i=0; i<filled; i++)); do bar+="█"; done
    for ((i=0; i<empty; i++)); do bar+="░"; done

    # Clear previous progress display (3 lines + header box)
    echo -ne "\033[6A\033[J"

    # Draw progress box
    echo -e "${CYAN}╔══════════════════════════════════════════════════════════════════════╗${NC}"
    echo -e "${CYAN}║  ${BOLD}PyTorch vs RustyTorch++ PINN Benchmark${NC}${CYAN}                              ║${NC}"
    echo -e "${CYAN}╠══════════════════════════════════════════════════════════════════════╣${NC}"
    printf "${CYAN}║${NC}  Progress: [${GREEN}%s${NC}] %3d%% (%d/%d steps)                  ${CYAN}║${NC}\n" "$bar" "$percent" "$step" "$TOTAL_STEPS"
    printf "${CYAN}║${NC}  Current:  %-54s ${CYAN}║${NC}\n" "$step_name"
    printf "${CYAN}║${NC}  Elapsed:  %-54s ${CYAN}║${NC}\n" "$(format_time $elapsed)"
    echo -e "${CYAN}╚══════════════════════════════════════════════════════════════════════╝${NC}"
}

# Initial progress display (placeholder before first update)
init_progress() {
    echo ""
    echo ""
    echo ""
    echo ""
    echo ""
    echo ""
    echo ""
}

# Script location and project root
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)"
PINN_DIR="$PROJECT_ROOT/examples/pinn_mre_helmholtz"

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

# Configuration defaults
QUICK_MODE=false
CPU_ONLY=false
GPU_ONLY=false
VENV_MODE=true  # Always use venv by default

# Auto-detect platform
if [[ "$(uname)" == "Darwin" ]]; then
    GPU_BACKEND="metal"
    PLATFORM="macOS"
else
    GPU_BACKEND="cuda"
    PLATFORM="Linux"
fi

# Hostname and date for output directory
HOSTNAME_SHORT=$(hostname | cut -d'.' -f1)
DATE_STR=$(date +"%m-%d-%Y")
OUTPUT_DIR="$PROJECT_ROOT/benchmark_reports/${HOSTNAME_SHORT}-rust-py-pinn-benchmark-${DATE_STR}"

# =============================================================================
# Parse command line arguments
# =============================================================================
show_help() {
    echo "Unified PyTorch vs RustyTorch++ PINN Benchmark"
    echo ""
    echo "Usage: $0 [OPTIONS]"
    echo ""
    echo "Options:"
    echo "  --quick       Run quick benchmarks (fewer iterations)"
    echo "  --cpu-only    Run only CPU benchmarks (skip GPU)"
    echo "  --gpu-only    Run only GPU benchmarks (skip CPU)"
    echo "  --no-venv     Use system Python instead of virtual environment"
    echo "  -h, --help    Show this help message"
    echo ""
    echo "Output:"
    echo "  benchmark_reports/<hostname>-rust-py-pinn-benchmark-<date>/"
}

while [[ $# -gt 0 ]]; do
    case $1 in
        --quick)
            QUICK_MODE=true
            shift
            ;;
        --cpu-only)
            CPU_ONLY=true
            shift
            ;;
        --gpu-only)
            GPU_ONLY=true
            shift
            ;;
        --no-venv)
            VENV_MODE=false
            shift
            ;;
        -h|--help)
            show_help
            exit 0
            ;;
        *)
            echo "Unknown option: $1"
            show_help
            exit 1
            ;;
    esac
done

# Determine what to run
RUN_CPU=true
RUN_GPU=true

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

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

# =============================================================================
# Banner and Initial Progress Display
# =============================================================================
echo -e "Host: ${YELLOW}${HOSTNAME_SHORT}${NC} | Platform: ${YELLOW}${PLATFORM}${NC} | GPU: ${YELLOW}${GPU_BACKEND}${NC}"
echo -e "Output: ${CYAN}${OUTPUT_DIR}${NC}"
echo ""

# Initialize progress display (creates space for progress box)
init_progress
show_progress 0 "Initializing..."

# =============================================================================
# Create output directory
# =============================================================================
mkdir -p "$OUTPUT_DIR"

# =============================================================================
# Phase 1: Collect System Information
# =============================================================================
show_progress 1 "Collecting system information..."

SYSTEM_INFO_FILE="$OUTPUT_DIR/system_info.json"

# Collect OS info based on platform
if [[ "$PLATFORM" == "macOS" ]]; then
    OS_NAME="macOS"
    OS_VERSION=$(sw_vers -productVersion 2>/dev/null || echo "Unknown")
    KERNEL=$(uname -r)
    CPU_MODEL=$(sysctl -n machdep.cpu.brand_string 2>/dev/null || echo "Apple Silicon")
    CPU_CORES=$(sysctl -n hw.physicalcpu 2>/dev/null || echo "Unknown")
    CPU_THREADS=$(sysctl -n hw.logicalcpu 2>/dev/null || echo "Unknown")
    TOTAL_MEM_BYTES=$(sysctl -n hw.memsize 2>/dev/null || echo "0")
    TOTAL_MEM="$((TOTAL_MEM_BYTES / 1024 / 1024 / 1024)) GB"
    GPU_NAME="Apple Silicon GPU"
    GPU_MEMORY="Unified Memory"
    GPU_DRIVER="Metal"
    GPU_COMPUTE="N/A"
else
    OS_NAME=$(cat /etc/os-release 2>/dev/null | grep "^NAME=" | cut -d'"' -f2 || echo "Unknown")
    OS_VERSION=$(cat /etc/os-release 2>/dev/null | grep "^VERSION=" | cut -d'"' -f2 || echo "Unknown")
    KERNEL=$(uname -r)
    CPU_MODEL=$(lscpu 2>/dev/null | grep "Model name" | cut -d':' -f2 | xargs || echo "Unknown")
    CPU_CORES=$(lscpu 2>/dev/null | grep "^CPU(s):" | cut -d':' -f2 | xargs || echo "Unknown")
    CPU_THREADS=$(lscpu 2>/dev/null | grep "Thread(s) per core" | cut -d':' -f2 | xargs || echo "1")
    TOTAL_MEM=$(free -h 2>/dev/null | grep "Mem:" | awk '{print $2}' || echo "Unknown")

    # GPU info
    if command -v nvidia-smi &> /dev/null; then
        GPU_INFO=$(nvidia-smi --query-gpu=name,memory.total,driver_version,compute_cap --format=csv,noheader 2>/dev/null || echo "")
        if [[ -n "$GPU_INFO" ]]; then
            GPU_NAME=$(echo "$GPU_INFO" | cut -d',' -f1 | xargs)
            GPU_MEMORY=$(echo "$GPU_INFO" | cut -d',' -f2 | xargs)
            GPU_DRIVER=$(echo "$GPU_INFO" | cut -d',' -f3 | xargs)
            GPU_COMPUTE=$(echo "$GPU_INFO" | cut -d',' -f4 | xargs)
        else
            GPU_NAME="Not detected"
            GPU_MEMORY="N/A"
            GPU_DRIVER="N/A"
            GPU_COMPUTE="N/A"
        fi
    else
        GPU_NAME="nvidia-smi not available"
        GPU_MEMORY="N/A"
        GPU_DRIVER="N/A"
        GPU_COMPUTE="N/A"
    fi
fi

# Get git info
GIT_COMMIT=$(cd "$PROJECT_ROOT" && git rev-parse --short HEAD 2>/dev/null || echo "unknown")
GIT_BRANCH=$(cd "$PROJECT_ROOT" && git branch --show-current 2>/dev/null || echo "unknown")

# Write system info JSON
cat > "$SYSTEM_INFO_FILE" << EOF
{
  "timestamp": "$(date -Iseconds)",
  "hostname": "$HOSTNAME_SHORT",
  "commit": "$GIT_COMMIT",
  "branch": "$GIT_BRANCH",
  "os": {
    "name": "$OS_NAME",
    "version": "$OS_VERSION",
    "kernel": "$KERNEL",
    "arch": "$(uname -m)"
  },
  "cpu": {
    "model": "$CPU_MODEL",
    "cores": "$CPU_CORES",
    "threads_per_core": "$CPU_THREADS"
  },
  "gpu": {
    "name": "$GPU_NAME",
    "memory": "$GPU_MEMORY",
    "driver": "$GPU_DRIVER",
    "compute_capability": "$GPU_COMPUTE"
  },
  "memory": {
    "total": "$TOTAL_MEM"
  }
}
EOF

# =============================================================================
# Phase 2: Setup Python Environment
# =============================================================================
show_progress 2 "Setting up Python environment..."

cd "$PINN_DIR"

if [[ "$VENV_MODE" == true ]]; then
    if [[ ! -d ".venv" ]]; then
        $PYTHON -m venv .venv
    fi
    source .venv/bin/activate
    PYTHON="python"

    # Install dependencies if needed
    if [[ -f "python/requirements.txt" ]]; then
        pip install -q -r python/requirements.txt 2>/dev/null || true
    fi
fi

# Check if Python and torch are available
if ! $PYTHON -c "import torch" 2>/dev/null; then
    SKIP_PYTHON=true
else
    SKIP_PYTHON=false
fi

# =============================================================================
# Phase 3: Run PyTorch Benchmarks
# =============================================================================
show_progress 3 "Running PyTorch benchmarks..."

if [[ "$SKIP_PYTHON" != true ]]; then
    if [[ "$RUN_CPU" == true ]]; then
        show_progress 3 "Running PyTorch CPU benchmarks..."
        $PYTHON python/benchmark_runner.py \
            --device cpu \
            --point-sizes "$POINT_SIZES" \
            --output "$OUTPUT_DIR/pytorch_cpu_results.json" >/dev/null 2>&1 || true
    fi

    if [[ "$RUN_GPU" == true ]]; then
        if $PYTHON -c "import torch; assert torch.cuda.is_available()" 2>/dev/null; then
            show_progress 3 "Running PyTorch GPU benchmarks..."
            $PYTHON python/benchmark_runner.py \
                --device cuda \
                --point-sizes "$POINT_SIZES" \
                --output "$OUTPUT_DIR/pytorch_gpu_results.json" >/dev/null 2>&1 || true
        fi
    fi
fi

# =============================================================================
# Phase 4: Run Rust Benchmarks
# =============================================================================
show_progress 4 "Running RustyTorch++ benchmarks..."

cd "$PINN_DIR"

if [[ "$RUN_CPU" == true ]]; then
    show_progress 4 "Running RustyTorch++ CPU benchmarks..."
    cargo bench --bench pinn_benchmark $RUST_BENCH_ARGS 2>&1 > "$OUTPUT_DIR/rust_cpu_output.txt" || true
fi

if [[ "$RUN_GPU" == true ]]; then
    if cargo check --features cuda 2>/dev/null; then
        show_progress 4 "Running RustyTorch++ GPU benchmarks..."
        cargo bench --bench pinn_benchmark --features cuda $RUST_BENCH_ARGS 2>&1 > "$OUTPUT_DIR/rust_gpu_output.txt" || true
    fi
fi

# =============================================================================
# Phase 5: Generate Report
# =============================================================================
show_progress 5 "Generating comparison report..."

cd "$PROJECT_ROOT"

# Build report generator arguments
REPORT_ARGS="--output-dir $OUTPUT_DIR"
REPORT_ARGS="$REPORT_ARGS --system-info $OUTPUT_DIR/system_info.json"

[[ -f "$OUTPUT_DIR/pytorch_cpu_results.json" ]] && REPORT_ARGS="$REPORT_ARGS --pytorch-cpu $OUTPUT_DIR/pytorch_cpu_results.json"
[[ -f "$OUTPUT_DIR/pytorch_gpu_results.json" ]] && REPORT_ARGS="$REPORT_ARGS --pytorch-gpu $OUTPUT_DIR/pytorch_gpu_results.json"
[[ -f "$OUTPUT_DIR/rust_cpu_output.txt" ]] && REPORT_ARGS="$REPORT_ARGS --rust-cpu $OUTPUT_DIR/rust_cpu_output.txt"
[[ -f "$OUTPUT_DIR/rust_gpu_output.txt" ]] && REPORT_ARGS="$REPORT_ARGS --rust-gpu $OUTPUT_DIR/rust_gpu_output.txt"

$PYTHON "$SCRIPT_DIR/generate_comparison_report.py" $REPORT_ARGS >/dev/null 2>&1

# =============================================================================
# Completion - Clear progress and show final summary
# =============================================================================
TOTAL_ELAPSED=$((SECONDS - START_SECONDS))

# Clear the progress display
echo -ne "\033[7A\033[J"

# Final completion banner
echo -e "${GREEN}╔════════════════════════════════════════════════════════════════════╗${NC}"
echo -e "${GREEN}║     ${BOLD}PyTorch vs RustyTorch++ PINN Benchmark Complete!${NC}${GREEN}              ║${NC}"
echo -e "${GREEN}╚════════════════════════════════════════════════════════════════════╝${NC}"
echo ""
echo -e "Output: ${CYAN}$OUTPUT_DIR${NC}"
echo ""
echo -e "Total time: ${YELLOW}$(format_time $TOTAL_ELAPSED)${NC}"
echo ""
echo "Files generated:"
ls -1 "$OUTPUT_DIR" | while read f; do
    echo -e "  - ${CYAN}$f${NC}"
done
echo ""
echo -e "View report: ${YELLOW}firefox $OUTPUT_DIR/comparison_report.html${NC}"
echo ""