//! Training Monitor for per-epoch tracking and validation
//!
//! Tracks rewards, actions, Q-values, episode lengths, and validates
//! training health (constant rewards, action diversity, Q-value balance).

use anyhow::Result;
use tracing::{debug, warn};

use crate::dqn::action_space::FactoredAction;

/// Training monitor for per-epoch tracking and validation
pub(crate) struct TrainingMonitor {
    pub(crate) epoch: usize,
    pub(crate) reward_history: Vec<f32>,
    pub(crate) action_counts: [usize; 9], // 5 exposure levels (DQN action space)
    pub(crate) q_value_sums: [f64; 9],    // Sum of Q-values per exposure action
    pub(crate) q_value_counts: [usize; 9], // Count of Q-values per exposure action
    pub(crate) order_type_counts: [usize; 3],  // Market, LimitMaker, IoC
    pub(crate) urgency_counts: [usize; 3],     // Patient, Normal, Aggressive
    /// Factored action counts: 9 exposure × 3 order × 3 urgency = 81 actions.
    /// Index = exposure * 9 + order * 3 + urgency (0-44).
    pub(crate) factored_action_counts: [usize; 81],
    pub(crate) consecutive_constant_epochs: usize,
    // Q-value range tracking (WAVE 9-11 production monitoring)
    pub(crate) q_value_min: f64,
    pub(crate) q_value_max: f64,
    pub(crate) q_value_history: Vec<f64>, // Per-step Q-values for mean calculation

    // WAVE P2: Episode length tracking
    pub(crate) episode_lengths: Vec<usize>,
    pub(crate) episode_start_step: usize,
    pub(crate) barrier_exit_counts: [usize; 4], // [profit, stop, time, boundary]
}

impl TrainingMonitor {
    pub(crate) fn new(epoch: usize) -> Self {
        Self {
            epoch,
            reward_history: Vec::new(),
            action_counts: [0; 9],
            q_value_sums: [0.0; 9],
            q_value_counts: [0; 9],
            order_type_counts: [0; 3],
            urgency_counts: [0; 3],
            factored_action_counts: [0; 81],
            consecutive_constant_epochs: 0,
            q_value_min: f64::INFINITY,
            q_value_max: f64::NEG_INFINITY,
            q_value_history: Vec::new(),

            // WAVE P2: Episode tracking
            episode_lengths: Vec::new(),
            episode_start_step: 0,
            barrier_exit_counts: [0; 4], // [profit=0, stop=1, time=2, boundary=3]
        }
    }

    /// Add reward to tracking (with bounded history)
    pub(crate) fn track_reward(&mut self, reward: f32) {
        self.reward_history.push(reward);
        // MEMORY LEAK FIX: Limit reward history to last 1000 entries per epoch
        // Each trial has ~10-50 epochs, so this limits to ~10-50K entries total
        // vs unbounded growth causing OOM at 10-30 trials
        if self.reward_history.len() > 1000 {
            self.reward_history.drain(0..500); // Remove oldest 500, keep newest 500
        }
    }

    /// Add action to tracking by exposure index (0-4)
    pub(crate) fn track_action_by_exposure(&mut self, exposure_idx: usize) {
        if exposure_idx < 5 {
            self.action_counts[exposure_idx] += 1;
        }
    }

    /// Add action to tracking (extracts exposure from FactoredAction)
    pub(crate) fn track_action(&mut self, action: &FactoredAction) {
        let exp_idx = action.exposure as usize; // 0-4
        self.action_counts[exp_idx] += 1;
        let order_idx = action.order as usize; // 0-2
        self.order_type_counts[order_idx] += 1;
        let urgency_idx = action.urgency as usize; // 0-2
        self.urgency_counts[urgency_idx] += 1;
        // Track composite factored action index (0-44)
        let factored_idx = exp_idx * 9 + order_idx * 3 + urgency_idx;
        if factored_idx < 81 {
            self.factored_action_counts[factored_idx] += 1;
        }
    }

    /// Track order type dimension (0=Market, 1=LimitMaker, 2=IoC)
    pub(crate) fn track_order_type(&mut self, idx: usize) {
        if idx < 3 {
            self.order_type_counts[idx] += 1;
        }
    }

    /// Track urgency dimension (0=Patient, 1=Normal, 2=Aggressive)
    pub(crate) fn track_urgency(&mut self, idx: usize) {
        if idx < 3 {
            self.urgency_counts[idx] += 1;
        }
    }

    /// Add Q-value to tracking by exposure index (0-4)
    pub(crate) fn track_q_value_by_exposure(&mut self, exposure_idx: usize, q_value: f64) {
        if exposure_idx < 5 {
            self.q_value_sums[exposure_idx] += q_value;
            self.q_value_counts[exposure_idx] += 1;
        }
    }

    /// Add Q-value to tracking (extracts exposure from FactoredAction)
    pub(crate) fn track_q_value(&mut self, action: &FactoredAction, q_value: f64) {
        let idx = action.exposure as usize; // 0-4
        self.q_value_sums[idx] += q_value;
        self.q_value_counts[idx] += 1;
    }

    /// Track Q-value range for monitoring (WAVE 9-11 production)
    pub(crate) fn track_q_value_range(&mut self, q_value: f64) {
        if q_value < self.q_value_min {
            self.q_value_min = q_value;
        }
        if q_value > self.q_value_max {
            self.q_value_max = q_value;
        }
        self.q_value_history.push(q_value);
        // MEMORY LEAK FIX: Limit Q-value history to last 1000 entries per epoch
        if self.q_value_history.len() > 1000 {
            self.q_value_history.drain(0..500); // Remove oldest 500, keep newest 500
        }
    }

    /// Get Q-value statistics (min, max, mean)
    pub(crate) fn get_q_value_stats(&self) -> (f64, f64, f64) {
        if self.q_value_history.is_empty() {
            return (0.0, 0.0, 0.0);
        }
        let mean = self.q_value_history.iter().sum::<f64>() / self.q_value_history.len() as f64;
        (self.q_value_min, self.q_value_max, mean)
    }

    /// Validate rewards are not constant
    pub(crate) fn validate_rewards(&mut self) -> Result<()> {
        // Need at least 2 samples for meaningful variance.  The GPU path pushes
        // a single aggregated mean_reward per epoch — std is trivially 0 with N=1.
        if self.reward_history.len() < 2 {
            return Ok(());
        }

        let mean = self.reward_history.iter().sum::<f32>() / self.reward_history.len() as f32;
        let variance = self
            .reward_history
            .iter()
            .map(|r| (r - mean).powi(2))
            .sum::<f32>()
            / self.reward_history.len() as f32;
        let std = variance.sqrt();

        // Check if all rewards are identical (std == 0) or nearly constant (std < 0.01)
        if std < 0.01 {
            self.consecutive_constant_epochs += 1;

            warn!(
                "⚠️  CONSTANT REWARDS DETECTED at epoch {}! std={:.6}, mean={:.4}, consecutive_epochs={}",
                self.epoch, std, mean, self.consecutive_constant_epochs
            );

            // Panic if constant for 5+ consecutive epochs (critical bug)
            if self.consecutive_constant_epochs >= 5 {
                return Err(anyhow::anyhow!(
                    "❌ CRITICAL: Constant rewards for {} consecutive epochs! std={:.6}, mean={:.4}\n\
                     This indicates a reward calculation bug. Training aborted.",
                    self.consecutive_constant_epochs, std, mean
                ));
            }
        } else {
            // Reset counter if variance is healthy
            self.consecutive_constant_epochs = 0;
        }

        Ok(())
    }

    /// Validate action diversity across all 45 factored actions
    pub(crate) fn validate_action_diversity(&self) -> Result<()> {
        let total_factored: usize = self.factored_action_counts.iter().sum();
        let total_exposure: usize = self.action_counts.iter().sum();

        if total_exposure == 0 {
            return Ok(()); // No actions yet, skip validation
        }

        // Check exposure diversity (5 levels)
        let exposure_names = ["S100", "S75", "S50", "S25", "Flat", "L25", "L50", "L75", "L100"];
        for (i, &count) in self.action_counts.iter().enumerate() {
            let percentage = (count as f64 / total_exposure as f64) * 100.0;
            if percentage < 5.0 {
                warn!(
                    "LOW EXPOSURE DIVERSITY at epoch {}: {} only {:.1}% ({}/{})",
                    self.epoch, exposure_names[i], percentage, count, total_exposure
                );
            }
        }

        // Check factored action diversity (45 actions)
        if total_factored > 0 {
            let unique_factored = self.factored_action_counts.iter().filter(|&&c| c > 0).count();
            if unique_factored < 5 {
                warn!(
                    "FACTORED ACTION COLLAPSE at epoch {}: only {}/45 composite actions used",
                    self.epoch, unique_factored
                );
            }
        }

        Ok(())
    }

    /// Validate Q-value balance across exposure actions
    pub(crate) fn validate_q_value_balance(&self) -> Result<()> {
        // Calculate average Q-value per exposure level
        let mut avg_q_values = [0.0_f64; 5];
        for i in 0..5 {
            if self.q_value_counts[i] > 0 {
                avg_q_values[i] = self.q_value_sums[i] / self.q_value_counts[i] as f64;
            }
        }

        // Check if any Q-values diverge > 1000 from others
        let max_q = avg_q_values.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
        let min_q = avg_q_values.iter().cloned().fold(f64::INFINITY, f64::min);

        if (max_q - min_q).abs() > 1000.0 {
            let exposure_names = ["S100", "S75", "S50", "S25", "Flat", "L25", "L50", "L75", "L100"];
            warn!(
                "Q-VALUE DIVERGENCE at epoch {}: {}",
                self.epoch,
                avg_q_values.iter().enumerate()
                    .map(|(i, q)| format!("{}={:.2}", exposure_names[i], q))
                    .collect::<Vec<_>>()
                    .join(", ")
            );
        }

        Ok(())
    }

    /// Log action distribution every 10 epochs
    pub(crate) fn log_action_distribution(&self) {
        if self.epoch % 10 == 0 {
            let total_actions: usize = self.action_counts.iter().sum();
            if total_actions > 0 {
                let exposure_names = ["S100", "S75", "S50", "S25", "Flat", "L25", "L50", "L75", "L100"];

                // Log all 5 exposure actions
                debug!(
                    "Action Distribution [Epoch {}] - 5 Exposure Actions:",
                    self.epoch
                );
                for (idx, &count) in self.action_counts.iter().enumerate() {
                    let pct = (count as f64 / total_actions as f64) * 100.0;
                    debug!("  [{}] {}: {} ({:.1}%)", idx, exposure_names[idx], count, pct);
                }

                // Log order type distribution (branching DQN)
                let total_order: usize = self.order_type_counts.iter().sum();
                if total_order > 0 {
                    let order_names = ["Market", "LimitMaker", "IoC"];
                    debug!(
                        "Order Type Distribution [Epoch {}]:",
                        self.epoch
                    );
                    for (idx, &count) in self.order_type_counts.iter().enumerate() {
                        let pct = (count as f64 / total_order as f64) * 100.0;
                        debug!("  [{}] {}: {} ({:.1}%)", idx, order_names[idx], count, pct);
                    }
                }

                // Log urgency distribution (branching DQN)
                let total_urgency: usize = self.urgency_counts.iter().sum();
                if total_urgency > 0 {
                    let urgency_names = ["Patient", "Normal", "Aggressive"];
                    debug!(
                        "Urgency Distribution [Epoch {}]:",
                        self.epoch
                    );
                    for (idx, &count) in self.urgency_counts.iter().enumerate() {
                        let pct = (count as f64 / total_urgency as f64) * 100.0;
                        debug!("  [{}] {}: {} ({:.1}%)", idx, urgency_names[idx], count, pct);
                    }
                }

                // Log factored 45-action diversity summary
                let total_factored: usize = self.factored_action_counts.iter().sum();
                if total_factored > 0 {
                    let unique_factored = self.factored_action_counts.iter()
                        .filter(|&&c| c > 0).count();
                    debug!(
                        "Factored Action Diversity [Epoch {}]: {}/45 actions used ({:.1}%)",
                        self.epoch,
                        unique_factored,
                        (unique_factored as f64 / 45.0) * 100.0
                    );
                }

                // Log average Q-values per exposure
                debug!("Average Q-values [Epoch {}]:", self.epoch);
                for idx in 0..5 {
                    if self.q_value_counts[idx] > 0 {
                        let avg_q = self.q_value_sums[idx] / self.q_value_counts[idx] as f64;
                        debug!("  [{}] {}: Q={:.4}", idx, exposure_names[idx], avg_q);
                    }
                }
            }
        }
    }

    /// Run all validations
    pub(crate) fn validate_all(&mut self) -> Result<()> {
        self.validate_rewards()?;
        self.validate_action_diversity()?;
        self.validate_q_value_balance()?;
        self.log_action_distribution();
        Ok(())
    }

    /// Track episode end and record length/exit reason
    pub(crate) fn track_episode_end(&mut self, current_step: usize, barrier_label: Option<i8>) {
        let episode_length = current_step - self.episode_start_step;
        self.episode_lengths.push(episode_length);

        // Count exit reason
        match barrier_label {
            Some(1) => self.barrier_exit_counts[0] += 1,  // Profit target
            Some(-1) => self.barrier_exit_counts[1] += 1, // Stop loss
            Some(0) => self.barrier_exit_counts[2] += 1,  // Time expiry
            _ => self.barrier_exit_counts[3] += 1,        // Data/time boundary
        }

        // Reset for next episode
        self.episode_start_step = current_step + 1;
    }

    /// Get episode statistics
    pub(crate) fn get_episode_stats(&self) -> (f64, f64, usize, usize, [usize; 4]) {
        if self.episode_lengths.is_empty() {
            return (0.0, 0.0, 0, 0, [0; 4]);
        }

        let total = self.episode_lengths.len();
        let mean = self.episode_lengths.iter().sum::<usize>() as f64 / total as f64;
        let min = self.episode_lengths.iter().min().copied().unwrap_or(0);
        let max = self.episode_lengths.iter().max().copied().unwrap_or(0);

        // Calculate std dev
        let variance = self.episode_lengths.iter()
            .map(|&len| {
                let diff = len as f64 - mean;
                diff * diff
            })
            .sum::<f64>() / total as f64;
        let std_dev = variance.sqrt();

        (mean, std_dev, min, max, self.barrier_exit_counts)
    }
}
