#!/bin/bash
# Enhanced DQN Checkpoint Backtest Script
# Agent 132 - 2025-10-14
#
# Usage:
#   ./backtest_dqn_trials_enhanced.sh --quick       # Test trial 35 only (10 min)
#   ./backtest_dqn_trials_enhanced.sh --sample      # Test 10 trials (1 hour)
#   ./backtest_dqn_trials_enhanced.sh --full        # Test all 36 trials (3-6 hours)

set -e

# Configuration
RESULTS_DIR="results/dqn_backtest"
RESULTS_FILE="$RESULTS_DIR/dqn_backtest_results.json"
DATA_FILE="test_data/ES.FUT.dbn"
START_DATE="2024-01-02"
SHARPE_THRESHOLD=1.5

# Sample trials (representative distribution)
SAMPLE_TRIALS=(0 4 8 12 16 20 24 28 32 35)

# Colors for output
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
RED='\033[0;31m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Create results directory
mkdir -p "$RESULTS_DIR"

# Function to backtest a single checkpoint
backtest_checkpoint() {
    local trial_num=$1
    local checkpoint_path="ml/tuning_checkpoints/trial_${trial_num}/checkpoint_epoch_50.safetensors"
    local output_file="$RESULTS_DIR/trial_${trial_num}_backtest.json"

    if [ ! -f "$checkpoint_path" ]; then
        echo -e "${RED}⚠️  Checkpoint not found: $checkpoint_path${NC}"
        return 1
    fi

    echo -e "${BLUE}🔬 Testing trial ${trial_num}...${NC}"

    # Check if backtest_dqn example exists
    if ! cargo run -p ml --example backtest_dqn -- --help &>/dev/null; then
        echo -e "${YELLOW}⚠️  backtest_dqn example not found${NC}"
        echo -e "${YELLOW}   Creating placeholder result...${NC}"

        # Create placeholder result (for testing infrastructure)
        cat > "$output_file" << EOF
{
  "trial_num": ${trial_num},
  "checkpoint": "${checkpoint_path}",
  "sharpe_ratio": 0.0,
  "total_return": 0.0,
  "max_drawdown": 0.0,
  "win_rate": 0.0,
  "num_trades": 0,
  "status": "backtest_example_not_implemented",
  "note": "Requires implementation of ml/examples/backtest_dqn.rs"
}
EOF
        return 0
    fi

    # Run actual backtest
    cargo run --release -p ml --example backtest_dqn -- \
        --checkpoint "$checkpoint_path" \
        --data "$DATA_FILE" \
        --start-date "$START_DATE" \
        --output "$output_file" 2>&1 | grep -E "(Sharpe|Return|Drawdown|Win)"

    # Extract Sharpe ratio
    if [ -f "$output_file" ]; then
        SHARPE=$(jq -r '.sharpe_ratio // 0' "$output_file")
        echo -e "${GREEN}   Sharpe: ${SHARPE}${NC}"
    fi
}

# Function to analyze results
analyze_results() {
    echo ""
    echo "=" | head -c 80
    echo ""
    echo -e "${BLUE}📊 ANALYZING BACKTEST RESULTS${NC}"
    echo "=" | head -c 80
    echo ""

    if [ ! -f "$RESULTS_FILE" ]; then
        echo -e "${RED}❌ Results file not found: $RESULTS_FILE${NC}"
        return 1
    fi

    # Use Python for analysis
    python3 << 'EOF'
import json
import sys

results_file = sys.argv[1]
try:
    with open(results_file, 'r') as f:
        results = json.load(f)
except Exception as e:
    print(f"❌ Error reading results: {e}")
    sys.exit(1)

if not results:
    print("❌ No results found")
    sys.exit(1)

# Filter valid results
valid_results = [r for r in results if r.get('sharpe_ratio', 0) > 0]

if not valid_results:
    print("⚠️  No valid backtest results (all Sharpe ratios are 0)")
    print("\n💡 This likely means the backtest_dqn example needs to be implemented")
    print("   See: ml/examples/backtest_dqn.rs")
    sys.exit(0)

# Sort by Sharpe ratio
sorted_results = sorted(valid_results, key=lambda x: x.get('sharpe_ratio', 0), reverse=True)

print(f"\n✅ Analyzed {len(valid_results)} successful backtests")
print(f"\n🏆 TOP 3 PERFORMING CHECKPOINTS:\n")

for i, result in enumerate(sorted_results[:3], 1):
    trial_num = result.get('trial_num', 'unknown')
    sharpe = result.get('sharpe_ratio', 0)
    ret = result.get('total_return', 0)
    dd = result.get('max_drawdown', 0)
    win = result.get('win_rate', 0)

    print(f"{i}. Trial {trial_num}")
    print(f"   Sharpe Ratio: {sharpe:.3f}")
    print(f"   Total Return: {ret:.2%}")
    print(f"   Max Drawdown: {dd:.2%}")
    print(f"   Win Rate:     {win:.2%}")
    print()

# Best checkpoint
best = sorted_results[0]
best_trial = best.get('trial_num')
best_sharpe = best.get('sharpe_ratio')

print("=" * 80)
print(f"✅ RECOMMENDATION: Use trial_{best_trial} checkpoint")
print(f"   Sharpe: {best_sharpe:.3f}")
print(f"   Path: ml/tuning_checkpoints/trial_{best_trial}/checkpoint_epoch_50.safetensors")
print("=" * 80)

# Save summary
summary = {
    "best_trial": best_trial,
    "best_sharpe": best_sharpe,
    "top_3": [
        {
            "trial_num": r.get('trial_num'),
            "sharpe_ratio": r.get('sharpe_ratio'),
            "total_return": r.get('total_return'),
            "max_drawdown": r.get('max_drawdown'),
            "win_rate": r.get('win_rate')
        }
        for r in sorted_results[:3]
    ]
}

import os
summary_file = os.path.join(os.path.dirname(results_file), "summary.json")
with open(summary_file, 'w') as f:
    json.dump(summary, f, indent=2)

print(f"\n💾 Summary saved to: {summary_file}")

EOF
    python3 - "$RESULTS_FILE"
}

# Main execution
echo "=" | head -c 80
echo ""
echo -e "${BLUE}DQN CHECKPOINT BACKTEST UTILITY${NC}"
echo "=" | head -c 80
echo ""

MODE=${1:-}

case $MODE in
    --quick)
        echo -e "${GREEN}Mode: QUICK TEST (trial 35 only)${NC}"
        echo -e "${GREEN}Expected time: 10 minutes${NC}"
        echo ""

        echo "[" > "$RESULTS_FILE"
        backtest_checkpoint 35
        cat "$RESULTS_DIR/trial_35_backtest.json" >> "$RESULTS_FILE"
        echo "]" >> "$RESULTS_FILE"

        # Check if meets threshold
        SHARPE=$(jq -r '.[0].sharpe_ratio // 0' "$RESULTS_FILE")
        if (( $(echo "$SHARPE > $SHARPE_THRESHOLD" | bc -l 2>/dev/null || echo 0) )); then
            echo ""
            echo -e "${GREEN}✅ Trial 35 exceeds threshold (Sharpe=$SHARPE > $SHARPE_THRESHOLD)${NC}"
            echo -e "${GREEN}   Recommendation: Use trial_35 for production${NC}"
        else
            echo ""
            echo -e "${YELLOW}⚠️  Trial 35 below threshold (Sharpe=$SHARPE < $SHARPE_THRESHOLD)${NC}"
            echo -e "${YELLOW}   Recommendation: Run --sample or --full backtest${NC}"
        fi
        ;;

    --sample)
        echo -e "${GREEN}Mode: SAMPLE TEST (10 trials)${NC}"
        echo -e "${GREEN}Expected time: 1 hour${NC}"
        echo -e "${GREEN}Trials: ${SAMPLE_TRIALS[@]}${NC}"
        echo ""

        echo "[" > "$RESULTS_FILE"
        first=true
        for trial_num in "${SAMPLE_TRIALS[@]}"; do
            if [ "$first" = false ]; then
                echo "," >> "$RESULTS_FILE"
            fi
            first=false

            backtest_checkpoint "$trial_num"
            cat "$RESULTS_DIR/trial_${trial_num}_backtest.json" >> "$RESULTS_FILE"
        done
        echo "" >> "$RESULTS_FILE"
        echo "]" >> "$RESULTS_FILE"

        analyze_results
        ;;

    --full)
        echo -e "${GREEN}Mode: FULL TEST (all 36 trials)${NC}"
        echo -e "${GREEN}Expected time: 3-6 hours${NC}"
        echo ""

        echo "[" > "$RESULTS_FILE"
        first=true
        for trial_num in {0..35}; do
            if [ "$first" = false ]; then
                echo "," >> "$RESULTS_FILE"
            fi
            first=false

            backtest_checkpoint "$trial_num"

            if [ -f "$RESULTS_DIR/trial_${trial_num}_backtest.json" ]; then
                cat "$RESULTS_DIR/trial_${trial_num}_backtest.json" >> "$RESULTS_FILE"
            fi
        done
        echo "" >> "$RESULTS_FILE"
        echo "]" >> "$RESULTS_FILE"

        analyze_results
        ;;

    *)
        echo -e "${RED}Usage:${NC}"
        echo "  $0 --quick    # Test trial 35 only (10 min)"
        echo "  $0 --sample   # Test 10 trials (1 hour)"
        echo "  $0 --full     # Test all 36 trials (3-6 hours)"
        echo ""
        echo -e "${YELLOW}No mode specified. Defaulting to --quick${NC}"
        echo ""
        exec "$0" --quick
        ;;
esac

echo ""
echo "=" | head -c 80
echo ""
echo -e "${GREEN}✅ BACKTEST COMPLETE${NC}"
echo "=" | head -c 80
echo ""
echo -e "Results: ${RESULTS_FILE}"
echo -e "Individual results: ${RESULTS_DIR}/trial_*_backtest.json"
echo ""