sklears_neural/

visualization.rs

//! Model Visualization Utilities
//!
//! This module provides comprehensive visualization capabilities for neural networks,
//! including architecture diagrams, training metrics, attention heatmaps, and
//! weight distributions.

use crate::NeuralResult;
use scirs2_core::ndarray::{Array2, Array3};
use sklears_core::error::SklearsError;
use sklears_core::types::FloatBounds;
use std::collections::HashMap;
use std::fs::File;
use std::io::Write;

/// Configuration for visualization output
#[derive(Debug, Clone)]
pub struct VisualizationConfig {
    /// Output directory for visualizations
    pub output_dir: String,
    /// Image format (SVG, PNG, etc.)
    pub format: ImageFormat,
    /// Color scheme
    pub color_scheme: ColorScheme,
    /// DPI for raster formats
    pub dpi: u32,
    /// Whether to show layer names
    pub show_layer_names: bool,
    /// Whether to show tensor shapes
    pub show_tensor_shapes: bool,
}

impl Default for VisualizationConfig {
    fn default() -> Self {
        Self {
            output_dir: "./visualizations".to_string(),
            format: ImageFormat::SVG,
            color_scheme: ColorScheme::Default,
            dpi: 300,
            show_layer_names: true,
            show_tensor_shapes: true,
        }
    }
}

/// Supported image formats
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum ImageFormat {
    /// Scalable Vector Graphics — lossless, suitable for diagrams
    SVG,
    /// Portable Network Graphics — raster format (SVG is generated first, then converted externally)
    PNG,
    /// Interactive HTML with embedded JavaScript visualizations
    HTML,
}

/// Color schemes for visualizations
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum ColorScheme {
    /// Default color scheme
    Default,
    /// Perceptually uniform sequential color map
    Viridis,
    /// High-contrast sequential color map
    Plasma,
    /// Black-to-white grayscale ramp
    Grayscale,
}

/// Model architecture visualizer
pub struct ModelVisualizer {
    config: VisualizationConfig,
}

impl ModelVisualizer {
    /// Create a new model visualizer
    pub fn new(config: VisualizationConfig) -> Self {
        Self { config }
    }

    /// Generate a model architecture diagram
    pub fn visualize_architecture(&self, layers: &[LayerInfo], filename: &str) -> NeuralResult<()> {
        let output_path = format!(
            "{}/{}.{}",
            self.config.output_dir,
            filename,
            self.format_extension()
        );

        match self.config.format {
            ImageFormat::SVG => self.generate_svg_architecture(layers, &output_path),
            ImageFormat::HTML => self.generate_html_architecture(layers, &output_path),
            ImageFormat::PNG => {
                // For PNG, we'll generate SVG first then mention it needs conversion
                self.generate_svg_architecture(layers, &output_path.replace(".png", ".svg"))?;
                println!("SVG generated. Use external tool to convert to PNG if needed.");
                Ok(())
            }
        }
    }

    /// Generate SVG architecture diagram
    fn generate_svg_architecture(
        &self,
        layers: &[LayerInfo],
        output_path: &str,
    ) -> NeuralResult<()> {
        let mut svg = String::new();

        // SVG header
        svg.push_str(&format!(
            "<svg width=\"800\" height=\"{}\" xmlns=\"http://www.w3.org/2000/svg\">\n            <defs>\n                <style>\n                    .layer-box {{ fill: #e1f5fe; stroke: #0277bd; stroke-width: 2; }}\n                    .layer-text {{ font-family: Arial, sans-serif; font-size: 12px; text-anchor: middle; }}\n                    .layer-name {{ font-weight: bold; }}\n                    .layer-shape {{ font-size: 10px; fill: #666; }}\n                    .connection {{ stroke: #424242; stroke-width: 2; marker-end: url(#arrowhead); }}\n                </style>\n                <marker id=\"arrowhead\" markerWidth=\"10\" markerHeight=\"7\" \n                    refX=\"10\" refY=\"3.5\" orient=\"auto\">\n                    <polygon points=\"0,0 10,3.5 0,7\" fill=\"#424242\" />\n                </marker>\n            </defs>\n            ", 
            layers.len() * 100 + 100
        ));

        // Draw layers
        for (i, layer) in layers.iter().enumerate() {
            let y = i * 100 + 50;
            let x = 400;

            // Layer box
            svg.push_str(&format!(
                r#"<rect x="{}" y="{}" width="200" height="60" class="layer-box" />
                "#,
                x - 100,
                y - 30
            ));

            // Layer name
            if self.config.show_layer_names {
                svg.push_str(&format!(
                    r#"<text x="{}" y="{}" class="layer-text layer-name">{}</text>
                    "#,
                    x,
                    y - 10,
                    layer.name
                ));
            }

            // Layer type
            svg.push_str(&format!(
                r#"<text x="{}" y="{}" class="layer-text">{}</text>
                "#,
                x,
                y + 5,
                layer.layer_type
            ));

            // Shape information
            if self.config.show_tensor_shapes {
                svg.push_str(&format!(
                    r#"<text x="{}" y="{}" class="layer-text layer-shape">{:?}</text>
                    "#,
                    x,
                    y + 20,
                    layer.output_shape
                ));
            }

            // Connection to next layer
            if i < layers.len() - 1 {
                svg.push_str(&format!(
                    r#"<line x1="{}" y1="{}" x2="{}" y2="{}" class="connection" />
                    "#,
                    x,
                    y + 30,
                    x,
                    y + 70
                ));
            }
        }

        svg.push_str("</svg>");

        // Write to file
        let mut file = File::create(output_path)
            .map_err(|e| SklearsError::InvalidInput(format!("Failed to create file: {}", e)))?;
        file.write_all(svg.as_bytes())
            .map_err(|e| SklearsError::InvalidInput(format!("Failed to write file: {}", e)))?;

        Ok(())
    }

    /// Generate HTML architecture diagram with interactive features
    fn generate_html_architecture(
        &self,
        layers: &[LayerInfo],
        output_path: &str,
    ) -> NeuralResult<()> {
        let mut html = String::new();

        html.push_str(
            r#"
        <!DOCTYPE html>
        <html>
        <head>
            <title>Neural Network Architecture</title>
            <style>
                body { font-family: Arial, sans-serif; margin: 20px; }
                .architecture { display: flex; flex-direction: column; align-items: center; }
                .layer { 
                    background: #e1f5fe; 
                    border: 2px solid #0277bd; 
                    border-radius: 8px;
                    padding: 15px; 
                    margin: 10px;
                    min-width: 200px;
                    text-align: center;
                    transition: all 0.3s ease;
                }
                .layer:hover { 
                    background: #b3e5fc; 
                    transform: scale(1.05);
                    box-shadow: 0 4px 8px rgba(0,0,0,0.2);
                }
                .layer-name { font-weight: bold; font-size: 16px; color: #0277bd; }
                .layer-type { font-size: 14px; color: #424242; margin: 5px 0; }
                .layer-shape { font-size: 12px; color: #666; }
                .arrow { 
                    font-size: 24px; 
                    color: #424242; 
                    margin: 5px 0;
                }
                .layer-details {
                    display: none;
                    margin-top: 10px;
                    padding: 10px;
                    background: #f5f5f5;
                    border-radius: 4px;
                    font-size: 12px;
                }
            </style>
            <script>
                function toggleDetails(layerId) {
                    const details = document.getElementById(layerId);
                    details.style.display = details.style.display === 'none' ? 'block' : 'none';
                }
            </script>
        </head>
        <body>
            <h1>Neural Network Architecture</h1>
            <div class="architecture">
        "#,
        );

        for (i, layer) in layers.iter().enumerate() {
            html.push_str(&format!(
                r#"
                <div class="layer" onclick="toggleDetails('details_{}')">
                    <div class="layer-name">{}</div>
                    <div class="layer-type">{}</div>
                    <div class="layer-shape">Shape: {:?}</div>
                    <div id="details_{}" class="layer-details">
                        <strong>Parameters:</strong> {}<br>
                        <strong>Activation:</strong> {}<br>
                        <strong>Trainable:</strong> {}
                    </div>
                </div>
                "#,
                i,
                layer.name,
                layer.layer_type,
                layer.output_shape,
                i,
                layer.num_parameters,
                layer.activation.as_deref().unwrap_or("None"),
                layer.trainable
            ));

            if i < layers.len() - 1 {
                html.push_str(r#"<div class="arrow">↓</div>"#);
            }
        }

        html.push_str(
            r#"
            </div>
        </body>
        </html>
        "#,
        );

        // Write to file
        let mut file = File::create(output_path)
            .map_err(|e| SklearsError::InvalidInput(format!("Failed to create file: {}", e)))?;
        file.write_all(html.as_bytes())
            .map_err(|e| SklearsError::InvalidInput(format!("Failed to write file: {}", e)))?;

        Ok(())
    }

    /// Get file extension for current format
    fn format_extension(&self) -> &'static str {
        match self.config.format {
            ImageFormat::SVG => "svg",
            ImageFormat::PNG => "png",
            ImageFormat::HTML => "html",
        }
    }
}

/// Information about a layer for visualization
#[derive(Debug, Clone)]
pub struct LayerInfo {
    /// Human-readable name identifying this layer in the architecture diagram
    pub name: String,
    /// String descriptor of the layer class (e.g., `"Dense"`, `"Conv2D"`)
    pub layer_type: String,
    /// Shape of the output tensor produced by this layer
    pub output_shape: Vec<usize>,
    /// Total number of trainable parameters in this layer
    pub num_parameters: usize,
    /// Name of the activation function applied after this layer, if any
    pub activation: Option<String>,
    /// Whether the layer's parameters are updated during training
    pub trainable: bool,
}

/// Training metrics visualizer
pub struct TrainingVisualizer {
    config: VisualizationConfig,
}

impl TrainingVisualizer {
    /// Create a new training visualizer
    pub fn new(config: VisualizationConfig) -> Self {
        Self { config }
    }

    /// Plot training history (loss, accuracy, etc.)
    pub fn plot_training_history(
        &self,
        metrics: &TrainingMetrics,
        filename: &str,
    ) -> NeuralResult<()> {
        let output_path = format!("{}/{}.html", self.config.output_dir, filename);

        let mut html = String::new();
        html.push_str(
            r#"
        <!DOCTYPE html>
        <html>
        <head>
            <title>Training History</title>
            <script src="https://cdn.plot.ly/plotly-latest.min.js"></script>
            <style>
                body { font-family: Arial, sans-serif; margin: 20px; }
                .plot-container { width: 100%; height: 400px; margin: 20px 0; }
            </style>
        </head>
        <body>
            <h1>Training History</h1>
        "#,
        );

        // Loss plot
        html.push_str(r#"<div id="loss-plot" class="plot-container"></div>"#);
        html.push_str(&format!(
            r#"
        <script>
            var lossData = [{{
                x: [{}],
                y: [{}],
                type: 'scatter',
                mode: 'lines',
                name: 'Training Loss',
                line: {{color: '#1f77b4'}}
            }}
        "#,
            (0..metrics.train_loss.len())
                .map(|i| i.to_string())
                .collect::<Vec<_>>()
                .join(","),
            metrics
                .train_loss
                .iter()
                .map(|x| x.to_string())
                .collect::<Vec<_>>()
                .join(",")
        ));

        if !metrics.val_loss.is_empty() {
            html.push_str(&format!(
                r#",{{
                x: [{}],
                y: [{}],
                type: 'scatter',
                mode: 'lines',
                name: 'Validation Loss',
                line: {{color: '#ff7f0e'}}
            }}"#,
                (0..metrics.val_loss.len())
                    .map(|i| i.to_string())
                    .collect::<Vec<_>>()
                    .join(","),
                metrics
                    .val_loss
                    .iter()
                    .map(|x| x.to_string())
                    .collect::<Vec<_>>()
                    .join(",")
            ));
        }

        html.push_str(
            r#"];
            var lossLayout = {
                title: 'Training Loss',
                xaxis: { title: 'Epoch' },
                yaxis: { title: 'Loss' }
            };
            Plotly.newPlot('loss-plot', lossData, lossLayout);
        </script>
        "#,
        );

        // Accuracy plot (if available)
        if !metrics.train_accuracy.is_empty() {
            html.push_str(r#"<div id="accuracy-plot" class="plot-container"></div>"#);
            html.push_str(&format!(
                r#"
            <script>
                var accuracyData = [{{
                    x: [{}],
                    y: [{}],
                    type: 'scatter',
                    mode: 'lines',
                    name: 'Training Accuracy',
                    line: {{color: '#2ca02c'}}
                }}
            "#,
                (0..metrics.train_accuracy.len())
                    .map(|i| i.to_string())
                    .collect::<Vec<_>>()
                    .join(","),
                metrics
                    .train_accuracy
                    .iter()
                    .map(|x| x.to_string())
                    .collect::<Vec<_>>()
                    .join(",")
            ));

            if !metrics.val_accuracy.is_empty() {
                html.push_str(&format!(
                    r#",{{
                    x: [{}],
                    y: [{}],
                    type: 'scatter',
                    mode: 'lines',
                    name: 'Validation Accuracy',
                    line: {{color: '#d62728'}}
                }}"#,
                    (0..metrics.val_accuracy.len())
                        .map(|i| i.to_string())
                        .collect::<Vec<_>>()
                        .join(","),
                    metrics
                        .val_accuracy
                        .iter()
                        .map(|x| x.to_string())
                        .collect::<Vec<_>>()
                        .join(",")
                ));
            }

            html.push_str(
                r#"];
                var accuracyLayout = {
                    title: 'Training Accuracy',
                    xaxis: { title: 'Epoch' },
                    yaxis: { title: 'Accuracy' }
                };
                Plotly.newPlot('accuracy-plot', accuracyData, accuracyLayout);
            </script>
            "#,
            );
        }

        html.push_str(
            r#"
        </body>
        </html>
        "#,
        );

        // Write to file
        let mut file = File::create(&output_path)
            .map_err(|e| SklearsError::InvalidInput(format!("Failed to create file: {}", e)))?;
        file.write_all(html.as_bytes())
            .map_err(|e| SklearsError::InvalidInput(format!("Failed to write file: {}", e)))?;

        println!("Training history saved to: {}", output_path);
        Ok(())
    }
}

/// Training metrics for visualization
#[derive(Debug, Clone, Default)]
pub struct TrainingMetrics {
    /// Per-epoch training loss values
    pub train_loss: Vec<f64>,
    /// Per-epoch validation loss values
    pub val_loss: Vec<f64>,
    /// Per-epoch training accuracy values
    pub train_accuracy: Vec<f64>,
    /// Per-epoch validation accuracy values
    pub val_accuracy: Vec<f64>,
    /// Learning rate schedule over epochs
    pub learning_rates: Vec<f64>,
}

/// Attention heatmap visualizer
pub struct AttentionVisualizer {
    config: VisualizationConfig,
}

impl AttentionVisualizer {
    /// Create a new attention visualizer
    pub fn new(config: VisualizationConfig) -> Self {
        Self { config }
    }

    /// Generate attention heatmap visualization
    pub fn visualize_attention_weights<T: FloatBounds>(
        &self,
        attention_weights: &Array3<T>,
        tokens: &[String],
        filename: &str,
    ) -> NeuralResult<()> {
        let output_path = format!("{}/{}.html", self.config.output_dir, filename);

        let mut html = String::new();
        html.push_str(
            r#"
        <!DOCTYPE html>
        <html>
        <head>
            <title>Attention Heatmap</title>
            <script src="https://cdn.plot.ly/plotly-latest.min.js"></script>
            <style>
                body { font-family: Arial, sans-serif; margin: 20px; }
                .heatmap-container { width: 100%; height: 600px; margin: 20px 0; }
            </style>
        </head>
        <body>
            <h1>Attention Weights Heatmap</h1>
            <div id="heatmap" class="heatmap-container"></div>
            <script>
        "#,
        );

        // Get the first head of the first layer for visualization
        let (batch_size, seq_len, _) = attention_weights.dim();
        if batch_size > 0 && seq_len > 0 {
            // Convert attention weights to JavaScript format
            let mut weights_js = String::new();
            weights_js.push('[');
            for i in 0..seq_len {
                weights_js.push('[');
                for j in 0..seq_len {
                    if j > 0 {
                        weights_js.push(',');
                    }
                    weights_js.push_str(
                        &attention_weights[[0, i, j]]
                            .to_f64()
                            .unwrap_or(0.0)
                            .to_string(),
                    );
                }
                weights_js.push(']');
                if i < seq_len - 1 {
                    weights_js.push(',');
                }
            }
            weights_js.push(']');

            // Convert tokens to JavaScript format
            let tokens_js = format!("[\"{}\"]", tokens.join("\",\""));

            html.push_str(&format!(
                r#"
                var data = [{{
                    z: {},
                    x: {},
                    y: {},
                    type: 'heatmap',
                    colorscale: 'Viridis'
                }}];
                
                var layout = {{
                    title: 'Attention Weights',
                    xaxis: {{ title: 'Key Tokens' }},
                    yaxis: {{ title: 'Query Tokens' }}
                }};
                
                Plotly.newPlot('heatmap', data, layout);
            "#,
                weights_js, tokens_js, tokens_js
            ));
        }

        html.push_str(
            r#"
            </script>
        </body>
        </html>
        "#,
        );

        // Write to file
        let mut file = File::create(&output_path)
            .map_err(|e| SklearsError::InvalidInput(format!("Failed to create file: {}", e)))?;
        file.write_all(html.as_bytes())
            .map_err(|e| SklearsError::InvalidInput(format!("Failed to write file: {}", e)))?;

        println!("Attention heatmap saved to: {}", output_path);
        Ok(())
    }
}

/// Weight distribution visualizer
pub struct WeightVisualizer {
    config: VisualizationConfig,
}

impl WeightVisualizer {
    /// Create a new weight visualizer
    pub fn new(config: VisualizationConfig) -> Self {
        Self { config }
    }

    /// Visualize weight distributions across layers
    pub fn visualize_weight_distributions<T: FloatBounds>(
        &self,
        weights: &HashMap<String, Array2<T>>,
        filename: &str,
    ) -> NeuralResult<()> {
        let output_path = format!("{}/{}.html", self.config.output_dir, filename);

        let mut html = String::new();
        html.push_str(
            r#"
        <!DOCTYPE html>
        <html>
        <head>
            <title>Weight Distributions</title>
            <script src="https://cdn.plot.ly/plotly-latest.min.js"></script>
            <style>
                body { font-family: Arial, sans-serif; margin: 20px; }
                .histogram-container { width: 100%; height: 400px; margin: 20px 0; }
            </style>
        </head>
        <body>
            <h1>Weight Distributions</h1>
        "#,
        );

        for (layer_name, weight_matrix) in weights.iter() {
            let div_id = format!("histogram-{}", layer_name.replace(".", "-"));
            html.push_str(&format!(r#"<h2>{}</h2>"#, layer_name));
            html.push_str(&format!(
                r#"<div id="{}" class="histogram-container"></div>"#,
                div_id
            ));

            // Flatten weights and convert to JavaScript array
            let flattened: Vec<f64> = weight_matrix
                .iter()
                .map(|&w| w.to_f64().unwrap_or(0.0))
                .collect();

            let weights_js = format!(
                "[{}]",
                flattened
                    .iter()
                    .map(|x| x.to_string())
                    .collect::<Vec<_>>()
                    .join(",")
            );

            html.push_str(&format!(
                r#"
            <script>
                var data_{} = [{{
                    x: {},
                    type: 'histogram',
                    nbinsx: 50,
                    name: '{}'
                }}];
                
                var layout_{} = {{
                    title: '{} Weight Distribution',
                    xaxis: {{ title: 'Weight Value' }},
                    yaxis: {{ title: 'Frequency' }}
                }};
                
                Plotly.newPlot('{}', data_{}, layout_{});
            </script>
            "#,
                div_id, weights_js, layer_name, div_id, layer_name, div_id, div_id, div_id
            ));
        }

        html.push_str(
            r#"
        </body>
        </html>
        "#,
        );

        // Write to file
        let mut file = File::create(&output_path)
            .map_err(|e| SklearsError::InvalidInput(format!("Failed to create file: {}", e)))?;
        file.write_all(html.as_bytes())
            .map_err(|e| SklearsError::InvalidInput(format!("Failed to write file: {}", e)))?;

        println!("Weight distributions saved to: {}", output_path);
        Ok(())
    }
}

/// Create output directory if it doesn't exist
pub fn ensure_output_directory(path: &str) -> NeuralResult<()> {
    std::fs::create_dir_all(path)
        .map_err(|e| SklearsError::InvalidInput(format!("Failed to create directory: {}", e)))
}

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

    #[test]
    fn test_visualization_config_default() {
        let config = VisualizationConfig::default();
        assert_eq!(config.format, ImageFormat::SVG);
        assert_eq!(config.color_scheme, ColorScheme::Default);
        assert_eq!(config.dpi, 300);
    }

    #[test]
    fn test_layer_info_creation() {
        let layer_info = LayerInfo {
            name: "dense_1".to_string(),
            layer_type: "Dense".to_string(),
            output_shape: vec![128, 64],
            num_parameters: 8256,
            activation: Some("ReLU".to_string()),
            trainable: true,
        };

        assert_eq!(layer_info.name, "dense_1");
        assert_eq!(layer_info.num_parameters, 8256);
    }

    #[test]
    fn test_training_metrics_default() {
        let metrics = TrainingMetrics::default();
        assert!(metrics.train_loss.is_empty());
        assert!(metrics.val_loss.is_empty());
        assert!(metrics.train_accuracy.is_empty());
    }
}