#!/bin/bash
# Mimalloc Allocator Performance Benchmark
# Compares training speed with and without mimalloc

set -e

echo "╔═══════════════════════════════════════════════════════════╗"
echo "║         Mimalloc Allocator Performance Benchmark         ║"
echo "╚═══════════════════════════════════════════════════════════╝"
echo ""

# Test configuration
PARQUET_FILE="test_data/ES_FUT_small.parquet"
EPOCHS=3
BATCH_SIZE=16  # Small batch to reduce total time

echo "Test Configuration:"
echo "  • Parquet file: $PARQUET_FILE"
echo "  • Epochs: $EPOCHS"
echo "  • Batch size: $BATCH_SIZE"
echo ""

# Build WITHOUT mimalloc (system allocator)
echo "═══════════════════════════════════════════════════════════"
echo "Building WITHOUT mimalloc (system allocator)..."
echo "═══════════════════════════════════════════════════════════"
cargo build --release -p ml --example train_tft_parquet --features cuda 2>&1 | grep -E "(Compiling|Finished)" | tail -5
echo ""

# Run benchmark WITHOUT mimalloc
echo "═══════════════════════════════════════════════════════════"
echo "Running TFT training WITHOUT mimalloc..."
echo "═══════════════════════════════════════════════════════════"
SYSTEM_START=$(date +%s)
timeout 180 ./target/release/examples/train_tft_parquet \
  --parquet-file "$PARQUET_FILE" \
  --epochs "$EPOCHS" \
  --batch-size "$BATCH_SIZE" \
  --use-gpu \
  2>&1 | tee /tmp/tft_system_allocator.log || true
SYSTEM_END=$(date +%s)
SYSTEM_DURATION=$((SYSTEM_END - SYSTEM_START))

echo ""
echo "System allocator duration: ${SYSTEM_DURATION}s"
echo ""

# Build WITH mimalloc
echo "═══════════════════════════════════════════════════════════"
echo "Building WITH mimalloc allocator..."
echo "═══════════════════════════════════════════════════════════"
cargo build --release -p ml --example train_tft_parquet --features "cuda,mimalloc-allocator" 2>&1 | grep -E "(Compiling|Finished)" | tail -5
echo ""

# Run benchmark WITH mimalloc
echo "═══════════════════════════════════════════════════════════"
echo "Running TFT training WITH mimalloc..."
echo "═══════════════════════════════════════════════════════════"
MIMALLOC_START=$(date +%s)
timeout 180 ./target/release/examples/train_tft_parquet \
  --parquet-file "$PARQUET_FILE" \
  --epochs "$EPOCHS" \
  --batch-size "$BATCH_SIZE" \
  --use-gpu \
  2>&1 | tee /tmp/tft_mimalloc.log || true
MIMALLOC_END=$(date +%s)
MIMALLOC_DURATION=$((MIMALLOC_END - MIMALLOC_START))

echo ""
echo "Mimalloc duration: ${MIMALLOC_DURATION}s"
echo ""

# Calculate improvement
if [ $SYSTEM_DURATION -gt 0 ]; then
    SPEEDUP=$(echo "scale=2; ($SYSTEM_DURATION - $MIMALLOC_DURATION) / $SYSTEM_DURATION * 100" | bc)
    RATIO=$(echo "scale=2; $SYSTEM_DURATION / $MIMALLOC_DURATION" | bc)

    echo "╔═══════════════════════════════════════════════════════════╗"
    echo "║                  Performance Results                     ║"
    echo "╚═══════════════════════════════════════════════════════════╝"
    echo ""
    echo "System Allocator Duration: ${SYSTEM_DURATION}s"
    echo "Mimalloc Duration:         ${MIMALLOC_DURATION}s"
    echo ""
    echo "Performance Improvement:   ${SPEEDUP}%"
    echo "Speedup Ratio:             ${RATIO}x"
    echo ""

    if (( $(echo "$SPEEDUP >= 10" | bc -l) )); then
        echo "✅ EXCELLENT: Mimalloc provides ${SPEEDUP}% speedup (≥10% target)"
    elif (( $(echo "$SPEEDUP >= 5" | bc -l) )); then
        echo "✅ GOOD: Mimalloc provides ${SPEEDUP}% speedup (≥5%)"
    else
        echo "⚠️  MODEST: Mimalloc provides ${SPEEDUP}% speedup (<5%)"
        echo "   This is expected for GPU-bound workloads."
    fi
    echo ""
else
    echo "❌ Error: System allocator test failed or took 0 seconds"
fi

# Extract training time from logs
echo "═══════════════════════════════════════════════════════════"
echo "Detailed Training Metrics"
echo "═══════════════════════════════════════════════════════════"
echo ""
echo "System Allocator:"
grep -E "(Training duration|Using|allocator)" /tmp/tft_system_allocator.log | head -5 || echo "No logs found"
echo ""
echo "Mimalloc Allocator:"
grep -E "(Training duration|Using mimalloc|allocator)" /tmp/tft_mimalloc.log | head -5 || echo "No logs found"
echo ""

echo "╔═══════════════════════════════════════════════════════════╗"
echo "║              Benchmark Complete                           ║"
echo "╚═══════════════════════════════════════════════════════════╝"
echo ""
echo "Full logs saved to:"
echo "  • /tmp/tft_system_allocator.log"
echo "  • /tmp/tft_mimalloc.log"