#!/bin/bash
# Train all 4 ML models (DQN, PPO, MAMBA-2, TFT) with REAL TRAINERS (not benchmarks)
# Saves .safetensors model files to disk

set -e

echo "🚀 Full Model Training - All 4 Models (REAL TRAINERS)"
echo "======================================================"
echo ""
echo "Models: DQN, PPO, MAMBA-2, TFT"
echo "Output: .safetensors files saved to ml/trained_models/"
echo "Epochs: 500 per model (production-scale training)"
echo ""

# Check GPU
if ! nvidia-smi > /dev/null 2>&1; then
    echo "❌ GPU not available"
    exit 1
fi

GPU_NAME=$(nvidia-smi --query-gpu=name --format=csv,noheader | head -1)
GPU_MEMORY=$(nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits | head -1)
echo "✅ GPU: $GPU_NAME ($GPU_MEMORY MB)"
echo ""

# Create output directory for trained models
MODEL_DIR="ml/trained_models"
mkdir -p "$MODEL_DIR"
echo "📁 Model output directory: $MODEL_DIR"
echo ""

# Training configuration
EPOCHS=500
LEARNING_RATE=0.0001
BATCH_SIZE_DQN=128      # DQN optimal
BATCH_SIZE_PPO=64       # PPO optimal
BATCH_SIZE_MAMBA=8      # MAMBA-2 memory-constrained (4GB VRAM)
BATCH_SIZE_TFT=32       # TFT memory-constrained

echo "⚙️  Training Configuration:"
echo "   • Epochs: $EPOCHS"
echo "   • Learning rate: $LEARNING_RATE"
echo "   • DQN batch size: $BATCH_SIZE_DQN"
echo "   • PPO batch size: $BATCH_SIZE_PPO"
echo "   • MAMBA-2 batch size: $BATCH_SIZE_MAMBA"
echo "   • TFT batch size: $BATCH_SIZE_TFT"
echo ""

# Training results log
RESULTS_FILE="$MODEL_DIR/training_results_$(date +%Y%m%d_%H%M%S).json"
echo "{" > "$RESULTS_FILE"
echo "  \"training_start\": \"$(date -Iseconds)\"," >> "$RESULTS_FILE"
echo "  \"configuration\": {" >> "$RESULTS_FILE"
echo "    \"epochs\": $EPOCHS," >> "$RESULTS_FILE"
echo "    \"learning_rate\": $LEARNING_RATE" >> "$RESULTS_FILE"
echo "  }," >> "$RESULTS_FILE"
echo "  \"models\": {" >> "$RESULTS_FILE"

# Function to train a model
train_model() {
    local MODEL_NAME=$1
    local MODEL_TYPE=$2  # dqn, ppo, mamba2, tft
    local BATCH_SIZE=$3

    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo "📊 Training $MODEL_NAME..."
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo ""

    local START_TIME=$(date +%s)
    local OUTPUT_DIR="$MODEL_DIR"

    # Run the REAL trainer (not benchmark)
    echo "   Training $MODEL_NAME with $EPOCHS epochs..."
    echo "   Batch size: $BATCH_SIZE"
    echo "   Output directory: $OUTPUT_DIR"
    echo ""

    # Use the proper training examples we just created
    cargo run -p ml --example "train_${MODEL_TYPE}" --release --features cuda -- \
        --epochs "$EPOCHS" \
        --learning-rate "$LEARNING_RATE" \
        --batch-size "$BATCH_SIZE" \
        --output-dir "$OUTPUT_DIR" \
        --verbose \
        2>&1 | tee "$MODEL_DIR/${MODEL_TYPE}_training.log"

    local EXIT_CODE=$?
    local END_TIME=$(date +%s)
    local DURATION=$((END_TIME - START_TIME))
    local HOURS=$((DURATION / 3600))
    local MINUTES=$(((DURATION % 3600) / 60))
    local SECONDS=$((DURATION % 60))

    if [ $EXIT_CODE -eq 0 ]; then
        echo ""
        echo "✅ $MODEL_NAME training complete!"
        echo "   Duration: ${HOURS}h ${MINUTES}m ${SECONDS}s"

        # Find the saved model files
        local MODEL_FILES=$(find "$OUTPUT_DIR" -name "${MODEL_TYPE}*.safetensors" -type f -mmin -$((DURATION / 60 + 5)) | head -5)
        if [ -n "$MODEL_FILES" ]; then
            echo "   Models saved:"
            echo "$MODEL_FILES" | while read -r file; do
                local SIZE=$(du -h "$file" | cut -f1)
                echo "     • $(basename "$file") ($SIZE)"
            done
        else
            echo "   ⚠️  Warning: No .safetensors files found (check logs)"
        fi
        echo ""

        # Record success
        echo "    \"$MODEL_TYPE\": {" >> "$RESULTS_FILE"
        echo "      \"model_name\": \"$MODEL_NAME\"," >> "$RESULTS_FILE"
        echo "      \"epochs\": $EPOCHS," >> "$RESULTS_FILE"
        echo "      \"batch_size\": $BATCH_SIZE," >> "$RESULTS_FILE"
        echo "      \"duration_seconds\": $DURATION," >> "$RESULTS_FILE"
        echo "      \"status\": \"success\"," >> "$RESULTS_FILE"
        echo "      \"log_file\": \"$MODEL_DIR/${MODEL_TYPE}_training.log\"" >> "$RESULTS_FILE"
        echo "    }," >> "$RESULTS_FILE"

        return 0
    else
        echo ""
        echo "❌ $MODEL_NAME training FAILED (exit code: $EXIT_CODE)"
        echo "   Duration: ${HOURS}h ${MINUTES}m ${SECONDS}s"
        echo "   Check logs: $MODEL_DIR/${MODEL_TYPE}_training.log"
        echo ""

        # Record failure
        echo "    \"$MODEL_TYPE\": {" >> "$RESULTS_FILE"
        echo "      \"model_name\": \"$MODEL_NAME\"," >> "$RESULTS_FILE"
        echo "      \"epochs\": $EPOCHS," >> "$RESULTS_FILE"
        echo "      \"batch_size\": $BATCH_SIZE," >> "$RESULTS_FILE"
        echo "      \"duration_seconds\": $DURATION," >> "$RESULTS_FILE"
        echo "      \"status\": \"failed\"," >> "$RESULTS_FILE"
        echo "      \"exit_code\": $EXIT_CODE," >> "$RESULTS_FILE"
        echo "      \"log_file\": \"$MODEL_DIR/${MODEL_TYPE}_training.log\"" >> "$RESULTS_FILE"
        echo "    }," >> "$RESULTS_FILE"

        return $EXIT_CODE
    fi
}

# Train all models sequentially
echo "🏋️  Starting training pipeline..."
echo ""

FAILED_COUNT=0

echo "1/4 Training DQN..."
if ! train_model "DQN (Deep Q-Network)" "dqn" "$BATCH_SIZE_DQN"; then
    FAILED_COUNT=$((FAILED_COUNT + 1))
fi

echo "2/4 Training PPO..."
if ! train_model "PPO (Proximal Policy Optimization)" "ppo" "$BATCH_SIZE_PPO"; then
    FAILED_COUNT=$((FAILED_COUNT + 1))
fi

echo "3/4 Training MAMBA-2..."
if ! train_model "MAMBA-2 (State Space Model)" "mamba2" "$BATCH_SIZE_MAMBA"; then
    FAILED_COUNT=$((FAILED_COUNT + 1))
fi

echo "4/4 Training TFT..."
if ! train_model "TFT (Temporal Fusion Transformer)" "tft" "$BATCH_SIZE_TFT"; then
    FAILED_COUNT=$((FAILED_COUNT + 1))
fi

# Close JSON
echo "    \"_end\": null" >> "$RESULTS_FILE"
echo "  }," >> "$RESULTS_FILE"
echo "  \"training_end\": \"$(date -Iseconds)\"," >> "$RESULTS_FILE"
echo "  \"failed_count\": $FAILED_COUNT" >> "$RESULTS_FILE"
echo "}" >> "$RESULTS_FILE"

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

if [ $FAILED_COUNT -eq 0 ]; then
    echo "🎉 All Models Trained Successfully!"
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    EXIT_STATUS=0
else
    echo "⚠️  Training Complete with $FAILED_COUNT Failures"
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    EXIT_STATUS=1
fi

echo ""
echo "📊 Training Summary:"
cat "$RESULTS_FILE" | grep -E "(model_name|duration_seconds|status)" | head -20
echo ""
echo "📁 Trained Models:"
find "$MODEL_DIR" -name "*.safetensors" -type f -mmin -300 | sort | while read -r file; do
    local SIZE=$(du -h "$file" | cut -f1)
    echo "  • $(basename "$file") ($SIZE)"
done
echo ""
echo "📄 Results saved to: $RESULTS_FILE"
echo ""
echo "✨ Next Steps:"
echo "   1. Load trained models in adaptive strategy tests"
echo "   2. Validate trading performance with real market data"
echo "   3. Benchmark Sharpe ratio and risk metrics"
echo "   4. Deploy to production for live trading"
echo ""

exit $EXIT_STATUS