#!/usr/bin/env bash
# ─────────────────────────────────────────────────────────────────────────────
# bench.sh — OSF-RS CPU vs GPU inference benchmark
#
# Detects platform, CPU, GPU, runs benchmarks for both backends,
# saves JSON results, and generates comparison charts.
#
# Usage:
#   ./bench.sh                          # default: 3 warmup, 10 runs
#   ./bench.sh 5 20                     # 5 warmup, 20 timed runs
#   WEIGHTS=path/to/osf.safetensors ./bench.sh
# ─────────────────────────────────────────────────────────────────────────────
set -euo pipefail

WARMUP="${1:-3}"
RUNS="${2:-10}"
WEIGHTS="${WEIGHTS:-data/osf_backbone.safetensors}"

if [ ! -f "$WEIGHTS" ]; then
    echo "Error: weights file not found: $WEIGHTS"
    echo "Run: python scripts/export_safetensors.py --input ../OSF-Base/osf_backbone.pth --output $WEIGHTS"
    exit 1
fi

# ─── Platform Detection ─────────────────────────────────────────────────────

OS="$(uname -s)"
ARCH="$(uname -m)"

detect_cpu() {
    case "$OS" in
        Darwin)
            sysctl -n machdep.cpu.brand_string 2>/dev/null || sysctl -n hw.model 2>/dev/null || echo "Unknown"
            ;;
        Linux)
            grep -m1 'model name' /proc/cpuinfo 2>/dev/null | sed 's/model name\s*:\s*//' || echo "Unknown"
            ;;
        *) echo "Unknown" ;;
    esac
}

detect_gpu() {
    case "$OS" in
        Darwin)
            local chip
            chip="$(sysctl -n machdep.cpu.brand_string 2>/dev/null || echo '')"
            if [[ "$chip" == *"Apple"* ]]; then echo "$chip GPU"
            else echo "Unknown"; fi
            ;;
        Linux)
            if command -v nvidia-smi &>/dev/null; then
                nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null | head -1
            else echo "Unknown"; fi
            ;;
        *) echo "Unknown" ;;
    esac
}

detect_cpu_features() {
    case "$OS" in
        Darwin)
            local cores mem
            cores="$(sysctl -n hw.logicalcpu 2>/dev/null || echo '?')"
            mem="$(( $(sysctl -n hw.memsize 2>/dev/null || echo 0) / 1073741824 ))"
            echo "${cores}T, ${mem}GB RAM"
            ;;
        Linux)
            local threads mem
            threads="$(grep -c '^processor' /proc/cpuinfo 2>/dev/null || echo '?')"
            mem="$(awk '/MemTotal/ {printf "%.0f", $2/1048576}' /proc/meminfo 2>/dev/null || echo '?')"
            echo "${threads}T, ${mem}GB RAM"
            ;;
        *) echo "unknown" ;;
    esac
}

detect_gpu_backend() {
    case "$OS" in
        Darwin) echo "metal" ;;
        Linux)
            if command -v nvidia-smi &>/dev/null || command -v vulkaninfo &>/dev/null; then
                echo "vulkan"
            else echo "wgpu"; fi
            ;;
        *) echo "wgpu" ;;
    esac
}

detect_cpu_backend_features() {
    case "$OS" in
        Darwin) echo "ndarray,blas-accelerate" ;;
        *)
            if ldconfig -p 2>/dev/null | grep -q libopenblas; then
                echo "ndarray,openblas-system"
            else echo "ndarray"; fi
            ;;
    esac
}

slugify() {
    echo "$1" | tr '[:upper:]' '[:lower:]' | sed 's/[^a-z0-9._-]/_/g' | sed 's/__*/_/g; s/^_//; s/_$//'
}

# ─── Gather info ─────────────────────────────────────────────────────────────

CPU_NAME="$(detect_cpu)"
GPU_NAME="$(detect_gpu)"
CPU_INFO="$(detect_cpu_features)"
GPU_BACKEND="$(detect_gpu_backend)"
CPU_FEATURES="$(detect_cpu_backend_features)"

CPU_SLUG="$(slugify "$CPU_NAME")"
TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
DATE_HUMAN="$(date '+%Y-%m-%d %H:%M:%S')"

RESULTS_DIR="bench_results"
RUN_ID="${TIMESTAMP}_${CPU_SLUG}"
RUN_DIR="${RESULTS_DIR}/${RUN_ID}"
mkdir -p "$RUN_DIR"

echo "╔══════════════════════════════════════════════════════════════╗"
echo "║  OSF-RS — CPU vs GPU Benchmark                              ║"
echo "╚══════════════════════════════════════════════════════════════╝"
echo ""
echo "  Platform:   $OS $ARCH"
echo "  CPU:        $CPU_NAME ($CPU_INFO)"
echo "  GPU:        $GPU_NAME"
echo "  Weights:    $WEIGHTS"
echo "  Warmup:     $WARMUP    Runs: $RUNS"
echo "  Output:     $RUN_DIR/"
echo ""

# ─── Write metadata ──────────────────────────────────────────────────────────

cat > "$RUN_DIR/meta.json" <<EOF
{
  "timestamp": "$DATE_HUMAN",
  "os": "$OS",
  "arch": "$ARCH",
  "cpu": "$CPU_NAME",
  "cpu_info": "$CPU_INFO",
  "gpu": "$GPU_NAME",
  "gpu_backend": "$GPU_BACKEND",
  "warmup": $WARMUP,
  "runs": $RUNS,
  "run_id": "$RUN_ID"
}
EOF

# ─── Build & Run CPU ─────────────────────────────────────────────────────────

echo "━━━ Building CPU backend (features: $CPU_FEATURES) ━━━"
cargo build --release --example benchmark --features "$CPU_FEATURES" 2>&1 | tail -3
echo ""

CPU_JSON="$RUN_DIR/cpu.json"
CPU_LOG="$RUN_DIR/cpu.log"
CPU_OK=false

echo "━━━ Running CPU benchmark ━━━"
if cargo run --release --example benchmark --features "$CPU_FEATURES" -- \
    --weights "$WEIGHTS" --warmup "$WARMUP" --runs "$RUNS" --json \
    > "$CPU_JSON" 2>"$CPU_LOG"; then
    if python3 -c "import json,sys; json.load(open(sys.argv[1]))" "$CPU_JSON" 2>/dev/null; then
        CPU_OK=true
        echo "  ✓ CPU results → $CPU_JSON"
    else
        echo "  ⚠  CPU benchmark produced invalid JSON"
    fi
else
    echo "  ⚠  CPU benchmark failed (see $CPU_LOG)"
    tail -5 "$CPU_LOG"
fi
echo ""

# ─── Build & Run GPU ─────────────────────────────────────────────────────────

GPU_FEATURE="${GPU_BACKEND}"
GPU_JSON="$RUN_DIR/gpu.json"
GPU_LOG="$RUN_DIR/gpu.log"
GPU_OK=false

echo "━━━ Building GPU backend (features: $GPU_FEATURE) ━━━"
if cargo build --release --example benchmark --no-default-features --features "$GPU_FEATURE" 2>&1 | tail -3; then
    echo ""
    echo "━━━ Running GPU benchmark ━━━"
    if cargo run --release --example benchmark --no-default-features --features "$GPU_FEATURE" -- \
        --weights "$WEIGHTS" --warmup "$WARMUP" --runs "$RUNS" --json \
        > "$GPU_JSON" 2>"$GPU_LOG"; then
        if python3 -c "import json,sys; json.load(open(sys.argv[1]))" "$GPU_JSON" 2>/dev/null; then
            GPU_OK=true
            echo "  ✓ GPU results → $GPU_JSON"
        else
            echo "  ⚠  GPU benchmark produced invalid JSON"
        fi
    else
        echo "  ⚠  GPU benchmark failed (see $GPU_LOG)"
        tail -5 "$GPU_LOG"
    fi
else
    echo "  ⚠  GPU build failed — skipping GPU benchmark"
fi
echo ""

# ─── Generate Charts ─────────────────────────────────────────────────────────

echo "━━━ Generating charts ━━━"

if ! $CPU_OK; then
    echo "  ⚠  No valid CPU results — cannot generate charts."
    [ -f "$CPU_LOG" ] && cat "$CPU_LOG"
    exit 1
fi

python3 scripts/generate_figures.py "$RUN_DIR" "$GPU_OK" "$CPU_NAME" "$GPU_NAME" "$DATE_HUMAN"

# Copy latest charts to ./figures for README
mkdir -p figures
cp "$RUN_DIR"/*.png figures/ 2>/dev/null || true

echo ""
echo "━━━ Results saved to $RUN_DIR/ ━━━"
echo "━━━ Charts copied to figures/ ━━━"
echo ""
ls -la "$RUN_DIR/"
echo ""
echo "Done."