#!/bin/bash
#
# DQN Hyperparameter Optimization Dry-Run Script
#
# Purpose: Quick validation of hyperopt setup before committing to full 100-trial campaign
# Duration: 5-10 minutes
# Cost: $0.02-$0.04 (RTX A4000)
#
# This script:
# 1. Runs 5 hyperopt trials with 10 epochs each
# 2. Captures all training logs (loss, action distribution, gradient norms)
# 3. Validates that Wave 11 bug fixes are working correctly
# 4. Provides pass/fail report with actionable next steps
#
# Usage:
#   ./scripts/hyperopt_dqn_dryrun.sh

set -e  # Exit on error
set -o pipefail  # Exit if any command in pipeline fails

# ==================== CONFIGURATION ====================

TRIALS=5
EPOCHS=10
PARQUET_FILE="test_data/ES_FUT_180d.parquet"
LOG_DIR="/tmp/dqn_hyperopt_logs"
TIMESTAMP=$(date +%Y%m%d_%H%M%S)
LOG_FILE="${LOG_DIR}/dqn_hyperopt_dryrun_${TIMESTAMP}.log"

# Create log directory
mkdir -p "$LOG_DIR"

# ==================== BANNER ====================

echo "========================================"
echo "DQN Hyperopt Dry-Run"
echo "========================================"
echo "Configuration:"
echo "  Trials: $TRIALS"
echo "  Epochs per trial: $EPOCHS"
echo "  Parquet file: $PARQUET_FILE"
echo "  Log file: $LOG_FILE"
echo ""
echo "Expected duration: 5-10 minutes"
echo "Expected cost: \$0.02-\$0.04 (RTX A4000)"
echo ""
echo "Wave 11 Validations:"
echo "  ✓ Gradient clipping active (max_norm=10.0)"
echo "  ✓ Portfolio tracking operational"
echo "  ✓ HOLD penalty enabled (0.01 weight)"
echo "  ✓ Close price extraction accurate"
echo "========================================"
echo ""

# ==================== PRE-FLIGHT CHECKS ====================

echo "[1/4] Pre-flight checks..."

# Check if parquet file exists
if [ ! -f "$PARQUET_FILE" ]; then
    echo "❌ ERROR: Parquet file not found: $PARQUET_FILE"
    exit 1
fi
echo "  ✓ Parquet file exists: $PARQUET_FILE"

# Check if CUDA is available
if command -v nvidia-smi &> /dev/null; then
    echo "  ✓ CUDA available (nvidia-smi found)"
    nvidia-smi --query-gpu=name,memory.total,memory.free --format=csv,noheader | head -1
else
    echo "  ⚠️  CUDA not available (will use CPU - slower)"
fi

echo ""

# ==================== RUN HYPEROPT ====================

echo "[2/4] Running hyperopt (this will take 5-10 minutes)..."
echo ""

# Run hyperopt and capture output to both file and stdout
cargo run --release -p ml --example hyperopt_dqn_demo --features cuda -- \
    --parquet-file "$PARQUET_FILE" \
    --trials $TRIALS \
    --epochs $EPOCHS \
    2>&1 | tee "$LOG_FILE"

EXIT_CODE=$?

if [ $EXIT_CODE -ne 0 ]; then
    echo ""
    echo "❌ HYPEROPT FAILED (exit code: $EXIT_CODE)"
    echo "   Check logs: $LOG_FILE"
    exit $EXIT_CODE
fi

echo ""

# ==================== VALIDATION CHECKS ====================

echo "[3/4] Running validation checks..."
echo ""

# Check for errors/panics
ERROR_COUNT=$(grep -i -c "error\|panic\|FAILED\|thread.*panicked" "$LOG_FILE" || true)
echo "  Errors/panics: $ERROR_COUNT"

# Check for gradient warnings (should be 0 with gradient clipping)
GRAD_WARNINGS=$(grep -c "Large gradient\|gradient explosion\|Q-VALUE DIVERGENCE" "$LOG_FILE" || true)
echo "  Gradient warnings: $GRAD_WARNINGS"

# Check for Q-value explosions
QVALUE_EXPLOSIONS=$(grep -c "Q-value explosion" "$LOG_FILE" || true)
echo "  Q-value explosions: $QVALUE_EXPLOSIONS"

# Check for action diversity (at least one trial should have varied actions)
ACTION_DIST_COUNT=$(grep -c "Action Distribution" "$LOG_FILE" || true)
echo "  Action distribution logs: $ACTION_DIST_COUNT"

# Extract last trial's action distribution
echo ""
echo "  Last trial action distribution:"
if grep -q "Action Distribution" "$LOG_FILE"; then
    grep "Action Distribution" "$LOG_FILE" | tail -1 | sed 's/^/    /'
else
    echo "    ⚠️  No action distribution logs found"
fi

# Check for HOLD penalty usage (we use hold_penalty_weight in config)
HOLD_PENALTY_CONFIG=$(grep -c "hold_penalty_weight: 0.01" "$LOG_FILE" || true)
echo ""
echo "  HOLD penalty weight in config: $HOLD_PENALTY_CONFIG occurrences"

# Check for portfolio tracking (feature vector should include portfolio state)
PORTFOLIO_FEATURES=$(grep -c "portfolio" "$LOG_FILE" || true)
echo "  Portfolio tracking references: $PORTFOLIO_FEATURES"

# Check for gradient clipping activation
GRAD_CLIP_COUNT=$(grep -c "gradient_clip_norm: Some(10.0)" "$LOG_FILE" || true)
echo "  Gradient clipping enabled: $GRAD_CLIP_COUNT occurrences"

echo ""

# ==================== RESULTS SUMMARY ====================

echo "[4/4] Validation summary"
echo "========================================"

# Count passes and failures
CHECKS_PASSED=0
CHECKS_FAILED=0

echo ""
echo "Critical Checks:"
echo "----------------"

# Check 1: Training completed without errors
if [ $ERROR_COUNT -eq 0 ]; then
    echo "✅ Training completed without errors"
    CHECKS_PASSED=$((CHECKS_PASSED + 1))
else
    echo "❌ Training had $ERROR_COUNT errors/panics"
    CHECKS_FAILED=$((CHECKS_FAILED + 1))
fi

# Check 2: Gradient clipping working (no explosions)
if [ $GRAD_WARNINGS -eq 0 ] && [ $QVALUE_EXPLOSIONS -eq 0 ]; then
    echo "✅ Gradient clipping working (no explosions/warnings)"
    CHECKS_PASSED=$((CHECKS_PASSED + 1))
else
    echo "❌ Gradient issues detected ($GRAD_WARNINGS warnings, $QVALUE_EXPLOSIONS explosions)"
    CHECKS_FAILED=$((CHECKS_FAILED + 1))
fi

# Check 3: Action distribution logged
if [ $ACTION_DIST_COUNT -gt 0 ]; then
    echo "✅ Action distribution logged ($ACTION_DIST_COUNT times)"
    CHECKS_PASSED=$((CHECKS_PASSED + 1))
else
    echo "❌ No action distribution logs found"
    CHECKS_FAILED=$((CHECKS_FAILED + 1))
fi

# Check 4: Gradient clipping config present
if [ $GRAD_CLIP_COUNT -gt 0 ]; then
    echo "✅ Gradient clipping enabled in config"
    CHECKS_PASSED=$((CHECKS_PASSED + 1))
else
    echo "⚠️  Gradient clipping config not found in logs"
    CHECKS_FAILED=$((CHECKS_FAILED + 1))
fi

echo ""
echo "Optional Checks:"
echo "----------------"

# Check 5: HOLD penalty weight configured
if [ $HOLD_PENALTY_CONFIG -gt 0 ]; then
    echo "✅ HOLD penalty weight configured (0.01)"
else
    echo "⚠️  HOLD penalty weight not found in logs (may still be working)"
fi

# Check 6: Portfolio tracking present
if [ $PORTFOLIO_FEATURES -gt 0 ]; then
    echo "✅ Portfolio tracking references found"
else
    echo "⚠️  Portfolio tracking references not found (may still be working)"
fi

echo ""
echo "========================================"
echo "Overall Result: $CHECKS_PASSED/4 critical checks passed"
echo "========================================"
echo ""

# Final verdict
if [ $CHECKS_FAILED -eq 0 ]; then
    echo "🎉 DRY-RUN PASSED"
    echo ""
    echo "✅ All critical validations passed"
    echo "✅ Wave 11 bug fixes appear operational"
    echo "✅ Ready for full hyperopt campaign"
    echo ""
    echo "Next Steps:"
    echo "  1. Review hyperopt results in logs: $LOG_FILE"
    echo "  2. Verify action diversity is reasonable (not 100% HOLD)"
    echo "  3. Proceed with full 100-trial campaign:"
    echo "     cargo run --release -p ml --example hyperopt_dqn_demo --features cuda -- \\"
    echo "       --parquet-file $PARQUET_FILE --trials 100 --epochs 50"
    echo ""
    echo "Estimated full campaign:"
    echo "  Duration: 60-90 minutes"
    echo "  Cost: \$0.25-\$0.38 (RTX A4000)"
    exit 0
else
    echo "⚠️  DRY-RUN COMPLETED WITH WARNINGS"
    echo ""
    echo "❌ $CHECKS_FAILED critical checks failed"
    echo "⚠️  Review issues above before proceeding"
    echo ""
    echo "Troubleshooting:"
    echo "  1. Check logs: $LOG_FILE"
    echo "  2. Look for specific error messages"
    echo "  3. Verify Wave 11 fixes are compiled correctly"
    echo "  4. Re-run dry-run after fixes"
    echo ""
    echo "Common Issues:"
    echo "  - Gradient explosions: Check gradient_clip_norm in DQNHyperparameters"
    echo "  - 100% HOLD: Check hold_penalty_weight and movement_threshold"
    echo "  - Portfolio errors: Check PortfolioTracker initialization"
    exit 1
fi