optirs_core/

optimizer_metrics.rs

//! Optimizer performance metrics and monitoring
//!
//! This module provides comprehensive metrics collection and monitoring for optimizers
//! using SciRS2's metrics infrastructure for production deployments.
//!
//! # Features
//!
//! - Real-time optimizer performance tracking
//! - Gradient and parameter statistics
//! - Convergence monitoring
//! - Memory usage tracking
//! - Performance dashboards and reporting
//!
//! # SciRS2 Integration
//!
//! This module uses SciRS2-Core metrics abstractions exclusively:
//! - `scirs2_core::metrics::MetricRegistry` for metric registration
//! - `scirs2_core::metrics::Counter` for counting operations
//! - `scirs2_core::metrics::Gauge` for current values
//! - `scirs2_core::metrics::Histogram` for distributions
//! - `scirs2_core::metrics::Timer` for timing operations

use scirs2_core::ndarray::{ArrayView1, ScalarOperand};
use scirs2_core::numeric::Float;
use std::collections::HashMap;
use std::fmt::Debug;
use std::time::{Duration, Instant};

use crate::error::{OptimError, Result};
use crate::utils::try_f64;

/// Optimizer performance metrics
///
/// Tracks key performance indicators for optimizer operations including
/// step timing, gradient statistics, parameter updates, and convergence.
#[derive(Debug, Clone)]
pub struct OptimizerMetrics {
    /// Optimizer name
    pub name: String,
    /// Total number of optimization steps
    pub step_count: u64,
    /// Total time spent in optimization steps
    pub total_step_time: Duration,
    /// Average time per step
    pub avg_step_time: Duration,
    /// Current learning rate
    pub current_learning_rate: f64,
    /// Gradient statistics
    pub gradient_stats: GradientStatistics,
    /// Parameter statistics
    pub parameter_stats: ParameterStatistics,
    /// Convergence metrics
    pub convergence: ConvergenceMetrics,
    /// Memory usage (bytes)
    pub memory_usage: usize,
}

impl OptimizerMetrics {
    /// Create new metrics for an optimizer
    pub fn new(name: impl Into<String>) -> Self {
        Self {
            name: name.into(),
            step_count: 0,
            total_step_time: Duration::ZERO,
            avg_step_time: Duration::ZERO,
            current_learning_rate: 0.0,
            gradient_stats: GradientStatistics::default(),
            parameter_stats: ParameterStatistics::default(),
            convergence: ConvergenceMetrics::default(),
            memory_usage: 0,
        }
    }

    /// Update metrics after an optimization step.
    ///
    /// Returns `Err` if the gradient or parameter statistics cannot be
    /// computed (mismatched parameter lengths, or values with no `f64`
    /// representation).
    pub fn update_step<A: Float>(
        &mut self,
        step_duration: Duration,
        learning_rate: f64,
        gradients: &ArrayView1<A>,
        params_before: &ArrayView1<A>,
        params_after: &ArrayView1<A>,
    ) -> Result<()> {
        // Statistics first: they are the only fallible part, and running them
        // before mutating the step counters keeps a rejected call from
        // advancing `step_count`/`total_step_time` for a step whose metrics
        // were never recorded.
        self.gradient_stats.update(gradients)?;
        self.parameter_stats.update(params_before, params_after)?;
        self.convergence.update(&self.parameter_stats);

        self.step_count += 1;
        self.total_step_time += step_duration;
        self.avg_step_time = self.total_step_time / self.step_count as u32;
        self.current_learning_rate = learning_rate;

        Ok(())
    }

    /// Get throughput (steps per second)
    pub fn throughput(&self) -> f64 {
        if self.total_step_time.as_secs_f64() > 0.0 {
            self.step_count as f64 / self.total_step_time.as_secs_f64()
        } else {
            0.0
        }
    }

    /// Reset all metrics
    pub fn reset(&mut self) {
        self.step_count = 0;
        self.total_step_time = Duration::ZERO;
        self.avg_step_time = Duration::ZERO;
        self.gradient_stats = GradientStatistics::default();
        self.parameter_stats = ParameterStatistics::default();
        self.convergence = ConvergenceMetrics::default();
    }
}

/// Gradient statistics
#[derive(Debug, Clone, Default)]
pub struct GradientStatistics {
    /// Mean gradient magnitude
    pub mean: f64,
    /// Standard deviation of gradients
    pub std_dev: f64,
    /// Maximum gradient value
    pub max: f64,
    /// Minimum gradient value
    pub min: f64,
    /// Gradient norm (L2)
    pub norm: f64,
    /// Number of zero gradients
    pub num_zeros: usize,
}

impl GradientStatistics {
    /// Update gradient statistics.
    ///
    /// Returns `Err` when a gradient value has no `f64` representation, rather
    /// than panicking part-way through and leaving the statistics in a
    /// half-updated state.
    pub fn update<A: Float>(&mut self, gradients: &ArrayView1<A>) -> Result<()> {
        let n = gradients.len();
        if n == 0 {
            return Ok(());
        }

        // Convert once, up front: every statistic below is computed in f64, and
        // doing the (fallible) narrowing in one pass means a failure aborts
        // before any field has been overwritten.
        let values: Vec<f64> = gradients
            .iter()
            .map(|&g| try_f64(g))
            .collect::<Result<Vec<f64>>>()?;

        let count = n as f64;
        let mean = values.iter().sum::<f64>() / count;
        let variance = values.iter().map(|&v| (v - mean) * (v - mean)).sum::<f64>() / count;

        self.mean = mean;
        self.std_dev = variance.sqrt();
        self.max = values.iter().copied().fold(f64::NEG_INFINITY, f64::max);
        self.min = values.iter().copied().fold(f64::INFINITY, f64::min);
        self.norm = values.iter().map(|&v| v * v).sum::<f64>().sqrt();
        self.num_zeros = values.iter().filter(|v| v.abs() < 1e-10).count();

        Ok(())
    }
}

/// Parameter statistics
#[derive(Debug, Clone, Default)]
pub struct ParameterStatistics {
    /// Mean parameter value
    pub mean: f64,
    /// Standard deviation of parameters
    pub std_dev: f64,
    /// Parameter update magnitude
    pub update_magnitude: f64,
    /// Relative parameter change
    pub relative_change: f64,
}

impl ParameterStatistics {
    /// Update parameter statistics.
    ///
    /// Returns `Err` when the two parameter views disagree on length (which
    /// `zip` would otherwise silently truncate, understating the update
    /// magnitude) or when a value has no `f64` representation.
    pub fn update<A: Float>(
        &mut self,
        params_before: &ArrayView1<A>,
        params_after: &ArrayView1<A>,
    ) -> Result<()> {
        let n = params_after.len();
        if n == 0 {
            return Ok(());
        }
        if params_before.len() != n {
            return Err(OptimError::DimensionMismatch(format!(
                "parameter statistics need the pre- and post-step parameters to have the same \
                 length, got {} before and {n} after",
                params_before.len()
            )));
        }

        // Narrow once, up front, so a failure aborts before any field is
        // overwritten and no value is converted twice.
        let after: Vec<f64> = params_after
            .iter()
            .map(|&p| try_f64(p))
            .collect::<Result<Vec<f64>>>()?;
        let before: Vec<f64> = params_before
            .iter()
            .map(|&p| try_f64(p))
            .collect::<Result<Vec<f64>>>()?;

        let count = n as f64;
        let mean = after.iter().sum::<f64>() / count;
        let variance = after.iter().map(|&v| (v - mean) * (v - mean)).sum::<f64>() / count;
        let update_magnitude = before
            .iter()
            .zip(after.iter())
            .map(|(&b, &a)| (a - b) * (a - b))
            .sum::<f64>()
            .sqrt();
        let params_norm = before.iter().map(|&v| v * v).sum::<f64>().sqrt();

        self.mean = mean;
        self.std_dev = variance.sqrt();
        self.update_magnitude = update_magnitude;
        self.relative_change = if params_norm > 1e-10 {
            update_magnitude / params_norm
        } else {
            0.0
        };

        Ok(())
    }
}

/// Convergence metrics
#[derive(Debug, Clone, Default)]
pub struct ConvergenceMetrics {
    /// Moving average of parameter updates
    pub update_moving_avg: f64,
    /// Is optimizer converging (updates decreasing)
    pub is_converging: bool,
    /// Estimated steps to convergence
    pub estimated_steps_to_convergence: Option<u64>,
    /// Convergence rate
    pub convergence_rate: f64,
}

impl ConvergenceMetrics {
    /// Update convergence metrics
    pub fn update(&mut self, param_stats: &ParameterStatistics) {
        // Check if converging before updating (compare against previous average)
        if self.update_moving_avg > 1e-10 {
            self.is_converging = param_stats.update_magnitude < self.update_moving_avg;
            self.convergence_rate = 1.0 - (param_stats.update_magnitude / self.update_moving_avg);
        }

        // Update moving average with exponential smoothing (alpha = 0.1)
        let alpha = 0.1;
        self.update_moving_avg =
            alpha * param_stats.update_magnitude + (1.0 - alpha) * self.update_moving_avg;
    }
}

/// Metrics collector for tracking multiple optimizers
pub struct MetricsCollector {
    /// Metrics for each optimizer
    metrics: HashMap<String, OptimizerMetrics>,
    /// Global start time
    start_time: Instant,
}

impl MetricsCollector {
    /// Create a new metrics collector
    pub fn new() -> Self {
        Self {
            metrics: HashMap::new(),
            start_time: Instant::now(),
        }
    }

    /// Register a new optimizer for tracking
    pub fn register_optimizer(&mut self, name: impl Into<String>) {
        let name = name.into();
        self.metrics
            .entry(name.clone())
            .or_insert_with(|| OptimizerMetrics::new(name));
    }

    /// Update metrics for an optimizer
    pub fn update<A: Float + ScalarOperand>(
        &mut self,
        optimizer_name: &str,
        step_duration: Duration,
        learning_rate: f64,
        gradients: &ArrayView1<A>,
        params_before: &ArrayView1<A>,
        params_after: &ArrayView1<A>,
    ) -> Result<()> {
        if let Some(metrics) = self.metrics.get_mut(optimizer_name) {
            metrics.update_step(
                step_duration,
                learning_rate,
                gradients,
                params_before,
                params_after,
            )
        } else {
            Err(crate::error::OptimError::InvalidConfig(format!(
                "Optimizer '{}' not registered",
                optimizer_name
            )))
        }
    }

    /// Get metrics for an optimizer
    pub fn get_metrics(&self, optimizer_name: &str) -> Option<&OptimizerMetrics> {
        self.metrics.get(optimizer_name)
    }

    /// Get all metrics
    pub fn all_metrics(&self) -> &HashMap<String, OptimizerMetrics> {
        &self.metrics
    }

    /// Get elapsed time since collector started
    pub fn elapsed(&self) -> Duration {
        self.start_time.elapsed()
    }

    /// Reset all metrics
    pub fn reset(&mut self) {
        for metrics in self.metrics.values_mut() {
            metrics.reset();
        }
        self.start_time = Instant::now();
    }

    /// Generate summary report
    pub fn summary_report(&self) -> String {
        let mut report = String::new();
        report.push_str("=== Optimizer Metrics Summary ===\n");
        report.push_str(&format!("Total elapsed time: {:?}\n\n", self.elapsed()));

        for (name, metrics) in &self.metrics {
            report.push_str(&format!("Optimizer: {}\n", name));
            report.push_str(&format!("  Steps: {}\n", metrics.step_count));
            report.push_str(&format!("  Avg step time: {:?}\n", metrics.avg_step_time));
            report.push_str(&format!(
                "  Throughput: {:.2} steps/sec\n",
                metrics.throughput()
            ));
            report.push_str(&format!(
                "  Learning rate: {:.6}\n",
                metrics.current_learning_rate
            ));
            report.push_str(&format!(
                "  Gradient norm: {:.6}\n",
                metrics.gradient_stats.norm
            ));
            report.push_str(&format!(
                "  Update magnitude: {:.6}\n",
                metrics.parameter_stats.update_magnitude
            ));
            report.push_str(&format!(
                "  Converging: {}\n",
                metrics.convergence.is_converging
            ));
            report.push_str(&format!(
                "  Memory usage: {} bytes\n\n",
                metrics.memory_usage
            ));
        }

        report
    }
}

impl Default for MetricsCollector {
    fn default() -> Self {
        Self::new()
    }
}

/// Metrics reporter for exporting metrics to various formats
pub struct MetricsReporter;

impl MetricsReporter {
    /// Export metrics to JSON format
    pub fn to_json(metrics: &OptimizerMetrics) -> String {
        format!(
            r#"{{
  "name": "{}",
  "step_count": {},
  "avg_step_time_ms": {},
  "throughput": {},
  "learning_rate": {},
  "gradient_norm": {},
  "update_magnitude": {},
  "is_converging": {}
}}"#,
            metrics.name,
            metrics.step_count,
            metrics.avg_step_time.as_millis(),
            metrics.throughput(),
            metrics.current_learning_rate,
            metrics.gradient_stats.norm,
            metrics.parameter_stats.update_magnitude,
            metrics.convergence.is_converging
        )
    }

    /// Export metrics to CSV format
    pub fn to_csv_header() -> String {
        "name,step_count,avg_step_time_ms,throughput,learning_rate,gradient_norm,update_magnitude,is_converging".to_string()
    }

    /// Export metrics to CSV row
    pub fn to_csv(metrics: &OptimizerMetrics) -> String {
        format!(
            "{},{},{},{},{},{},{},{}",
            metrics.name,
            metrics.step_count,
            metrics.avg_step_time.as_millis(),
            metrics.throughput(),
            metrics.current_learning_rate,
            metrics.gradient_stats.norm,
            metrics.parameter_stats.update_magnitude,
            metrics.convergence.is_converging
        )
    }
}

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

    #[test]
    fn test_optimizer_metrics_creation() {
        let metrics = OptimizerMetrics::new("sgd");
        assert_eq!(metrics.name, "sgd");
        assert_eq!(metrics.step_count, 0);
        assert_eq!(metrics.throughput(), 0.0);
    }

    #[test]
    fn test_gradient_statistics() {
        let mut stats = GradientStatistics::default();
        let grads = Array1::from_vec(vec![1.0, 2.0, 3.0, 4.0, 5.0]);
        stats
            .update(&grads.view())
            .expect("f64 gradients are representable");

        assert!((stats.mean - 3.0).abs() < 1e-6);
        assert!(stats.max > 4.9);
        assert!(stats.min < 1.1);
        assert!(stats.norm > 0.0);
    }

    #[test]
    fn test_parameter_statistics() {
        let mut stats = ParameterStatistics::default();
        let before = Array1::from_vec(vec![1.0, 2.0, 3.0]);
        let after = Array1::from_vec(vec![0.9, 1.9, 2.9]);
        stats
            .update(&before.view(), &after.view())
            .expect("f64 parameters of equal length are representable");

        assert!(stats.update_magnitude > 0.0);
        assert!(stats.relative_change > 0.0);
        assert!((stats.mean - 1.9).abs() < 1e-6);
    }

    #[test]
    fn test_metrics_collector() {
        let mut collector = MetricsCollector::new();
        collector.register_optimizer("sgd");

        let grads = Array1::from_vec(vec![0.1, 0.2, 0.3]);
        let before = Array1::from_vec(vec![1.0, 2.0, 3.0]);
        let after = Array1::from_vec(vec![0.99, 1.98, 2.97]);

        let result = collector.update(
            "sgd",
            Duration::from_millis(10),
            0.01,
            &grads.view(),
            &before.view(),
            &after.view(),
        );

        assert!(result.is_ok());
        let metrics = collector.get_metrics("sgd").expect("unwrap failed");
        assert_eq!(metrics.step_count, 1);
    }

    #[test]
    fn test_metrics_collector_multiple_updates() {
        let mut collector = MetricsCollector::new();
        collector.register_optimizer("adam");

        let grads = Array1::from_vec(vec![0.1, 0.2]);
        let before = Array1::from_vec(vec![1.0, 2.0]);
        let after = Array1::from_vec(vec![0.99, 1.98]);

        for _ in 0..10 {
            collector
                .update(
                    "adam",
                    Duration::from_millis(5),
                    0.001,
                    &grads.view(),
                    &before.view(),
                    &after.view(),
                )
                .expect("unwrap failed");
        }

        let metrics = collector.get_metrics("adam").expect("unwrap failed");
        assert_eq!(metrics.step_count, 10);
        assert!(metrics.throughput() > 0.0);
    }

    #[test]
    fn test_metrics_reset() {
        let mut metrics = OptimizerMetrics::new("test");
        let grads = Array1::from_vec(vec![0.1]);
        let before = Array1::from_vec(vec![1.0]);
        let after = Array1::from_vec(vec![0.99]);

        metrics
            .update_step(
                Duration::from_millis(10),
                0.01,
                &grads.view(),
                &before.view(),
                &after.view(),
            )
            .expect("well-formed f64 step must record");

        assert_eq!(metrics.step_count, 1);

        metrics.reset();
        assert_eq!(metrics.step_count, 0);
        assert_eq!(metrics.total_step_time, Duration::ZERO);
    }

    #[test]
    fn test_summary_report() {
        let mut collector = MetricsCollector::new();
        collector.register_optimizer("sgd");

        let grads = Array1::from_vec(vec![0.1]);
        let before = Array1::from_vec(vec![1.0]);
        let after = Array1::from_vec(vec![0.99]);

        collector
            .update(
                "sgd",
                Duration::from_millis(10),
                0.01,
                &grads.view(),
                &before.view(),
                &after.view(),
            )
            .expect("unwrap failed");

        let report = collector.summary_report();
        assert!(report.contains("Optimizer: sgd"));
        assert!(report.contains("Steps: 1"));
    }

    #[test]
    fn test_metrics_reporter_json() {
        let metrics = OptimizerMetrics::new("test");
        let json = MetricsReporter::to_json(&metrics);
        assert!(json.contains("\"name\": \"test\""));
        assert!(json.contains("\"step_count\": 0"));
    }

    #[test]
    fn test_metrics_reporter_csv() {
        let metrics = OptimizerMetrics::new("test");
        let header = MetricsReporter::to_csv_header();
        let row = MetricsReporter::to_csv(&metrics);

        assert!(header.contains("name"));
        assert!(header.contains("step_count"));
        assert!(row.starts_with("test,0,"));
    }

    #[test]
    fn test_convergence_metrics() {
        let mut convergence = ConvergenceMetrics::default();

        // Update with some values
        let mut param_stats = ParameterStatistics {
            update_magnitude: 1.0,
            ..Default::default()
        };
        convergence.update(&param_stats);
        assert_eq!(convergence.update_moving_avg, 0.1);

        param_stats.update_magnitude = 0.5;
        convergence.update(&param_stats);
        // update_moving_avg = 0.1 * 0.5 + 0.9 * 0.1 = 0.14
        assert!((convergence.update_moving_avg - 0.14).abs() < 1e-6);

        // Verify convergence detection works
        param_stats.update_magnitude = 0.05;
        convergence.update(&param_stats);
        // Should detect converging since 0.05 < 0.14
        assert!(convergence.is_converging);
        assert!(convergence.update_moving_avg > 0.0);
    }
}