//! Optimized Logging Utilities for DQN Training
//!
//! This module provides structured, performance-optimized logging for ML training.
//! Follows Rust best practices with proper log levels and sampling.
//!
//! ## Log Level Guidelines
//! - **error!**: Unrecoverable failures requiring immediate attention
//! - **warn!**: Degraded state but training continues
//! - **info!**: Major milestones only (epoch start/end, checkpoints)
//! - **debug!**: Diagnostic details for development
//! - **trace!**: High-frequency internals (disabled in production)
//!
//! ## Performance Considerations
//! - Trace/debug have 200-1000ns overhead per call
//! - Use sampling for high-frequency events (every N batches)
//! - Aggregate metrics over windows, don't log every step

use std::time::{Duration, Instant};
use tracing::{debug, info, trace, warn};

/// Configuration for training logging behavior
#[derive(Debug, Clone)]
pub struct LoggingConfig {
    /// Log training metrics every N batches (default: 100)
    pub batch_log_interval: usize,
    /// Log epoch summaries (default: true)
    pub log_epoch_summaries: bool,
    /// Enable trace-level logging for forward passes (default: false)
    pub trace_forward_pass: bool,
    /// Log gradient statistics every N steps (default: 500)
    pub gradient_log_interval: usize,
    /// Minimum duration between metric logs to prevent spam (default: 1s)
    pub min_log_interval: Duration,
    /// Enable network diagnostics logging (default: false)
    pub enable_network_diagnostics: bool,
}

impl Default for LoggingConfig {
    fn default() -> Self {
        Self {
            batch_log_interval: 100,
            log_epoch_summaries: true,
            trace_forward_pass: false,
            gradient_log_interval: 500,
            min_log_interval: Duration::from_secs(1),
            enable_network_diagnostics: false,
        }
    }
}

/// Metrics aggregator for windowed training statistics
#[derive(Debug, Default)]
pub struct MetricsAggregator {
    losses: Vec<f32>,
    q_values: Vec<f32>,
    rewards: Vec<f32>,
    gradient_norms: Vec<f32>,
    last_log_time: Option<Instant>,
    batch_count: usize,
}

impl MetricsAggregator {
    /// Create a new metrics aggregator
    pub fn new() -> Self {
        Self::default()
    }

    /// Record a training step's metrics
    pub fn record(&mut self, loss: f32, q_value: f32, reward: f32) {
        self.losses.push(loss);
        self.q_values.push(q_value);
        self.rewards.push(reward);
        self.batch_count += 1;
    }

    /// Record gradient norm
    pub fn record_gradient(&mut self, norm: f32) {
        self.gradient_norms.push(norm);
    }

    /// Get the current batch count
    pub const fn batch_count(&self) -> usize {
        self.batch_count
    }

    /// Check if we should log based on interval and time
    pub fn should_log(&mut self, config: &LoggingConfig) -> bool {
        if self.batch_count % config.batch_log_interval != 0 {
            return false;
        }

        // Rate limiting check
        if let Some(last_time) = self.last_log_time {
            if last_time.elapsed() < config.min_log_interval {
                return false;
            }
        }

        self.last_log_time = Some(Instant::now());
        true
    }

    /// Get aggregated statistics and clear the buffer
    pub fn aggregate_and_clear(&mut self) -> AggregatedMetrics {
        let metrics = AggregatedMetrics {
            mean_loss: mean(&self.losses),
            std_loss: std_dev(&self.losses),
            mean_q_value: mean(&self.q_values),
            std_q_value: std_dev(&self.q_values),
            mean_reward: mean(&self.rewards),
            mean_gradient_norm: mean(&self.gradient_norms),
            batch_count: self.batch_count,
        };

        // Clear buffers but keep batch count
        self.losses.clear();
        self.q_values.clear();
        self.rewards.clear();
        self.gradient_norms.clear();

        metrics
    }

    /// Reset the aggregator completely
    pub fn reset(&mut self) {
        self.losses.clear();
        self.q_values.clear();
        self.rewards.clear();
        self.gradient_norms.clear();
        self.last_log_time = None;
        self.batch_count = 0;
    }
}

/// Aggregated metrics from a window of training steps
#[derive(Debug, Clone)]
pub struct AggregatedMetrics {
    pub mean_loss: f32,
    pub std_loss: f32,
    pub mean_q_value: f32,
    pub std_q_value: f32,
    pub mean_reward: f32,
    pub mean_gradient_norm: f32,
    pub batch_count: usize,
}

/// Sampled logger for high-frequency events
#[derive(Debug)]
pub struct SampledLogger {
    sample_rate: usize,
    counter: usize,
}

impl SampledLogger {
    /// Create a new sampled logger that logs every N calls
    pub fn new(sample_rate: usize) -> Self {
        Self {
            sample_rate: sample_rate.max(1),
            counter: 0,
        }
    }

    /// Check if this call should be logged
    pub const fn should_log(&mut self) -> bool {
        self.counter += 1;
        if self.counter >= self.sample_rate {
            self.counter = 0;
            true
        } else {
            false
        }
    }

    /// Get the current sample rate
    pub const fn sample_rate(&self) -> usize {
        self.sample_rate
    }

    /// Reset the counter
    pub const fn reset(&mut self) {
        self.counter = 0;
    }
}

// ============================================================================
// Structured Logging Functions
// ============================================================================

/// Log epoch start (info level - major milestone)
pub fn log_epoch_start(epoch: usize, total_epochs: usize, learning_rate: f64) {
    info!(
        epoch = epoch,
        total_epochs = total_epochs,
        learning_rate = %format!("{:.2e}", learning_rate),
        "Starting epoch"
    );
}

/// Log epoch end with summary (info level - major milestone)
pub fn log_epoch_end(epoch: usize, metrics: &AggregatedMetrics, duration_secs: f64) {
    info!(
        epoch = epoch,
        mean_loss = %format!("{:.6}", metrics.mean_loss),
        mean_q_value = %format!("{:.4}", metrics.mean_q_value),
        mean_reward = %format!("{:.4}", metrics.mean_reward),
        batches = metrics.batch_count,
        duration_secs = %format!("{:.2}", duration_secs),
        "Epoch complete"
    );
}

/// Log training progress (debug level - diagnostic)
pub fn log_training_progress(metrics: &AggregatedMetrics) {
    debug!(
        mean_loss = %format!("{:.6}", metrics.mean_loss),
        std_loss = %format!("{:.6}", metrics.std_loss),
        mean_q = %format!("{:.4}", metrics.mean_q_value),
        std_q = %format!("{:.4}", metrics.std_q_value),
        batches = metrics.batch_count,
        "Training progress"
    );
}

/// Log gradient statistics (debug level - diagnostic)
pub fn log_gradient_stats(norm: f32, max_norm: f32, clip_threshold: f32) {
    debug!(
        gradient_norm = %format!("{:.4}", norm),
        max_norm = %format!("{:.4}", max_norm),
        clip_threshold = %format!("{:.2}", clip_threshold),
        clipped = norm > clip_threshold,
        "Gradient statistics"
    );
}

/// Log network forward pass (trace level - high frequency)
pub fn log_forward_pass(batch_size: usize, input_dim: usize, output_dim: usize) {
    trace!(
        batch_size = batch_size,
        input_dim = input_dim,
        output_dim = output_dim,
        "Network forward pass"
    );
}

/// Log replay buffer sampling (trace level - high frequency)
pub fn log_replay_sample(buffer_size: usize, batch_size: usize, beta: f32) {
    trace!(
        buffer_size = buffer_size,
        batch_size = batch_size,
        beta = %format!("{:.3}", beta),
        "Replay buffer sample"
    );
}

/// Log target network update (debug level - periodic)
pub fn log_target_update(update_type: &str, tau: f64) {
    debug!(
        update_type = update_type,
        tau = %format!("{:.4}", tau),
        "Target network update"
    );
}

/// Log early stopping trigger (warn level - degraded state)
pub fn log_early_stopping(reason: &str, epoch: usize, metric_value: f32) {
    warn!(
        reason = reason,
        epoch = epoch,
        metric_value = %format!("{:.6}", metric_value),
        "Early stopping triggered"
    );
}

/// Log checkpoint save (info level - major milestone)
pub fn log_checkpoint_save(path: &str, epoch: usize, val_loss: f32) {
    info!(
        path = path,
        epoch = epoch,
        val_loss = %format!("{:.6}", val_loss),
        "Checkpoint saved"
    );
}

/// Log Q-value statistics (debug level - diagnostic)
pub fn log_q_value_stats(mean: f32, std: f32, min: f32, max: f32) {
    debug!(
        mean = %format!("{:.4}", mean),
        std = %format!("{:.4}", std),
        min = %format!("{:.4}", min),
        max = %format!("{:.4}", max),
        "Q-value statistics"
    );
}

// ============================================================================
// Helper Functions
// ============================================================================

fn mean(values: &[f32]) -> f32 {
    if values.is_empty() {
        return 0.0;
    }
    values.iter().sum::<f32>() / values.len() as f32
}

fn std_dev(values: &[f32]) -> f32 {
    if values.len() < 2 {
        return 0.0;
    }
    let m = mean(values);
    let variance = values.iter().map(|x| (x - m).powi(2)).sum::<f32>() / (values.len() - 1) as f32;
    variance.sqrt()
}

// ============================================================================
// Tests
// ============================================================================

#[cfg(test)]
mod tests {
    use super::*;

    // ========================================================================
    // LoggingConfig Tests
    // ========================================================================

    #[test]
    fn test_logging_config_default() {
        let config = LoggingConfig::default();
        assert_eq!(config.batch_log_interval, 100);
        assert!(config.log_epoch_summaries);
        assert!(!config.trace_forward_pass);
        assert_eq!(config.gradient_log_interval, 500);
        assert_eq!(config.min_log_interval, Duration::from_secs(1));
        assert!(!config.enable_network_diagnostics);
    }

    #[test]
    fn test_logging_config_custom() {
        let config = LoggingConfig {
            batch_log_interval: 50,
            log_epoch_summaries: false,
            trace_forward_pass: true,
            gradient_log_interval: 100,
            min_log_interval: Duration::from_millis(500),
            enable_network_diagnostics: true,
        };
        assert_eq!(config.batch_log_interval, 50);
        assert!(!config.log_epoch_summaries);
        assert!(config.trace_forward_pass);
        assert_eq!(config.gradient_log_interval, 100);
        assert_eq!(config.min_log_interval, Duration::from_millis(500));
        assert!(config.enable_network_diagnostics);
    }

    // ========================================================================
    // MetricsAggregator Tests
    // ========================================================================

    #[test]
    fn test_metrics_aggregator_new() {
        let agg = MetricsAggregator::new();
        assert_eq!(agg.batch_count(), 0);
    }

    #[test]
    fn test_metrics_aggregator_record() {
        let mut agg = MetricsAggregator::new();
        agg.record(0.5, 1.0, 0.1);
        assert_eq!(agg.batch_count(), 1);

        agg.record(0.4, 1.2, 0.2);
        assert_eq!(agg.batch_count(), 2);
    }

    #[test]
    fn test_metrics_aggregator_record_gradient() {
        let mut agg = MetricsAggregator::new();
        agg.record_gradient(1.5);
        agg.record_gradient(2.0);

        let metrics = agg.aggregate_and_clear();
        assert!((metrics.mean_gradient_norm - 1.75).abs() < 0.01);
    }

    #[test]
    fn test_metrics_aggregator_should_log_interval() {
        let mut agg = MetricsAggregator::new();
        let config = LoggingConfig {
            batch_log_interval: 10,
            min_log_interval: Duration::from_millis(0), // Disable rate limiting
            ..Default::default()
        };

        // Should not log before interval
        for _ in 0..9 {
            agg.record(0.5, 1.0, 0.1);
            assert!(!agg.should_log(&config));
        }

        // Should log at interval
        agg.record(0.5, 1.0, 0.1);
        assert!(agg.should_log(&config));
    }

    #[test]
    fn test_metrics_aggregator_aggregate_and_clear() {
        let mut agg = MetricsAggregator::new();
        agg.record(1.0, 2.0, 0.5);
        agg.record(2.0, 4.0, 1.0);
        agg.record(3.0, 6.0, 1.5);

        let metrics = agg.aggregate_and_clear();

        // Mean loss: (1 + 2 + 3) / 3 = 2.0
        assert!((metrics.mean_loss - 2.0).abs() < 0.01);
        // Mean Q-value: (2 + 4 + 6) / 3 = 4.0
        assert!((metrics.mean_q_value - 4.0).abs() < 0.01);
        // Mean reward: (0.5 + 1.0 + 1.5) / 3 = 1.0
        assert!((metrics.mean_reward - 1.0).abs() < 0.01);
        // Batch count preserved
        assert_eq!(metrics.batch_count, 3);

        // Internal buffers should be cleared
        let metrics2 = agg.aggregate_and_clear();
        assert_eq!(metrics2.mean_loss, 0.0);
        assert_eq!(metrics2.batch_count, 3); // Batch count not reset by aggregate_and_clear
    }

    #[test]
    fn test_metrics_aggregator_reset() {
        let mut agg = MetricsAggregator::new();
        agg.record(1.0, 2.0, 0.5);
        agg.record(2.0, 4.0, 1.0);

        agg.reset();

        assert_eq!(agg.batch_count(), 0);
        let metrics = agg.aggregate_and_clear();
        assert_eq!(metrics.mean_loss, 0.0);
        assert_eq!(metrics.batch_count, 0);
    }

    #[test]
    fn test_metrics_aggregator_std_dev() {
        let mut agg = MetricsAggregator::new();
        // Values: 1, 2, 3 - std dev should be 1.0 (sample std dev)
        agg.record(1.0, 0.0, 0.0);
        agg.record(2.0, 0.0, 0.0);
        agg.record(3.0, 0.0, 0.0);

        let metrics = agg.aggregate_and_clear();
        // Sample std dev of [1, 2, 3] = sqrt(((1-2)^2 + (2-2)^2 + (3-2)^2) / 2) = 1.0
        assert!((metrics.std_loss - 1.0).abs() < 0.01);
    }

    #[test]
    fn test_metrics_aggregator_empty() {
        let mut agg = MetricsAggregator::new();
        let metrics = agg.aggregate_and_clear();

        assert_eq!(metrics.mean_loss, 0.0);
        assert_eq!(metrics.std_loss, 0.0);
        assert_eq!(metrics.mean_q_value, 0.0);
        assert_eq!(metrics.batch_count, 0);
    }

    // ========================================================================
    // SampledLogger Tests
    // ========================================================================

    #[test]
    fn test_sampled_logger_new() {
        let logger = SampledLogger::new(100);
        assert_eq!(logger.sample_rate(), 100);
    }

    #[test]
    fn test_sampled_logger_new_zero_rate() {
        // Zero rate should be clamped to 1
        let logger = SampledLogger::new(0);
        assert_eq!(logger.sample_rate(), 1);
    }

    #[test]
    fn test_sampled_logger_should_log() {
        let mut logger = SampledLogger::new(5);

        // First 4 calls should not log
        assert!(!logger.should_log());
        assert!(!logger.should_log());
        assert!(!logger.should_log());
        assert!(!logger.should_log());

        // 5th call should log
        assert!(logger.should_log());

        // Next 4 calls should not log again
        assert!(!logger.should_log());
        assert!(!logger.should_log());
        assert!(!logger.should_log());
        assert!(!logger.should_log());

        // 10th call should log
        assert!(logger.should_log());
    }

    #[test]
    fn test_sampled_logger_reset() {
        let mut logger = SampledLogger::new(5);

        // Advance counter
        logger.should_log();
        logger.should_log();
        logger.should_log();

        // Reset
        logger.reset();

        // Should need 5 more calls to log
        assert!(!logger.should_log());
        assert!(!logger.should_log());
        assert!(!logger.should_log());
        assert!(!logger.should_log());
        assert!(logger.should_log());
    }

    // ========================================================================
    // AggregatedMetrics Tests
    // ========================================================================

    #[test]
    fn test_aggregated_metrics_clone() {
        let metrics = AggregatedMetrics {
            mean_loss: 0.5,
            std_loss: 0.1,
            mean_q_value: 2.0,
            std_q_value: 0.5,
            mean_reward: 0.1,
            mean_gradient_norm: 1.0,
            batch_count: 100,
        };

        let cloned = metrics.clone();
        assert_eq!(metrics.mean_loss, cloned.mean_loss);
        assert_eq!(metrics.batch_count, cloned.batch_count);
    }

    // ========================================================================
    // Helper Function Tests
    // ========================================================================

    #[test]
    fn test_mean_empty() {
        let values: Vec<f32> = vec![];
        assert_eq!(mean(&values), 0.0);
    }

    #[test]
    fn test_mean_single() {
        let values = vec![5.0];
        assert_eq!(mean(&values), 5.0);
    }

    #[test]
    fn test_mean_multiple() {
        let values = vec![1.0, 2.0, 3.0, 4.0, 5.0];
        assert_eq!(mean(&values), 3.0);
    }

    #[test]
    fn test_std_dev_empty() {
        let values: Vec<f32> = vec![];
        assert_eq!(std_dev(&values), 0.0);
    }

    #[test]
    fn test_std_dev_single() {
        let values = vec![5.0];
        assert_eq!(std_dev(&values), 0.0);
    }

    #[test]
    fn test_std_dev_multiple() {
        // [2, 4, 4, 4, 5, 5, 7, 9] - sample std dev = 2.138...
        let values = vec![2.0, 4.0, 4.0, 4.0, 5.0, 5.0, 7.0, 9.0];
        let result = std_dev(&values);
        assert!((result - 2.138).abs() < 0.01);
    }

    // ========================================================================
    // Integration Tests
    // ========================================================================

    #[test]
    fn test_logging_workflow() {
        // Simulate a training loop workflow
        let config = LoggingConfig {
            batch_log_interval: 10,
            min_log_interval: Duration::from_millis(0),
            ..Default::default()
        };

        let mut aggregator = MetricsAggregator::new();
        let mut forward_logger = SampledLogger::new(100);

        // Simulate 100 training steps
        for step in 0..100 {
            // Record metrics every step
            aggregator.record(0.5 - (step as f32 * 0.001), 1.0, 0.1);

            // Sample forward pass logging
            if forward_logger.should_log() {
                // Would call log_forward_pass here
            }

            // Aggregate and log every batch_log_interval
            if aggregator.should_log(&config) {
                let metrics = aggregator.aggregate_and_clear();
                // Would call log_training_progress here
                assert!(metrics.mean_loss > 0.0);
            }
        }

        assert_eq!(aggregator.batch_count(), 100);
    }

    #[test]
    fn test_log_level_appropriateness() {
        // Test that log functions compile and have appropriate signatures
        // This is a compile-time check more than runtime

        let metrics = AggregatedMetrics {
            mean_loss: 0.5,
            std_loss: 0.1,
            mean_q_value: 2.0,
            std_q_value: 0.5,
            mean_reward: 0.1,
            mean_gradient_norm: 1.0,
            batch_count: 100,
        };

        // Info level - major milestones
        log_epoch_start(1, 100, 0.001);
        log_epoch_end(1, &metrics, 60.0);
        log_checkpoint_save("/path/to/checkpoint", 1, 0.5);

        // Debug level - diagnostics
        log_training_progress(&metrics);
        log_gradient_stats(1.0, 2.0, 1.0);
        log_target_update("soft", 0.001);
        log_q_value_stats(1.0, 0.5, -1.0, 3.0);

        // Trace level - high frequency
        log_forward_pass(32, 54, 5);
        log_replay_sample(10000, 32, 0.4);

        // Warn level - degraded state
        log_early_stopping("gradient_collapse", 50, 0.001);
    }
}
