#!/bin/bash
# GPU Batch Size Optimization Script for RTX 3050 Ti (4GB VRAM)
#
# Tests optimal batch sizes for TFT, MAMBA-2, and Liquid models
# Generates BATCH_SIZE_OPTIMIZATION_REPORT.md with recommendations

set -e

echo "==================================="
echo "GPU Batch Size Optimization"
echo "==================================="
echo ""

# Check if nvidia-smi is available
if ! command -v nvidia-smi &> /dev/null; then
    echo "ERROR: nvidia-smi not found. This script requires NVIDIA GPU."
    exit 1
fi

# Display GPU info
echo "GPU Information:"
nvidia-smi --query-gpu=name,memory.total,driver_version --format=csv,noheader
echo ""

# Check CUDA availability
echo "Checking CUDA setup..."
if [ -d "/usr/local/cuda" ]; then
    echo "✓ CUDA found at /usr/local/cuda"
    nvcc --version | head -n 1
else
    echo "⚠ CUDA not found at /usr/local/cuda (will fall back to CPU)"
fi
echo ""

# Build the optimization tool in release mode
echo "Building optimization tool (release mode)..."
cargo build -p ml --example optimize_batch_sizes --release

if [ $? -ne 0 ]; then
    echo "ERROR: Failed to build optimization tool"
    exit 1
fi

echo "✓ Build complete"
echo ""

# Run the optimization
echo "Running batch size optimization..."
echo "This will test:"
echo "  - TFT: batch sizes [16, 32, 64, 128]"
echo "  - MAMBA-2: batch sizes [8, 16, 32]"
echo "  - Liquid: batch sizes [16, 32, 64]"
echo ""
echo "Estimated runtime: 3-5 minutes"
echo ""

# Monitor VRAM usage in background
echo "Starting VRAM monitor..."
(
    while true; do
        nvidia-smi --query-gpu=memory.used,memory.total --format=csv,noheader,nounits | \
        awk '{printf "VRAM: %d MB / %d MB (%.1f%%)\r", $1, $2, ($1/$2)*100}'
        sleep 1
    done
) &
MONITOR_PID=$!

# Run the optimization
cargo run -p ml --example optimize_batch_sizes --release

# Kill the monitor
kill $MONITOR_PID 2>/dev/null || true
echo ""

# Check if report was generated
if [ -f "BATCH_SIZE_OPTIMIZATION_REPORT.md" ]; then
    echo ""
    echo "==================================="
    echo "Optimization Complete!"
    echo "==================================="
    echo ""
    echo "Report generated: BATCH_SIZE_OPTIMIZATION_REPORT.md"
    echo ""

    # Extract recommendations
    echo "Recommended Batch Sizes:"
    grep -A 10 "## Optimization Summary" BATCH_SIZE_OPTIMIZATION_REPORT.md | \
    grep -E "^\| (TFT|MAMBA-2|Liquid)" | \
    awk -F'|' '{print "  " $2 " -> batch_size = " $3}' | \
    sed 's/  //' | sed 's/^ *//'

    echo ""
    echo "Next Steps:"
    echo "1. Review BATCH_SIZE_OPTIMIZATION_REPORT.md for detailed results"
    echo "2. Update model configurations with recommended batch sizes"
    echo "3. Test training with optimized batch sizes"
else
    echo ""
    echo "ERROR: Report not generated"
    exit 1
fi