#!/bin/bash
# Quick test: Train DQN for 10 epochs to verify .safetensors file is created

set -e

echo "🧪 DQN Training Test (10 epochs)"
echo "================================"
echo ""

# Check GPU
if ! nvidia-smi > /dev/null 2>&1; then
    echo "⚠️  GPU not available, using CPU"
    GPU_FLAG=""
else
    GPU_NAME=$(nvidia-smi --query-gpu=name --format=csv,noheader | head -1)
    echo "✅ GPU: $GPU_NAME"
    GPU_FLAG="--features cuda"
fi
echo ""

# Create test output directory
TEST_DIR="ml/trained_models/test"
rm -rf "$TEST_DIR"
mkdir -p "$TEST_DIR"

echo "📁 Output directory: $TEST_DIR"
echo ""

# Run short training test
echo "🏋️  Training DQN for 10 epochs..."
echo ""

cargo run -p ml --example train_dqn --release $GPU_FLAG -- \
    --epochs 10 \
    --batch-size 128 \
    --learning-rate 0.0001 \
    --checkpoint-frequency 5 \
    --output-dir "$TEST_DIR" \
    --verbose

EXIT_CODE=$?

echo ""
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"

if [ $EXIT_CODE -eq 0 ]; then
    echo "✅ Training completed successfully!"
    echo ""
    echo "📊 Output files:"
    ls -lh "$TEST_DIR"/*.safetensors 2>/dev/null || echo "⚠️  No .safetensors files found"
    echo ""

    # Count files
    FILE_COUNT=$(find "$TEST_DIR" -name "*.safetensors" -type f | wc -l)
    if [ $FILE_COUNT -gt 0 ]; then
        echo "✅ SUCCESS: $FILE_COUNT .safetensors files created!"
        echo ""
        echo "File details:"
        find "$TEST_DIR" -name "*.safetensors" -type f -exec ls -lh {} \; | while read -r line; do
            echo "  $line"
        done
    else
        echo "❌ FAILED: No .safetensors files created"
        exit 1
    fi
else
    echo "❌ Training failed with exit code: $EXIT_CODE"
    exit $EXIT_CODE
fi

echo ""
echo "🎉 Test passed! DQN trainer saves models correctly."
echo ""