#!/bin/bash
# DQN Production Training Script - Wave 1/2 Complete + Trial #68 Hyperopt
# 
# Status: ✅ PRODUCTION READY
# Generated: 2025-11-04
# 
# This script consolidates:
# - Wave 1: Huber loss, HOLD penalty, Double DQN, gradient clipping
# - Wave 2: Validation system, target network optimization, replay buffer
# - Trial #68: Best hyperparameters from 116 trials (objective: 0.0006354887)

set -e  # Exit on any error

echo "════════════════════════════════════════════════════════════════"
echo "🚀 DQN Production Training v2.0"
echo "   Wave 1/2 Complete + Trial #68 Hyperopt Optimized"
echo "════════════════════════════════════════════════════════════════"
echo ""

# ─────────────────────────────────────────────────────────────────────
# Configuration
# ─────────────────────────────────────────────────────────────────────

MODEL_NAME="dqn_v2_production_$(date +%Y%m%d_%H%M%S)"
EPOCHS=500
DATA_FILE="test_data/ES_FUT_180d.parquet"
OUTPUT_DIR="ml/trained_models/${MODEL_NAME}"
LOG_FILE="/tmp/${MODEL_NAME}.log"

# Trial #68 Hyperparameters (Best of 116 trials)
LEARNING_RATE=0.00055        # 5.5x higher than conservative default
BATCH_SIZE=230               # 7.2x higher than old hyperopt (max GPU capacity)
GAMMA=0.99                   # Long-term reward focus (top 5 avg=0.9887)
EPSILON_START=1.0            # Full exploration initially
EPSILON_END=0.01             # Minimal final exploration
EPSILON_DECAY=0.99           # Faster exploration→exploitation transition
BUFFER_SIZE=1000000          # 9.6x higher than old hyperopt (top 3 trials all used max)
MIN_REPLAY_SIZE=2000         # 2x batch_size

# Wave 1 Features
USE_DOUBLE_DQN=true          # Reduce Q-value overestimation
USE_HUBER_LOSS=true          # Robust to outliers (Q-values: -87K to +142K)
HUBER_DELTA=1.0              # Standard threshold
GRADIENT_CLIP_NORM=1.0       # Prevent gradient explosions
HOLD_PENALTY_WEIGHT=0.01     # 1% penalty per 1% excess movement
MOVEMENT_THRESHOLD=0.02      # 2% deadzone before penalty applies

# Wave 2 Features
TARGET_UPDATE_FREQ=500       # 2x faster than old default (1000)
VALIDATION_SPLIT=0.2         # 20% holdout set
VALIDATION_PATIENCE=5        # Early stop after 5 epochs val loss increase
MIN_EPOCHS_BEFORE_STOPPING=50  # Prevent premature stopping

# Training Configuration
CHECKPOINT_FREQ=10           # Save checkpoint every 10 epochs
VALIDATION_LOG_FREQ=10       # Log validation metrics every 10 epochs

echo "📋 Configuration Summary"
echo "────────────────────────────────────────────────────────────────"
echo "Model:       ${MODEL_NAME}"
echo "Epochs:      ${EPOCHS}"
echo "Data:        ${DATA_FILE}"
echo "Output:      ${OUTPUT_DIR}"
echo "Log:         ${LOG_FILE}"
echo ""
echo "Trial #68 Hyperparameters:"
echo "  • Learning Rate:   ${LEARNING_RATE} (was 0.0001, 5.5x improvement)"
echo "  • Batch Size:      ${BATCH_SIZE} (was 32, 7.2x improvement)"
echo "  • Gamma:           ${GAMMA} (was 0.9626, long-term focus)"
echo "  • Epsilon Decay:   ${EPSILON_DECAY} (was 0.995, faster convergence)"
echo "  • Buffer Size:     ${BUFFER_SIZE} (was 104,346, 9.6x improvement)"
echo ""
echo "Wave 1 Features (Address 99.4% HOLD + outliers + gradient explosions):"
echo "  • Double DQN:      ${USE_DOUBLE_DQN}"
echo "  • Huber Loss:      ${USE_HUBER_LOSS} (delta=${HUBER_DELTA})"
echo "  • Gradient Clip:   ${GRADIENT_CLIP_NORM}"
echo "  • HOLD Penalty:    ${HOLD_PENALTY_WEIGHT} (threshold=${MOVEMENT_THRESHOLD})"
echo ""
echo "Wave 2 Features (Faster convergence + validation):"
echo "  • Target Update:   ${TARGET_UPDATE_FREQ} (was 1000, 2x faster)"
echo "  • Validation Split: ${VALIDATION_SPLIT}"
echo "  • Val Patience:    ${VALIDATION_PATIENCE} epochs"
echo ""

# ─────────────────────────────────────────────────────────────────────
# Pre-flight Checks
# ─────────────────────────────────────────────────────────────────────

echo "🔍 Pre-flight Checks"
echo "────────────────────────────────────────────────────────────────"

# Check data file exists
if [ ! -f "${DATA_FILE}" ]; then
    echo "❌ ERROR: Data file not found: ${DATA_FILE}"
    echo "   Run: python3 scripts/python/data/download_es_90d_multi_contract.py"
    exit 1
fi
echo "✅ Data file exists: ${DATA_FILE}"

# Check GPU availability
if ! command -v nvidia-smi &> /dev/null; then
    echo "⚠️  WARNING: nvidia-smi not found. GPU may not be available."
else
    echo "✅ GPU detected:"
    nvidia-smi --query-gpu=name,memory.total,memory.free --format=csv,noheader | head -n 1
fi

# Check sufficient disk space (need ~5GB for model + logs)
AVAILABLE_SPACE=$(df -BG . | tail -1 | awk '{print $4}' | sed 's/G//')
if [ "$AVAILABLE_SPACE" -lt 5 ]; then
    echo "⚠️  WARNING: Low disk space (${AVAILABLE_SPACE}GB available, recommend 5GB+)"
fi
echo "✅ Disk space: ${AVAILABLE_SPACE}GB available"

# Create output directory
mkdir -p "${OUTPUT_DIR}"
echo "✅ Output directory created: ${OUTPUT_DIR}"
echo ""

# ─────────────────────────────────────────────────────────────────────
# Training
# ─────────────────────────────────────────────────────────────────────

echo "🏋️  Starting Production Training"
echo "────────────────────────────────────────────────────────────────"
echo "Estimated Duration: 60 minutes (500 epochs @ 7-8s/epoch)"
echo "Expected Cost: $0.25 (Runpod RTX A4000) or Free (local RTX 3050 Ti)"
echo ""

START_TIME=$(date +%s)

cargo run -p ml --example train_dqn --release --features cuda -- \
  --epochs ${EPOCHS} \
  --learning-rate ${LEARNING_RATE} \
  --batch-size ${BATCH_SIZE} \
  --gamma ${GAMMA} \
  --epsilon-start ${EPSILON_START} \
  --epsilon-end ${EPSILON_END} \
  --epsilon-decay ${EPSILON_DECAY} \
  --buffer-size ${BUFFER_SIZE} \
  --min-replay-size ${MIN_REPLAY_SIZE} \
  --use-double-dqn=${USE_DOUBLE_DQN} \
  --use-huber-loss=${USE_HUBER_LOSS} \
  --huber-delta ${HUBER_DELTA} \
  --gradient-clip-norm ${GRADIENT_CLIP_NORM} \
  --hold-penalty-weight ${HOLD_PENALTY_WEIGHT} \
  --movement-threshold ${MOVEMENT_THRESHOLD} \
  --validation-split ${VALIDATION_SPLIT} \
  --validation-patience ${VALIDATION_PATIENCE} \
  --validation-log-frequency ${VALIDATION_LOG_FREQ} \
  --min-epochs-before-stopping ${MIN_EPOCHS_BEFORE_STOPPING} \
  --checkpoint-frequency ${CHECKPOINT_FREQ} \
  --output-dir "${OUTPUT_DIR}" \
  --parquet-file "${DATA_FILE}" \
  2>&1 | tee "${LOG_FILE}"

EXIT_CODE=$?
END_TIME=$(date +%s)
DURATION=$((END_TIME - START_TIME))
DURATION_MIN=$((DURATION / 60))

echo ""
echo "════════════════════════════════════════════════════════════════"

if [ $EXIT_CODE -eq 0 ]; then
    echo "✅ Training Complete!"
    echo "   Duration: ${DURATION_MIN} minutes"
    echo "   Output:   ${OUTPUT_DIR}"
    echo "   Log:      ${LOG_FILE}"
else
    echo "❌ Training Failed (exit code: ${EXIT_CODE})"
    echo "   Check log: ${LOG_FILE}"
    exit $EXIT_CODE
fi

echo "════════════════════════════════════════════════════════════════"
echo ""

# ─────────────────────────────────────────────────────────────────────
# Post-Training Analysis
# ─────────────────────────────────────────────────────────────────────

echo "📊 Post-Training Analysis"
echo "────────────────────────────────────────────────────────────────"

# Count checkpoints
CHECKPOINT_COUNT=$(ls -1 "${OUTPUT_DIR}"/*.safetensors 2>/dev/null | wc -l)
echo "Checkpoints saved: ${CHECKPOINT_COUNT}"

# Find best model
if [ -f "${OUTPUT_DIR}/dqn_best_model.safetensors" ]; then
    BEST_MODEL="${OUTPUT_DIR}/dqn_best_model.safetensors"
    BEST_SIZE=$(du -h "${BEST_MODEL}" | cut -f1)
    echo "Best model: ${BEST_MODEL} (${BEST_SIZE})"
else
    echo "⚠️  WARNING: Best model not found (early stopping may have failed)"
fi

# Extract final metrics from log (if available)
if grep -q "Epoch ${EPOCHS}:" "${LOG_FILE}"; then
    echo ""
    echo "Final Epoch Metrics:"
    grep "Epoch ${EPOCHS}:" "${LOG_FILE}" | tail -1
fi

echo ""
echo "────────────────────────────────────────────────────────────────"
echo "🔍 Next Steps"
echo "────────────────────────────────────────────────────────────────"
echo ""
echo "1. Backtest on Unseen Data:"
echo "   cargo run -p ml --example backtest_dqn --release --features cuda -- \\"
echo "     --model-path \"${OUTPUT_DIR}/dqn_best_model.safetensors\" \\"
echo "     --data-file test_data/ES_FUT_unseen.parquet \\"
echo "     --output-json /tmp/${MODEL_NAME}_backtest.json"
echo ""
echo "2. Compare to Trial #35 Baseline:"
echo "   python3 scripts/python/compare_backtest_results.py \\"
echo "     /tmp/${MODEL_NAME}_backtest.json \\"
echo "     /tmp/dqn_trial35_backtest_results.json"
echo ""
echo "3. Deploy to Production (if metrics pass):"
echo "   - Sharpe Ratio > 1.5 ✓"
echo "   - Win Rate > 50% ✓"
echo "   - HOLD % < 70% ✓"
echo "   - Max Drawdown < 20% ✓"
echo ""
echo "════════════════════════════════════════════════════════════════"
echo "🎉 DQN Production Training Complete!"
echo "════════════════════════════════════════════════════════════════"