ruviz 0.4.3

//! Heatmap example demonstrating various heatmap features
//!
//! Run with: cargo run --example heatmap_example

use ruviz::prelude::*;
use std::f64::consts::PI;

fn main() -> Result<()> {
    // Example 1: Basic heatmap with default colormap
    basic_heatmap()?;

    // Example 2: Correlation matrix style heatmap with annotations
    correlation_matrix()?;

    // Example 3: Scientific heatmap with diverging colormap
    scientific_heatmap()?;

    // Example 4: Log-scaled heatmap with masked zeros
    masked_log_heatmap()?;

    // Example 5: Large heatmap demonstration
    large_heatmap()?;

    println!("All heatmap examples saved successfully!");
    Ok(())
}

/// Basic heatmap with default viridis colormap
fn basic_heatmap() -> Result<()> {
    // Create a 5x5 grid of values
    let data = vec![
        vec![1.0, 2.0, 3.0, 4.0, 5.0],
        vec![2.0, 3.0, 4.0, 5.0, 6.0],
        vec![3.0, 4.0, 5.0, 6.0, 7.0],
        vec![4.0, 5.0, 6.0, 7.0, 8.0],
        vec![5.0, 6.0, 7.0, 8.0, 9.0],
    ];

    Plot::new()
        .heatmap(&data, Some(HeatmapConfig::default()))
        .title("Basic Heatmap")
        .save("generated/examples/heatmap_basic.png")?;

    println!("Saved: generated/examples/heatmap_basic.png");
    Ok(())
}

/// Correlation matrix with annotations and custom colormap
fn correlation_matrix() -> Result<()> {
    // Simulated correlation matrix
    let data = vec![
        vec![1.0, 0.8, 0.6, -0.2, -0.4],
        vec![0.8, 1.0, 0.5, 0.0, -0.3],
        vec![0.6, 0.5, 1.0, 0.3, 0.1],
        vec![-0.2, 0.0, 0.3, 1.0, 0.7],
        vec![-0.4, -0.3, 0.1, 0.7, 1.0],
    ];

    let config = HeatmapConfig::new()
        .colormap(ColorMap::coolwarm()) // Diverging colormap
        .vmin(-1.0)
        .vmax(1.0)
        .annotate(true)
        .colorbar(true)
        .colorbar_label("Correlation");

    Plot::new()
        .heatmap(&data, Some(config))
        .title("Correlation Matrix")
        .save("generated/examples/heatmap_correlation.png")?;

    println!("Saved: generated/examples/heatmap_correlation.png");
    Ok(())
}

/// Scientific heatmap with a function surface
fn scientific_heatmap() -> Result<()> {
    // Create a 2D sine wave pattern
    let rows = 20;
    let cols = 30;
    let mut data = vec![vec![0.0; cols]; rows];

    for (i, row) in data.iter_mut().enumerate().take(rows) {
        for (j, value) in row.iter_mut().enumerate().take(cols) {
            let x = j as f64 * 2.0 * PI / cols as f64;
            let y = i as f64 * 2.0 * PI / rows as f64;
            *value = (x.sin() * y.cos()).sin();
        }
    }

    let config = HeatmapConfig::new()
        .colormap(ColorMap::plasma())
        .colorbar(true)
        .colorbar_label("sin(sin(x)*cos(y))")
        .aspect(1.0);

    Plot::new()
        .heatmap(&data, Some(config))
        .title("2D Sine Wave Surface")
        .xlabel("X")
        .ylabel("Y")
        .save("generated/examples/heatmap_scientific.png")?;

    println!("Saved: generated/examples/heatmap_scientific.png");
    Ok(())
}

/// Log-scaled heatmap demonstrating masked zero-value cutouts
fn masked_log_heatmap() -> Result<()> {
    let rows = 72usize;
    let cols = 96usize;
    let center_col = (cols.saturating_sub(1)) as f64 / 2.0;
    let mut data = vec![vec![0.0; cols]; rows];

    for (row, row_values) in data.iter_mut().enumerate() {
        let depth = 1.0 - row as f64 / (rows.saturating_sub(1).max(1)) as f64;

        for (col, value) in row_values.iter_mut().enumerate() {
            let x = (col as f64 - center_col) / cols as f64;
            let beam_core = 250.0 * (-(x / 0.018).powi(2)).exp() * (-(depth / 0.12)).exp();
            let halo = 9.0 * (-(x / 0.18).powi(2)).exp() * (-(depth / 0.55)).exp();
            let buildup =
                2.5 * (-(x / 0.25).powi(2)).exp() * (-((depth - 0.32) / 0.035).powi(2)).exp();
            let background = 2.0e-4 * (-(x / 0.34).powi(2)).exp() * (-(depth / 0.9)).exp();
            let masked_edge = col < 8 || col >= cols.saturating_sub(8);
            let sparse_mask =
                (col < 14 || col >= cols.saturating_sub(14)) && ((row + col * 5) % 13 < 5);

            *value = if masked_edge || sparse_mask {
                0.0
            } else {
                beam_core + halo + buildup + background
            };
        }
    }

    let config = HeatmapConfig::new()
        .value_scale(AxisScale::Log)
        .colorbar(true)
        .colorbar_log_subticks(true)
        .colorbar_label("Absorbed Energy");

    Plot::new()
        .heatmap(&data, Some(config))
        .title("Masked Log Heatmap")
        .xlabel("Position in x (cells)")
        .ylabel("Depth (cells)")
        .ylim(rows as f64, 0.0)
        .save("generated/examples/heatmap_log_masked.png")?;

    println!("Saved: generated/examples/heatmap_log_masked.png");
    Ok(())
}

/// Large heatmap for performance demonstration
fn large_heatmap() -> Result<()> {
    // Create a 50x50 Gaussian-like pattern
    let size = 50;
    let mut data = vec![vec![0.0; size]; size];

    for (i, row) in data.iter_mut().enumerate().take(size) {
        for (j, value) in row.iter_mut().enumerate().take(size) {
            let x = (j as f64 - size as f64 / 2.0) / 10.0;
            let y = (i as f64 - size as f64 / 2.0) / 10.0;
            *value = (-x * x - y * y).exp();
        }
    }

    let config = HeatmapConfig::new()
        .colormap(ColorMap::inferno())
        .colorbar(true)
        .colorbar_label("Intensity")
        .vmin(0.0)
        .vmax(1.0);

    Plot::new()
        .heatmap(&data, Some(config))
        .title("2D Gaussian Distribution")
        .xlabel("X")
        .ylabel("Y")
        .save("generated/examples/heatmap_large.png")?;

    println!("Saved: generated/examples/heatmap_large.png");
    Ok(())
}