trustformers-optim 0.1.1

#![allow(clippy::all)]
//! Advanced Memory Optimization Analyzer for TrustformeRS Optimizers
//!
//! This tool provides comprehensive memory analysis with:
//! - Memory usage pattern detection
//! - Memory leak detection and prevention
//! - Buffer optimization recommendations
//! - Memory allocation efficiency analysis
//! - Memory fragmentation detection

#![allow(unused_imports, unused_variables, dead_code)]
//! - Optimizer-specific memory optimization strategies

use serde::{Deserialize, Serialize};
use std::collections::{HashMap, VecDeque};
use std::time::{Duration, Instant};
use trustformers_core::traits::Optimizer;
use trustformers_core::Tensor;
use trustformers_core::TrustformersError;
use trustformers_optim::*;

/// Memory usage sample at a specific point in time
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MemorySample {
    pub timestamp: Duration,
    pub total_memory_mb: f64,
    pub heap_memory_mb: f64,
    pub optimizer_state_mb: f64,
    pub parameter_memory_mb: f64,
    pub gradient_memory_mb: f64,
}

/// Memory optimization recommendations
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MemoryOptimizationReport {
    pub optimizer_name: String,
    pub total_memory_usage_mb: f64,
    pub peak_memory_usage_mb: f64,
    pub memory_efficiency_score: f64,
    pub memory_fragmentation_score: f64,
    pub potential_memory_savings_mb: f64,
    pub recommendations: Vec<String>,
    pub memory_patterns: Vec<MemoryPattern>,
}

/// Detected memory usage patterns
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum MemoryPattern {
    ConstantGrowth {
        rate_mb_per_iteration: f64,
    },
    MemoryLeaks {
        leak_rate_mb_per_iteration: f64,
    },
    MemorySpikes {
        spike_frequency: usize,
        avg_spike_size_mb: f64,
    },
    EfficientUsage {
        stability_score: f64,
    },
    BufferThrashing {
        thrash_frequency: usize,
    },
}

/// Advanced memory optimization analyzer
pub struct MemoryOptimizationAnalyzer {
    memory_samples: VecDeque<MemorySample>,
    max_samples: usize,
    start_time: Instant,
    optimization_reports: HashMap<String, MemoryOptimizationReport>,
    memory_threshold_mb: f64,
}

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

impl MemoryOptimizationAnalyzer {
    /// Create a new memory optimization analyzer
    pub fn new() -> Self {
        Self {
            memory_samples: VecDeque::new(),
            max_samples: 1000,
            start_time: Instant::now(),
            optimization_reports: HashMap::new(),
            memory_threshold_mb: 100.0, // Default threshold
        }
    }

    /// Configure memory analysis parameters
    pub fn with_threshold(mut self, threshold_mb: f64) -> Self {
        self.memory_threshold_mb = threshold_mb;
        self
    }

    /// Analyze memory usage patterns for an optimizer
    pub fn analyze_optimizer<T: Optimizer>(
        &mut self,
        optimizer_name: &str,
        mut optimizer: T,
        param_size: usize,
        iterations: usize,
    ) -> Result<MemoryOptimizationReport, TrustformersError> {
        println!(
            "🧠 Analyzing memory patterns for {} ({} params, {} iterations)...",
            optimizer_name, param_size, iterations
        );

        // Reset for new analysis
        self.memory_samples.clear();
        self.start_time = Instant::now();

        // Create test tensors
        let mut params = Tensor::randn(&[param_size])?;
        let gradients = Tensor::randn(&[param_size])?;

        // Initial memory sample
        self.record_memory_sample(&optimizer, &params, &gradients)?;

        // Run optimizer iterations with memory tracking
        for i in 0..iterations {
            optimizer.zero_grad();
            let _ = optimizer.update(&mut params, &gradients);
            optimizer.step();

            // Record memory sample every few iterations
            if i % 5 == 0 || i == iterations - 1 {
                self.record_memory_sample(&optimizer, &params, &gradients)?;
            }

            // Progress indicator
            if (i + 1) % (iterations / 10).max(1) == 0 {
                print!(".");
                std::io::Write::flush(&mut std::io::stdout()).expect("Failed to flush stdout");
            }
        }

        println!(" ✅ Complete!");

        // Analyze patterns and generate report
        let report = self.generate_optimization_report(optimizer_name)?;
        self.optimization_reports.insert(optimizer_name.to_string(), report.clone());

        Ok(report)
    }

    /// Record a memory usage sample
    fn record_memory_sample<T: Optimizer>(
        &mut self,
        optimizer: &T,
        params: &Tensor,
        gradients: &Tensor,
    ) -> Result<(), TrustformersError> {
        let timestamp = self.start_time.elapsed();

        // Simulate memory measurements (in practice, would use actual memory tracking)
        let param_data_size = params.data()?.len() * 4; // f32 = 4 bytes
        let grad_data_size = gradients.data()?.len() * 4;
        let optimizer_memory = self.estimate_optimizer_memory(params);

        let sample = MemorySample {
            timestamp,
            total_memory_mb: (param_data_size + grad_data_size + optimizer_memory) as f64
                / (1024.0 * 1024.0),
            heap_memory_mb: self.get_simulated_heap_memory(),
            optimizer_state_mb: optimizer_memory as f64 / (1024.0 * 1024.0),
            parameter_memory_mb: param_data_size as f64 / (1024.0 * 1024.0),
            gradient_memory_mb: grad_data_size as f64 / (1024.0 * 1024.0),
        };

        self.memory_samples.push_back(sample);

        // Keep only recent samples
        if self.memory_samples.len() > self.max_samples {
            self.memory_samples.pop_front();
        }

        Ok(())
    }

    /// Generate comprehensive memory optimization report
    fn generate_optimization_report(
        &self,
        optimizer_name: &str,
    ) -> Result<MemoryOptimizationReport, TrustformersError> {
        if self.memory_samples.is_empty() {
            return Err(TrustformersError::tensor_op_error(
                "No memory samples available",
                "memory_analysis",
            ));
        }

        let samples: Vec<_> = self.memory_samples.iter().collect();

        // Calculate basic statistics
        let total_memory_usage =
            samples.last().expect("Memory samples should not be empty").total_memory_mb;
        let peak_memory_usage = samples.iter().map(|s| s.total_memory_mb).fold(0.0, f64::max);

        // Analyze memory patterns
        let patterns = self.detect_memory_patterns(&samples);

        // Calculate memory efficiency
        let avg_memory =
            samples.iter().map(|s| s.total_memory_mb).sum::<f64>() / samples.len() as f64;
        let memory_variance =
            samples.iter().map(|s| (s.total_memory_mb - avg_memory).powi(2)).sum::<f64>()
                / samples.len() as f64;
        let memory_efficiency_score = 1.0 / (1.0 + memory_variance / avg_memory);

        // Calculate fragmentation score
        let fragmentation_score = self.calculate_fragmentation_score(&samples);

        // Generate recommendations
        let recommendations =
            self.generate_recommendations(optimizer_name, &patterns, memory_efficiency_score);

        // Estimate potential savings
        let potential_savings = self.estimate_memory_savings(&patterns, peak_memory_usage);

        Ok(MemoryOptimizationReport {
            optimizer_name: optimizer_name.to_string(),
            total_memory_usage_mb: total_memory_usage,
            peak_memory_usage_mb: peak_memory_usage,
            memory_efficiency_score,
            memory_fragmentation_score: fragmentation_score,
            potential_memory_savings_mb: potential_savings,
            recommendations,
            memory_patterns: patterns,
        })
    }

    /// Detect memory usage patterns
    fn detect_memory_patterns(&self, samples: &[&MemorySample]) -> Vec<MemoryPattern> {
        let mut patterns = Vec::new();

        if samples.len() < 10 {
            return patterns;
        }

        // Check for constant growth
        let growth_rates: Vec<f64> = samples
            .windows(2)
            .map(|w| (w[1].total_memory_mb - w[0].total_memory_mb) / w[1].timestamp.as_secs_f64())
            .collect();

        let avg_growth_rate = growth_rates.iter().sum::<f64>() / growth_rates.len() as f64;

        if avg_growth_rate > 0.1 {
            patterns.push(MemoryPattern::ConstantGrowth {
                rate_mb_per_iteration: avg_growth_rate,
            });
        }

        // Check for memory leaks (continuously increasing memory without release)
        let leak_threshold = 0.05; // MB per iteration
        if avg_growth_rate > leak_threshold {
            let leak_samples = samples.len() / 4; // Check last quarter
            let recent_samples = &samples[samples.len() - leak_samples..];
            let recent_growth = (recent_samples
                .last()
                .expect("Recent samples should not be empty")
                .total_memory_mb
                - recent_samples
                    .first()
                    .expect("Recent samples should not be empty")
                    .total_memory_mb)
                / leak_samples as f64;

            if recent_growth > leak_threshold {
                patterns.push(MemoryPattern::MemoryLeaks {
                    leak_rate_mb_per_iteration: recent_growth,
                });
            }
        }

        // Check for memory spikes
        let memory_values: Vec<f64> = samples.iter().map(|s| s.total_memory_mb).collect();
        let mean_memory = memory_values.iter().sum::<f64>() / memory_values.len() as f64;
        let std_dev = (memory_values.iter().map(|&x| (x - mean_memory).powi(2)).sum::<f64>()
            / memory_values.len() as f64)
            .sqrt();

        let spike_threshold = mean_memory + 2.0 * std_dev;
        let spikes: Vec<_> = memory_values
            .iter()
            .enumerate()
            .filter(|(_, &value)| value > spike_threshold)
            .collect();

        if !spikes.is_empty() {
            let avg_spike_size = spikes.iter().map(|(_, &value)| value - mean_memory).sum::<f64>()
                / spikes.len() as f64;

            patterns.push(MemoryPattern::MemorySpikes {
                spike_frequency: spikes.len(),
                avg_spike_size_mb: avg_spike_size,
            });
        }

        // Check for efficient usage (low variance)
        if std_dev / mean_memory < 0.1 && avg_growth_rate < 0.01 {
            let stability_score = 1.0 - (std_dev / mean_memory);
            patterns.push(MemoryPattern::EfficientUsage { stability_score });
        }

        // Check for buffer thrashing (frequent large changes)
        let large_changes = memory_values
            .windows(2)
            .map(|w| (w[1] - w[0]).abs())
            .filter(|&change| change > mean_memory * 0.1)
            .count();

        if large_changes > samples.len() / 4 {
            patterns.push(MemoryPattern::BufferThrashing {
                thrash_frequency: large_changes,
            });
        }

        patterns
    }

    /// Calculate memory fragmentation score
    fn calculate_fragmentation_score(&self, samples: &[&MemorySample]) -> f64 {
        if samples.len() < 2 {
            return 0.0;
        }

        // Simulate fragmentation based on memory allocation patterns
        let memory_changes: Vec<f64> = samples
            .windows(2)
            .map(|w| (w[1].total_memory_mb - w[0].total_memory_mb).abs())
            .collect();

        let avg_change = memory_changes.iter().sum::<f64>() / memory_changes.len() as f64;
        let total_memory = samples.last().expect("Samples should not be empty").total_memory_mb;

        // Higher ratio indicates more fragmentation
        avg_change / total_memory.max(1.0)
    }

    /// Generate optimization recommendations
    fn generate_recommendations(
        &self,
        optimizer_name: &str,
        patterns: &[MemoryPattern],
        efficiency_score: f64,
    ) -> Vec<String> {
        let mut recommendations = Vec::new();

        for pattern in patterns {
            match pattern {
                MemoryPattern::ConstantGrowth {
                    rate_mb_per_iteration,
                } => {
                    recommendations.push(format!(
                        "⚠️  Detected constant memory growth ({:.3} MB/iter). Consider using gradient accumulation or memory-efficient variants.",
                        rate_mb_per_iteration
                    ));
                },
                MemoryPattern::MemoryLeaks {
                    leak_rate_mb_per_iteration,
                } => {
                    recommendations.push(format!(
                        "🚨 Potential memory leak detected ({:.3} MB/iter). Review tensor lifecycle and implement proper cleanup.",
                        leak_rate_mb_per_iteration
                    ));
                },
                MemoryPattern::MemorySpikes {
                    spike_frequency,
                    avg_spike_size_mb,
                } => {
                    recommendations.push(format!(
                        "📈 Memory spikes detected ({} spikes, avg {:.1} MB). Consider pre-allocating buffers or using memory pools.",
                        spike_frequency, avg_spike_size_mb
                    ));
                },
                MemoryPattern::EfficientUsage { stability_score } => {
                    recommendations.push(format!(
                        "✅ Efficient memory usage detected (stability: {:.3}). Current configuration is well-optimized.",
                        stability_score
                    ));
                },
                MemoryPattern::BufferThrashing { thrash_frequency } => {
                    recommendations.push(format!(
                        "🔄 Buffer thrashing detected ({} large changes). Consider using persistent buffers or batch processing.",
                        thrash_frequency
                    ));
                },
            }
        }

        // Optimizer-specific recommendations
        match optimizer_name {
            "Adam" | "AdamW" => {
                if efficiency_score < 0.7 {
                    recommendations.push("💡 For Adam/AdamW: Consider using 8-bit optimizers (Adam8bit) for memory savings.".to_string());
                }
            },
            "BGE-Adam" => {
                recommendations.push(
                    "🚀 BGE-Adam: Use OptimizedBGEAdam for 3-5x better memory efficiency."
                        .to_string(),
                );
            },
            "SGD" => {
                recommendations.push("💾 SGD: Already memory-efficient. Consider Nesterov momentum for better convergence.".to_string());
            },
            _ => {},
        }

        // General recommendations based on efficiency
        if efficiency_score < 0.5 {
            recommendations.push("🔧 Low memory efficiency detected. Consider CPU offloading or gradient checkpointing.".to_string());
        }

        if recommendations.is_empty() {
            recommendations.push(
                "🎉 No major memory optimization issues detected. Current usage is efficient."
                    .to_string(),
            );
        }

        recommendations
    }

    /// Estimate potential memory savings
    fn estimate_memory_savings(&self, patterns: &[MemoryPattern], peak_memory_mb: f64) -> f64 {
        let mut total_savings = 0.0;

        for pattern in patterns {
            match pattern {
                MemoryPattern::MemorySpikes {
                    avg_spike_size_mb, ..
                } => {
                    total_savings += avg_spike_size_mb * 0.8; // 80% of spike size recoverable
                },
                MemoryPattern::BufferThrashing { .. } => {
                    total_savings += peak_memory_mb * 0.15; // 15% savings from buffer optimization
                },
                MemoryPattern::ConstantGrowth {
                    rate_mb_per_iteration,
                } => {
                    total_savings += rate_mb_per_iteration * 10.0; // Savings over 10 iterations
                },
                _ => {},
            }
        }

        total_savings.min(peak_memory_mb * 0.5) // Cap at 50% of peak memory
    }

    /// Generate comprehensive memory analysis report
    pub fn generate_report(&self) -> String {
        let mut report = String::new();
        report.push_str("🧠 TrustformeRS Memory Optimization Analysis\n");
        report.push_str("========================================\n\n");

        if self.optimization_reports.is_empty() {
            report.push_str("No analysis results available.\n");
            return report;
        }

        // Summary table
        report.push_str("📊 Memory Usage Summary\n");
        report.push_str("-----------------------\n");
        report.push_str(&format!(
            "{:<15} {:<12} {:<12} {:<15} {:<18} {:<15}\n",
            "Optimizer", "Total MB", "Peak MB", "Efficiency", "Fragmentation", "Savings MB"
        ));
        report.push_str(&format!("{}\n", "─".repeat(90)));

        let mut sorted_reports: Vec<_> = self.optimization_reports.values().collect();
        sorted_reports.sort_by(|a, b| {
            b.memory_efficiency_score
                .partial_cmp(&a.memory_efficiency_score)
                .expect("Cannot compare NaN values in memory efficiency scores")
        });

        for report_data in &sorted_reports {
            report.push_str(&format!(
                "{:<15} {:<12.1} {:<12.1} {:<15.3} {:<18.3} {:<15.1}\n",
                report_data.optimizer_name,
                report_data.total_memory_usage_mb,
                report_data.peak_memory_usage_mb,
                report_data.memory_efficiency_score,
                report_data.memory_fragmentation_score,
                report_data.potential_memory_savings_mb,
            ));
        }

        // Detailed recommendations
        report.push_str("\n💡 Optimization Recommendations\n");
        report.push_str("-------------------------------\n");

        for report_data in &sorted_reports {
            report.push_str(&format!(
                "\n📋 {} Recommendations:\n",
                report_data.optimizer_name
            ));
            for rec in &report_data.recommendations {
                report.push_str(&format!("   {}\n", rec));
            }
        }

        // Memory patterns analysis
        report.push_str("\n🔍 Memory Pattern Analysis\n");
        report.push_str("--------------------------\n");

        for report_data in &sorted_reports {
            if !report_data.memory_patterns.is_empty() {
                report.push_str(&format!("\n📈 {} Patterns:\n", report_data.optimizer_name));
                for pattern in &report_data.memory_patterns {
                    report.push_str(&format!("   {:?}\n", pattern));
                }
            }
        }

        report
    }

    /// Export detailed results to JSON
    pub fn export_json(&self, filename: &str) -> Result<(), Box<dyn std::error::Error>> {
        let json = serde_json::to_string_pretty(&self.optimization_reports)?;
        std::fs::write(filename, json)?;
        println!("📁 Memory analysis results exported to {}", filename);
        Ok(())
    }

    /// Simulate heap memory usage
    fn get_simulated_heap_memory(&self) -> f64 {
        // Simulate heap memory based on current time for demo purposes
        60.0 + (std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .expect("SystemTime is before UNIX EPOCH")
            .as_millis()
            % 50) as f64
            * 0.2
    }

    /// Estimate optimizer memory usage based on parameter size
    fn estimate_optimizer_memory(&self, params: &Tensor) -> usize {
        // Simulate optimizer memory usage (typically 2x parameter size for Adam-like optimizers)
        let param_size = params.data().map_or(0, |data| data.len() * 4); // f32 = 4 bytes
        param_size * 2 // Momentum + velocity buffers
    }
}

/// Main memory analysis function
fn main() -> Result<(), TrustformersError> {
    println!("🧠 TrustformeRS Memory Optimization Analyzer");
    println!("===========================================");

    let mut analyzer = MemoryOptimizationAnalyzer::new().with_threshold(50.0);

    let param_size = 5000;
    let iterations = 100;

    println!("\n🎯 Analyzing memory patterns for optimizers");
    println!("{}", "═".repeat(45));

    // Analyze different optimizers
    let _ = analyzer.analyze_optimizer(
        "Adam",
        Adam::new(0.001, (0.9, 0.999), 1e-8, 0.0),
        param_size,
        iterations,
    )?;

    let _ = analyzer.analyze_optimizer(
        "AdamW",
        AdamW::new(0.001, (0.9, 0.999), 1e-8, 0.01),
        param_size,
        iterations,
    )?;

    let _ = analyzer.analyze_optimizer(
        "SGD",
        SGD::new(0.01, 0.9, 0.0, false),
        param_size,
        iterations,
    )?;

    let _ = analyzer.analyze_optimizer(
        "BGE-Adam",
        BGEAdam::new(0.001, (0.9, 0.999), 1e-8, 0.01, 0.1, 0.05, 0.05),
        param_size,
        iterations,
    )?;

    // Generate and display report
    let report = analyzer.generate_report();
    println!("\n{}", report);

    // Export results
    if let Err(e) = analyzer.export_json("memory_analysis_results.json") {
        println!("⚠️  Warning: Could not export results to JSON: {}", e);
    }

    println!("\n✅ Memory optimization analysis complete!");

    Ok(())
}

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

    #[test]
    fn test_analyzer_creation() {
        let analyzer = MemoryOptimizationAnalyzer::new();
        assert_eq!(analyzer.memory_threshold_mb, 100.0);
        assert!(analyzer.memory_samples.is_empty());
    }

    #[test]
    fn test_analyzer_configuration() {
        let analyzer = MemoryOptimizationAnalyzer::new().with_threshold(200.0);
        assert_eq!(analyzer.memory_threshold_mb, 200.0);
    }

    #[test]
    fn test_pattern_detection() {
        let analyzer = MemoryOptimizationAnalyzer::new();

        // Create sample data
        let samples = vec![
            MemorySample {
                timestamp: Duration::from_secs(0),
                total_memory_mb: 10.0,
                heap_memory_mb: 5.0,
                optimizer_state_mb: 2.0,
                parameter_memory_mb: 2.0,
                gradient_memory_mb: 1.0,
            },
            MemorySample {
                timestamp: Duration::from_secs(1),
                total_memory_mb: 15.0,
                heap_memory_mb: 7.0,
                optimizer_state_mb: 3.0,
                parameter_memory_mb: 3.0,
                gradient_memory_mb: 2.0,
            },
            MemorySample {
                timestamp: Duration::from_secs(2),
                total_memory_mb: 20.0,
                heap_memory_mb: 9.0,
                optimizer_state_mb: 4.0,
                parameter_memory_mb: 4.0,
                gradient_memory_mb: 3.0,
            },
        ];

        let sample_refs: Vec<_> = samples.iter().collect();
        let patterns = analyzer.detect_memory_patterns(&sample_refs);

        // Should detect constant growth
        assert!(!patterns.is_empty());
    }
}