#!/bin/bash
# Track VRAM during MAMBA-2 training
# Measures actual memory usage vs theoretical predictions

set -e

echo "========================================="
echo "MAMBA-2 VRAM Usage Analysis"
echo "========================================="
echo ""

# Test batch sizes
BATCH_SIZES=(32 64 96 128 144 160 180 200 220)

# Store results
declare -A VRAM_USAGE
declare -A TEST_STATUS

echo "Measuring VRAM usage for different batch sizes..."
echo "This will take approximately 10-15 minutes"
echo ""

for BS in "${BATCH_SIZES[@]}"; do
  echo "----------------------------------------"
  echo "Testing batch_size=$BS"
  echo "----------------------------------------"

  # Clear VRAM first
  if command -v nvidia-smi &> /dev/null; then
    nvidia-smi --gpu-reset-ecc-errors &> /dev/null || true
  fi
  sleep 2

  # Measure baseline VRAM
  BASELINE_VRAM=$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits | head -1)
  echo "  Baseline VRAM: ${BASELINE_VRAM}MB"

  # Start training in background
  timeout 90s cargo run -p ml --example hyperopt_mamba2_demo --release --features cuda -- \
    --parquet-file test_data/ES_FUT_180d.parquet \
    --trials 1 \
    --epochs 1 \
    --batch-size-min $BS \
    --batch-size-max $BS 2>&1 | grep -E "VRAM|Creating|Training|error|CUDA" &

  PID=$!

  # Wait for training to start
  sleep 15

  # Measure peak VRAM during training
  PEAK_VRAM=0
  for i in {1..10}; do
    if ps -p $PID > /dev/null 2>&1; then
      CURRENT_VRAM=$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits | head -1)
      if [ $CURRENT_VRAM -gt $PEAK_VRAM ]; then
        PEAK_VRAM=$CURRENT_VRAM
      fi
      sleep 2
    else
      break
    fi
  done

  # Kill training if still running
  kill $PID 2>/dev/null || true
  wait $PID 2>/dev/null || true

  # Calculate net usage
  NET_VRAM=$((PEAK_VRAM - BASELINE_VRAM))

  # Store results
  VRAM_USAGE[$BS]=$NET_VRAM

  if [ $NET_VRAM -gt 0 ]; then
    TEST_STATUS[$BS]="SUCCESS"
    echo "  Peak VRAM: ${PEAK_VRAM}MB"
    echo "  Net usage: ${NET_VRAM}MB"
    echo "  Status: ✓ Success"
  else
    TEST_STATUS[$BS]="FAILED"
    echo "  Status: ✗ Failed to measure"
  fi

  echo ""

  # Cool down between tests
  sleep 5
done

echo ""
echo "========================================="
echo "VRAM Usage Summary"
echo "========================================="
echo ""
printf "%-12s | %-12s | %-10s\n" "Batch Size" "VRAM Usage" "Status"
printf "%-12s-+-%-12s-+-%-10s\n" "------------" "------------" "----------"

for BS in "${BATCH_SIZES[@]}"; do
  VRAM=${VRAM_USAGE[$BS]:-"N/A"}
  STATUS=${TEST_STATUS[$BS]:-"UNKNOWN"}

  if [ "$STATUS" = "SUCCESS" ]; then
    STATUS_ICON="✓"
  else
    STATUS_ICON="✗"
  fi

  printf "%-12s | %-12s | %-10s\n" "$BS" "${VRAM}MB" "$STATUS_ICON $STATUS"
done

echo ""
echo "========================================="
echo "Analysis"
echo "========================================="
echo ""

# Calculate linear regression A + B * batch_size
# Using successful measurements only

VALID_POINTS=()
for BS in "${BATCH_SIZES[@]}"; do
  if [ "${TEST_STATUS[$BS]}" = "SUCCESS" ]; then
    VRAM=${VRAM_USAGE[$BS]}
    if [ $VRAM -gt 0 ]; then
      VALID_POINTS+=("$BS,$VRAM")
    fi
  fi
done

if [ ${#VALID_POINTS[@]} -ge 2 ]; then
  echo "Valid measurements: ${#VALID_POINTS[@]}"
  echo ""
  echo "Formula derivation: VRAM = A + B × batch_size"
  echo ""

  # Simple two-point calculation (first and last)
  FIRST_POINT=(${VALID_POINTS[0]//,/ })
  LAST_POINT=(${VALID_POINTS[-1]//,/ })

  BS1=${FIRST_POINT[0]}
  VRAM1=${FIRST_POINT[1]}
  BS2=${LAST_POINT[0]}
  VRAM2=${LAST_POINT[1]}

  # Calculate slope B = (VRAM2 - VRAM1) / (BS2 - BS1)
  B=$(echo "scale=2; ($VRAM2 - $VRAM1) / ($BS2 - $BS1)" | bc)

  # Calculate intercept A = VRAM1 - B * BS1
  A=$(echo "scale=2; $VRAM1 - $B * $BS1" | bc)

  echo "  Slope (B): ${B}MB per batch_size"
  echo "  Intercept (A): ${A}MB"
  echo ""
  echo "  New formula: VRAM = ${A}MB + ${B}MB × batch_size"
  echo ""

  # Calculate safe maximum for 16GB GPU (14.4GB with 10% margin)
  TARGET_VRAM=14400
  MAX_BATCH_SIZE=$(echo "scale=0; ($TARGET_VRAM - $A) / $B" | bc)

  echo "Safe maximum batch size for 16GB GPU (14.4GB with 10% margin):"
  echo "  Max batch_size: $MAX_BATCH_SIZE"
  echo ""

  # Compare with old formula
  echo "Comparison with old formula (predicted 13.2GB @ batch_size=144):"
  PREDICTED_OLD=13200
  ACTUAL_144=${VRAM_USAGE[144]:-"N/A"}

  if [ "$ACTUAL_144" != "N/A" ]; then
    DISCREPANCY=$((PREDICTED_OLD - ACTUAL_144))
    PERCENT_ERROR=$(echo "scale=1; 100 * $DISCREPANCY / $PREDICTED_OLD" | bc)
    echo "  Old formula predicted: 13.2GB"
    echo "  Actual measured @ BS=144: ${ACTUAL_144}MB"
    echo "  Discrepancy: ${DISCREPANCY}MB (${PERCENT_ERROR}% error)"
  fi
else
  echo "Insufficient valid measurements for analysis"
fi

echo ""
echo "========================================="
echo "Recommendations"
echo "========================================="
echo ""
echo "1. Update batch_size_max from 144 to $MAX_BATCH_SIZE"
echo "2. Verify formula: VRAM = ${A}MB + ${B}MB × batch_size"
echo "3. Test edge cases near maximum"
echo "4. Monitor for memory fragmentation"
echo ""