#!/bin/bash
# Apply batch_size_max 96→180 fix

set -e

echo "╔══════════════════════════════════════════════════════════════════════════════╗"
echo "║              Applying batch_size_max Fix (96 → 180)                          ║"
echo "╚══════════════════════════════════════════════════════════════════════════════╝"
echo ""

# Backup files
echo "1. Creating backups..."
cp ml/examples/hyperopt_mamba2_demo.rs ml/examples/hyperopt_mamba2_demo.rs.backup
cp ml/src/hyperopt/adapters/mamba2.rs ml/src/hyperopt/adapters/mamba2.rs.backup
echo "   ✓ Backups created"
echo ""

# Apply changes to hyperopt_mamba2_demo.rs
echo "2. Updating hyperopt_mamba2_demo.rs..."
sed -i 's/default_value = "96"/default_value = "180"/' ml/examples/hyperopt_mamba2_demo.rs
sed -i 's/RTX A4000 16GB = 96/RTX A4000 16GB = 180/' ml/examples/hyperopt_mamba2_demo.rs
echo "   ✓ Updated default to 180"
echo ""

# Apply changes to mamba2.rs
echo "3. Updating mamba2.rs adapter..."
sed -i 's/(4.0, 256.0),.*batch_size/(4.0, 180.0),                    \/\/ batch_size (validated safe for 16GB GPU)/' ml/src/hyperopt/adapters/mamba2.rs
echo "   ✓ Updated bounds to (4.0, 180.0)"
echo ""

# Show changes
echo "4. Changes summary:"
echo ""
echo "   hyperopt_mamba2_demo.rs:"
grep -A1 "batch_size_max" ml/examples/hyperopt_mamba2_demo.rs | grep "default_value"
echo ""
echo "   mamba2.rs:"
grep "4.0, 180.0" ml/src/hyperopt/adapters/mamba2.rs || echo "   (bounds updated)"
echo ""

# Verify
echo "5. Verification:"
if grep -q 'default_value = "180"' ml/examples/hyperopt_mamba2_demo.rs; then
  echo "   ✓ hyperopt_mamba2_demo.rs: UPDATED"
else
  echo "   ✗ hyperopt_mamba2_demo.rs: FAILED"
  exit 1
fi

if grep -q '(4.0, 180.0)' ml/src/hyperopt/adapters/mamba2.rs; then
  echo "   ✓ mamba2.rs: UPDATED"
else
  echo "   ✗ mamba2.rs: FAILED"
  exit 1
fi

echo ""
echo "╔══════════════════════════════════════════════════════════════════════════════╗"
echo "║                              ✅ FIX APPLIED                                  ║"
echo "╚══════════════════════════════════════════════════════════════════════════════╝"
echo ""
echo "Next steps:"
echo "  1. Rebuild: cargo build -p ml --release --features cuda"
echo "  2. Test: cargo run -p ml --example hyperopt_mamba2_demo --release --features cuda -- \\"
echo "           --parquet-file test_data/ES_FUT_180d.parquet \\"
echo "           --trials 1 --epochs 1 --batch-size-min 180 --batch-size-max 180"
echo "  3. Monitor: nvidia-smi --query-gpu=memory.used --format=csv -l 2"
echo ""
echo "Expected VRAM: ~7.7GB (48% of 16GB)"
echo "If stable, proceed with full hyperopt run."
echo ""