//! DQN Checkpoint Analysis Tool
//!
//! Analyzes all 51 DQN checkpoints to identify best performers and predict
//! which epochs will trade more actively in backtesting.
//!
//! Usage:
//!   cargo run -p ml --example analyze_dqn_checkpoints --release
//!
//! Outputs:
//! - Checkpoint rankings by Q-value magnitude
//! - Size analysis for model complexity
//! - Recommendations for top 10 candidates
//! - Script to test multiple checkpoints systematically

use ml::checkpoint::{CheckpointConfig, CheckpointManager};
use ml::trainers::dqn::{DQNTrainer, DQNHyperparameters};
use std::collections::HashMap;
use std::path::PathBuf;
use tch::{Device, Kind, Tensor};

#[derive(Debug, Clone)]
struct CheckpointAnalysis {
    epoch: u32,
    file_path: PathBuf,
    file_size: u64,
    q_value_magnitude: f64,
    weight_norm: f64,
    actions_sampled: Vec<i64>,
    trading_activity_score: f64,
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("🔍 DQN Checkpoint Analysis Tool");
    println!("=====================================\n");

    let checkpoint_dir = PathBuf::from("ml/trained_models/production/dqn_real_data");

    // Phase 1: Discover all checkpoints
    println!("📂 Phase 1: Discovering checkpoints...");
    let checkpoints = discover_checkpoints(&checkpoint_dir)?;
    println!("   Found {} checkpoint files\n", checkpoints.len());

    // Phase 2: Analyze each checkpoint
    println!("🔬 Phase 2: Analyzing checkpoints (this may take a few minutes)...");
    let mut analyses = Vec::new();

    for (epoch, file_path) in checkpoints.iter() {
        print!("   Analyzing epoch {}... ", epoch);
        match analyze_checkpoint(*epoch, file_path).await {
            Ok(analysis) => {
                println!("✅ Q-mag: {:.4}, Activity: {:.2}",
                    analysis.q_value_magnitude,
                    analysis.trading_activity_score
                );
                analyses.push(analysis);
            }
            Err(e) => {
                println!("❌ Error: {}", e);
            }
        }
    }

    println!("\n📊 Phase 3: Rankings and Recommendations\n");

    // Phase 3a: Rank by Q-value magnitude (higher = more aggressive trading)
    let mut by_q_value = analyses.clone();
    by_q_value.sort_by(|a, b| b.q_value_magnitude.partial_cmp(&a.q_value_magnitude).unwrap());

    println!("🏆 Top 10 by Q-Value Magnitude (Most Likely to Trade):");
    println!("    Rank | Epoch | Q-Magnitude | Activity | Size  | Notes");
    println!("    -----|-------|-------------|----------|-------|------");
    for (i, analysis) in by_q_value.iter().take(10).enumerate() {
        let note = if analysis.epoch < 50 {
            "Early epoch - exploring"
        } else if analysis.epoch < 150 {
            "Mid training - learning"
        } else if analysis.epoch < 300 {
            "Late training - refining"
        } else {
            "Final phase - converged"
        };
        println!("    {:>4} | {:>5} | {:>11.4} | {:>8.2} | {:>4}K | {}",
            i + 1,
            analysis.epoch,
            analysis.q_value_magnitude,
            analysis.trading_activity_score,
            analysis.file_size / 1024,
            note
        );
    }

    // Phase 3b: Rank by trading activity score
    let mut by_activity = analyses.clone();
    by_activity.sort_by(|a, b| b.trading_activity_score.partial_cmp(&a.trading_activity_score).unwrap());

    println!("\n📈 Top 10 by Trading Activity Score (Action Diversity):");
    println!("    Rank | Epoch | Activity | Q-Magnitude | Actions Distribution");
    println!("    -----|-------|----------|-------------|---------------------");
    for (i, analysis) in by_activity.iter().take(10).enumerate() {
        let action_dist = format!("Buy:{} Hold:{} Sell:{}",
            analysis.actions_sampled.iter().filter(|&&a| a == 0).count(),
            analysis.actions_sampled.iter().filter(|&&a| a == 1).count(),
            analysis.actions_sampled.iter().filter(|&&a| a == 2).count()
        );
        println!("    {:>4} | {:>5} | {:>8.2} | {:>11.4} | {}",
            i + 1,
            analysis.epoch,
            analysis.trading_activity_score,
            analysis.q_value_magnitude,
            action_dist
        );
    }

    // Phase 3c: Identify early vs late epoch patterns
    println!("\n📊 Training Phase Analysis:");

    let early_checkpoints: Vec<_> = analyses.iter().filter(|a| a.epoch <= 100).collect();
    let mid_checkpoints: Vec<_> = analyses.iter().filter(|a| a.epoch > 100 && a.epoch <= 300).collect();
    let late_checkpoints: Vec<_> = analyses.iter().filter(|a| a.epoch > 300).collect();

    let avg_q_early = early_checkpoints.iter().map(|a| a.q_value_magnitude).sum::<f64>() / early_checkpoints.len() as f64;
    let avg_q_mid = mid_checkpoints.iter().map(|a| a.q_value_magnitude).sum::<f64>() / mid_checkpoints.len() as f64;
    let avg_q_late = late_checkpoints.iter().map(|a| a.q_value_magnitude).sum::<f64>() / late_checkpoints.len() as f64;

    println!("    Early epochs (1-100):   Avg Q-magnitude: {:.4}", avg_q_early);
    println!("    Mid epochs   (101-300): Avg Q-magnitude: {:.4}", avg_q_mid);
    println!("    Late epochs  (301-500): Avg Q-magnitude: {:.4}", avg_q_late);

    if avg_q_early > avg_q_late {
        println!("\n    💡 Insight: Early epochs have higher Q-values → More aggressive trading");
    } else {
        println!("\n    💡 Insight: Later epochs have higher Q-values → Better strategy learned");
    }

    // Phase 4: Generate testing script
    println!("\n🚀 Phase 4: Test Strategy Recommendations\n");

    println!("📝 Top 10 Candidates for Systematic Backtesting:");
    println!("    (Balanced selection across training phases)\n");

    // Select diverse checkpoints
    let mut top_candidates = Vec::new();

    // Add top 3 by Q-value
    for analysis in by_q_value.iter().take(3) {
        top_candidates.push((analysis.epoch, "High Q-value"));
    }

    // Add 3 from early training (epochs 10-100)
    let early_high_q: Vec<_> = by_q_value.iter().filter(|a| a.epoch >= 10 && a.epoch <= 100).take(3).collect();
    for analysis in early_high_q {
        if !top_candidates.iter().any(|(e, _)| *e == analysis.epoch) {
            top_candidates.push((analysis.epoch, "Early exploration"));
        }
    }

    // Add 2 from mid training (epochs 100-250)
    let mid_high_q: Vec<_> = by_q_value.iter().filter(|a| a.epoch > 100 && a.epoch <= 250).take(2).collect();
    for analysis in mid_high_q {
        if !top_candidates.iter().any(|(e, _)| *e == analysis.epoch) {
            top_candidates.push((analysis.epoch, "Mid learning"));
        }
    }

    // Add 2 from late training (epochs 400-500)
    let late_checkpoints_sorted: Vec<_> = by_q_value.iter().filter(|a| a.epoch >= 400).take(2).collect();
    for analysis in late_checkpoints_sorted {
        if !top_candidates.iter().any(|(e, _)| *e == analysis.epoch) {
            top_candidates.push((analysis.epoch, "Converged model"));
        }
    }

    // Sort by epoch for clearer presentation
    top_candidates.sort_by_key(|(e, _)| *e);

    for (i, (epoch, reason)) in top_candidates.iter().take(10).enumerate() {
        println!("    {}. Epoch {:>3} - {}", i + 1, epoch, reason);
    }

    // Generate shell script
    println!("\n📜 Generating test script: test_top_dqn_checkpoints.sh");
    generate_test_script(&top_candidates)?;

    println!("\n✅ Analysis complete!");
    println!("\n🎯 Next Steps:");
    println!("   1. Run: bash test_top_dqn_checkpoints.sh");
    println!("   2. Compare backtest results (PnL, trades, Sharpe)");
    println!("   3. Select best performer for production");

    Ok(())
}

/// Discover all checkpoint files in the directory
fn discover_checkpoints(dir: &PathBuf) -> Result<Vec<(u32, PathBuf)>, Box<dyn std::error::Error>> {
    let mut checkpoints = Vec::new();

    for entry in std::fs::read_dir(dir)? {
        let entry = entry?;
        let path = entry.path();

        if let Some(filename) = path.file_name().and_then(|n| n.to_str()) {
            if filename.starts_with("dqn_epoch_") && filename.ends_with(".safetensors") {
                // Extract epoch number from filename like "dqn_epoch_100.safetensors"
                if let Some(epoch_str) = filename.strip_prefix("dqn_epoch_").and_then(|s| s.strip_suffix(".safetensors")) {
                    if let Ok(epoch) = epoch_str.parse::<u32>() {
                        checkpoints.push((epoch, path));
                    }
                }
            }
        }
    }

    checkpoints.sort_by_key(|(epoch, _)| *epoch);
    Ok(checkpoints)
}

/// Analyze a single checkpoint
async fn analyze_checkpoint(epoch: u32, file_path: &PathBuf) -> Result<CheckpointAnalysis, Box<dyn std::error::Error>> {
    // Get file size
    let metadata = std::fs::metadata(file_path)?;
    let file_size = metadata.len();

    // Load checkpoint
    let device = Device::cuda_if_available(0);
    let hyperparams = DQNHyperparameters {
        state_dim: 16, // From production training config
        action_dim: 3, // Buy, Hold, Sell
        learning_rate: 0.0001,
        gamma: 0.99,
        epsilon_start: 1.0,
        epsilon_end: 0.1,
        epsilon_decay: 0.995,
        batch_size: 64,
        memory_capacity: 10000,
        target_update_freq: 10,
    };

    let mut trainer = DQNTrainer::new(hyperparams)?;

    // Load checkpoint data
    let checkpoint_data = std::fs::read(file_path)?;
    trainer.deserialize_model(&checkpoint_data).await?;

    // Sample Q-values on diverse test states
    let test_states = generate_test_states(100, 16);
    let mut q_values = Vec::new();
    let mut actions_sampled = Vec::new();

    for state in test_states {
        let state_tensor = Tensor::of_slice(&state).to_kind(Kind::Float).to_device(device);
        let q_vals = trainer.get_q_values(&state_tensor).await?;

        // Get action with highest Q-value
        let action = q_vals.argmax(-1, false).int64_value(&[]);
        actions_sampled.push(action);

        // Store Q-values for magnitude calculation
        let q_vals_vec: Vec<f64> = Vec::try_from(q_vals)?;
        q_values.extend(q_vals_vec);
    }

    // Calculate Q-value magnitude (mean absolute value)
    let q_value_magnitude = q_values.iter().map(|q| q.abs()).sum::<f64>() / q_values.len() as f64;

    // Calculate weight norm (for model complexity)
    let weight_norm = calculate_weight_norm(&trainer)?;

    // Calculate trading activity score (action diversity)
    let buy_count = actions_sampled.iter().filter(|&&a| a == 0).count();
    let hold_count = actions_sampled.iter().filter(|&&a| a == 1).count();
    let sell_count = actions_sampled.iter().filter(|&&a| a == 2).count();

    // Entropy-based diversity: higher = more diverse actions
    let total = actions_sampled.len() as f64;
    let buy_prob = buy_count as f64 / total;
    let hold_prob = hold_count as f64 / total;
    let sell_prob = sell_count as f64 / total;

    let entropy = -(
        buy_prob * buy_prob.ln().max(0.0) +
        hold_prob * hold_prob.ln().max(0.0) +
        sell_prob * sell_prob.ln().max(0.0)
    );

    // Trading activity: prefer non-hold actions
    let non_hold_ratio = (buy_count + sell_count) as f64 / total;
    let trading_activity_score = entropy * non_hold_ratio * 100.0;

    Ok(CheckpointAnalysis {
        epoch,
        file_path: file_path.clone(),
        file_size,
        q_value_magnitude,
        weight_norm,
        actions_sampled,
        trading_activity_score,
    })
}

/// Generate diverse test states for Q-value sampling
fn generate_test_states(count: usize, state_dim: usize) -> Vec<Vec<f32>> {
    let mut states = Vec::new();

    for i in 0..count {
        let mut state = vec![0.0; state_dim];

        // Generate diverse patterns
        match i % 5 {
            0 => {
                // Bullish pattern (positive features)
                for j in 0..state_dim {
                    state[j] = (0.5 + (j as f32 * 0.1)) % 1.0;
                }
            }
            1 => {
                // Bearish pattern (negative features)
                for j in 0..state_dim {
                    state[j] = -(0.5 + (j as f32 * 0.1)) % 1.0;
                }
            }
            2 => {
                // Neutral pattern (near zero)
                for j in 0..state_dim {
                    state[j] = (j as f32 * 0.01) - 0.08;
                }
            }
            3 => {
                // High volatility pattern (large values)
                for j in 0..state_dim {
                    state[j] = ((j as f32).sin() * 2.0) % 1.5;
                }
            }
            4 => {
                // Low volatility pattern (small values)
                for j in 0..state_dim {
                    state[j] = (j as f32 * 0.001) - 0.01;
                }
            }
            _ => unreachable!()
        }

        states.push(state);
    }

    states
}

/// Calculate L2 norm of model weights
fn calculate_weight_norm(trainer: &DQNTrainer) -> Result<f64, Box<dyn std::error::Error>> {
    // Placeholder - would need to access internal model weights
    // For now, return a default value
    Ok(1.0)
}

/// Generate shell script to test top checkpoints
fn generate_test_script(candidates: &[(u32, &str)]) -> Result<(), Box<dyn std::error::Error>> {
    let script = format!(
r#"#!/bin/bash
# DQN Checkpoint Testing Script
# Generated by analyze_dqn_checkpoints
# Tests top {} checkpoint candidates

set -e

CHECKPOINT_DIR="ml/trained_models/production/dqn_real_data"
RESULTS_DIR="ml/backtest_results"
mkdir -p $RESULTS_DIR

echo "🧪 Testing {} DQN Checkpoints"
echo "=================================="
echo ""

{}

echo ""
echo "✅ All tests complete!"
echo "📊 Results saved to: $RESULTS_DIR"
echo ""
echo "🎯 Compare results with:"
echo "   ls -lh $RESULTS_DIR/dqn_epoch_*.json"
"#,
        candidates.len(),
        candidates.len(),
        candidates.iter().map(|(epoch, reason)| {
            format!(
r#"echo "Testing Epoch {} - {}"
cargo run -p backtesting_service --example backtest_dqn --release -- \
  --checkpoint $CHECKPOINT_DIR/dqn_epoch_{}.safetensors \
  --output $RESULTS_DIR/dqn_epoch_{}.json \
  --data test_data/real/databento/ml_training/ZN.FUT_ohlcv-1m_2024-04-17.dbn
echo "  ✅ Epoch {} complete"
echo """#,
                epoch, reason, epoch, epoch, epoch
            )
        }).collect::<Vec<_>>().join("\n")
    );

    std::fs::write("test_top_dqn_checkpoints.sh", script)?;

    // Make executable
    #[cfg(unix)]
    {
        use std::os::unix::fs::PermissionsExt;
        let mut perms = std::fs::metadata("test_top_dqn_checkpoints.sh")?.permissions();
        perms.set_mode(0o755);
        std::fs::set_permissions("test_top_dqn_checkpoints.sh", perms)?;
    }

    Ok(())
}