crater/serde/threejs/

field_visualization.rs

//! Scalar field visualization for Three.js
//!
//! This module provides high-quality, performant visualization of scalar fields
//! with support for continuous gradients, volume rendering, and multiple color schemes.

use crate::csg::prelude::*;
use burn::prelude::*;
use serde_json::{Value, json};
use std::error::Error;

/// Configuration for scalar field visualization
#[derive(Debug, Clone)]
pub struct FieldVisualizationConfig {
    /// Sampling resolution in each dimension [x, y, z]
    pub resolution: [usize; 3],
    /// Spatial bounds [[min_x, min_y, min_z], [max_x, max_y, max_z]]
    pub bounds: [[f32; 3]; 2],
    /// Color mapping scheme
    pub color_scheme: ColorScheme,
    /// Overall opacity for the visualization
    pub opacity: f32,
    /// Rendering method
    pub mode: RenderingMode,
}

/// Color schemes for field visualization
#[derive(Debug, Clone)]
pub enum ColorScheme {
    /// Perceptually uniform viridis color map (recommended)
    Viridis,
    /// Classic blue-to-red heat map
    Thermal,
    /// Grayscale mapping
    Grayscale,
    /// Custom three-point gradient
    Custom {
        negative: String,
        zero: String,
        positive: String,
    },
}

/// Field rendering methods
#[derive(Debug, Clone)]
pub enum RenderingMode {
    /// Continuous volume gradient (for smooth fields)
    VolumeGradient,
    /// Discrete point cloud
    PointCloud { point_size: f32 },
    /// Gradient vector arrows (direction and magnitude visualization)
    GradientArrows {
        arrow_scale: f32,
        density_factor: f32,
    },
}

impl Default for FieldVisualizationConfig {
    fn default() -> Self {
        Self {
            resolution: [32, 32, 32],
            bounds: [[-2.0, -2.0, -2.0], [2.0, 2.0, 2.0]],
            color_scheme: ColorScheme::Viridis,
            opacity: 0.6,
            mode: RenderingMode::VolumeGradient,
        }
    }
}

/// Complete field visualization data package
#[derive(Debug)]
pub struct FieldVisualization {
    /// Three.js-compatible rendering data
    pub render_data: Value,
    /// Field statistics and metadata
    pub metadata: FieldMetadata,
}

/// Statistical information about the field
#[derive(Debug)]
pub struct FieldMetadata {
    pub dimension: usize,
    pub sample_count: usize,
    pub value_range: (f32, f32),
    pub resolution: [usize; 3],
    pub bounds: [[f32; 3]; 2],
}

/// Main entry point for field visualization
pub fn visualize_field<B: Backend, const N: usize>(
    field: &dyn ScalarField<N, B>,
    config: &FieldVisualizationConfig,
) -> Result<FieldVisualization, Box<dyn Error>> {
    // Validate configuration
    validate_config::<N>(config)?;

    // Sample the field on a regular grid
    let samples = sample_field::<B, N>(field, config)?;

    // Generate visualization based on rendering mode
    let render_data = match &config.mode {
        RenderingMode::VolumeGradient => VolumeRenderer::new(config).render::<N>(&samples)?,
        RenderingMode::PointCloud { point_size } => {
            PointCloudRenderer::new(config, *point_size).render::<N>(&samples)?
        }
        RenderingMode::GradientArrows {
            arrow_scale,
            density_factor,
        } => GradientArrowRenderer::new(config, *arrow_scale, *density_factor)
            .render::<B, N>(field, &samples)?,
    };

    let metadata = FieldMetadata {
        dimension: N,
        sample_count: samples.values.len(),
        value_range: (samples.min_value, samples.max_value),
        resolution: config.resolution,
        bounds: config.bounds,
    };

    Ok(FieldVisualization {
        render_data,
        metadata,
    })
}

/// Field sampling results
#[derive(Debug)]
struct FieldSamples {
    positions: Vec<f32>,
    values: Vec<f32>,
    min_value: f32,
    max_value: f32,
}

/// Sample the scalar field on a regular grid
fn sample_field<B: Backend, const N: usize>(
    field: &dyn ScalarField<N, B>,
    config: &FieldVisualizationConfig,
) -> Result<FieldSamples, Box<dyn Error>> {
    let grid_points = generate_grid::<B, N>(config)?;
    let field_values = field.evaluate(grid_points.clone());

    // Convert to host data
    let positions_data = grid_points.to_data();
    let values_data = field_values.to_data();

    let positions: Vec<f32> = positions_data.iter::<f32>().collect();
    let values: Vec<f32> = values_data.iter::<f32>().collect();

    // Compute value statistics
    let min_value = values.iter().fold(f32::INFINITY, |a, &b| a.min(b));
    let max_value = values.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b));

    Ok(FieldSamples {
        positions,
        values,
        min_value,
        max_value,
    })
}

/// Generate regular sampling grid
fn generate_grid<B: Backend, const N: usize>(
    config: &FieldVisualizationConfig,
) -> Result<Tensor<B, 2>, Box<dyn Error>> {
    let device = B::Device::default();
    let total_points = config.resolution.iter().take(N).product::<usize>();
    let mut points = Vec::with_capacity(total_points * N);

    match N {
        2 => generate_2d_grid(&mut points, config),
        3 => generate_3d_grid(&mut points, config),
        _ => return Err("Only 2D and 3D fields are supported".into()),
    }

    let tensor = Tensor::<B, 1>::from_data(points.as_slice(), &device);
    Ok(tensor.reshape([total_points, N]))
}

fn generate_2d_grid(points: &mut Vec<f32>, config: &FieldVisualizationConfig) {
    let [nx, ny, _] = config.resolution;
    let [[x_min, y_min, _], [x_max, y_max, _]] = config.bounds;

    for j in 0..ny {
        for i in 0..nx {
            let x = x_min + (i as f32 / (nx - 1) as f32) * (x_max - x_min);
            let y = y_min + (j as f32 / (ny - 1) as f32) * (y_max - y_min);
            points.extend_from_slice(&[x, y]);
        }
    }
}

fn generate_3d_grid(points: &mut Vec<f32>, config: &FieldVisualizationConfig) {
    let [nx, ny, nz] = config.resolution;
    let [[x_min, y_min, z_min], [x_max, y_max, z_max]] = config.bounds;

    for k in 0..nz {
        for j in 0..ny {
            for i in 0..nx {
                let x = x_min + (i as f32 / (nx - 1) as f32) * (x_max - x_min);
                let y = y_min + (j as f32 / (ny - 1) as f32) * (y_max - y_min);
                let z = z_min + (k as f32 / (nz - 1) as f32) * (z_max - z_min);
                points.extend_from_slice(&[x, y, z]);
            }
        }
    }
}

/// Volume gradient renderer for continuous field visualization
struct VolumeRenderer<'a> {
    config: &'a FieldVisualizationConfig,
}

impl<'a> VolumeRenderer<'a> {
    fn new(config: &'a FieldVisualizationConfig) -> Self {
        Self { config }
    }

    fn render<const N: usize>(&self, samples: &FieldSamples) -> Result<Value, Box<dyn Error>> {
        let color_mapper = ColorMapper::new(&self.config.color_scheme);
        let mut texture_data = Vec::new();

        for &value in samples.values.iter() {
            if !value.is_finite() {
                continue;
            }

            let (r, g, b) = color_mapper.map_value(value, samples.min_value, samples.max_value);
            let alpha = self.compute_alpha(value, samples.min_value, samples.max_value);

            texture_data.push(json!({
                "r": r,
                "g": g,
                "b": b,
                "a": alpha,
                "value": value
            }));
        }

        Ok(json!({
            "type": "volume",
            "mode": "gradient",
            "resolution": self.config.resolution,
            "bounds": self.config.bounds,
            "textureData": texture_data,
            "opacity": self.config.opacity,
            "colorScheme": format!("{:?}", self.config.color_scheme)
        }))
    }

    fn compute_alpha(&self, value: f32, min_val: f32, max_val: f32) -> u8 {
        let range = max_val - min_val;
        if range <= 0.0 {
            return (self.config.opacity * 255.0) as u8;
        }

        // Simple, smooth gradient based on distance from zero isosurface
        let abs_value = value.abs();
        let max_abs = (min_val.abs()).max(max_val.abs());

        if max_abs <= 0.0 {
            return (self.config.opacity * 255.0) as u8;
        }

        // Fade out as we get further from the isosurface (value = 0)
        // This creates a subtle visualization focused on the boundary
        let distance_factor = 1.0 - (abs_value / max_abs).min(1.0);
        let smooth_factor = distance_factor * distance_factor * distance_factor * distance_factor; // Quartic falloff for smoother appearance

        let alpha = self.config.opacity * smooth_factor;
        (alpha * 255.0) as u8
    }
}

/// Point cloud renderer for discrete visualization
struct PointCloudRenderer<'a> {
    config: &'a FieldVisualizationConfig,
    point_size: f32,
}

impl<'a> PointCloudRenderer<'a> {
    fn new(config: &'a FieldVisualizationConfig, point_size: f32) -> Self {
        Self { config, point_size }
    }

    fn render<const N: usize>(&self, samples: &FieldSamples) -> Result<Value, Box<dyn Error>> {
        let color_mapper = ColorMapper::new(&self.config.color_scheme);
        let mut points = Vec::new();

        for (i, &value) in samples.values.iter().enumerate() {
            if !value.is_finite() {
                continue;
            }

            let position = self.extract_position::<N>(&samples.positions, i);
            let color = color_mapper.map_to_hex(value, samples.min_value, samples.max_value)?;

            points.push(json!({
                "position": position,
                "value": value,
                "color": color,
                "size": self.point_size
            }));
        }

        Ok(json!({
            "type": "points",
            "data": points
        }))
    }

    fn extract_position<const N: usize>(&self, positions: &[f32], index: usize) -> [f32; 3] {
        match N {
            2 => [positions[index * 2], positions[index * 2 + 1], 0.0],
            3 => [
                positions[index * 3],
                positions[index * 3 + 1],
                positions[index * 3 + 2],
            ],
            _ => [0.0, 0.0, 0.0],
        }
    }
}

/// Gradient arrow renderer for vector field visualization
struct GradientArrowRenderer<'a> {
    config: &'a FieldVisualizationConfig,
    arrow_scale: f32,
    density_factor: f32,
}

impl<'a> GradientArrowRenderer<'a> {
    fn new(config: &'a FieldVisualizationConfig, arrow_scale: f32, density_factor: f32) -> Self {
        Self {
            config,
            arrow_scale,
            density_factor,
        }
    }

    fn render<B: Backend, const N: usize>(
        &self,
        field: &dyn ScalarField<N, B>,
        samples: &FieldSamples,
    ) -> Result<Value, Box<dyn Error>> {
        let mut arrows = Vec::new();
        let _color_mapper = ColorMapper::new(&self.config.color_scheme);

        // Compute gradients numerically
        let gradients = self.compute_gradients::<B, N>(field, samples)?;

        for (i, (&value, gradient)) in samples.values.iter().zip(gradients.iter()).enumerate() {
            if !value.is_finite() {
                continue;
            }

            // Skip points with very small gradient magnitude (no interesting direction)
            let grad_magnitude = gradient.iter().map(|x| x * x).sum::<f32>().sqrt();
            if grad_magnitude < 1e-6 {
                continue;
            }

            // Density filtering: favor points near isosurface (low absolute field values)
            let abs_value = value.abs();
            // Invert the density threshold - higher density near zero field values
            let density_threshold = 1.0 - (abs_value * self.density_factor).min(0.8);

            // Use deterministic sampling based on position for consistent results
            let position = self.extract_position::<N>(&samples.positions, i);
            let hash = ((position[0] * 73.0
                + position[1] * 151.0
                + position.get(2).unwrap_or(&0.0) * 233.0)
                .abs()
                * 1000.0) as u32;
            let random_val = (hash % 1000) as f32 / 1000.0;

            if random_val > density_threshold {
                continue;
            }

            // Fixed arrow length - all arrows same size, only direction varies
            let arrow_length = self.arrow_scale;

            // Normalize gradient direction and flip if field is positive (point toward zero)
            let direction_sign = if value > 0.0 { -1.0 } else { 1.0 };
            let direction: Vec<f32> = gradient
                .iter()
                .map(|x| direction_sign * x / grad_magnitude)
                .collect();

            // Arrow endpoint
            let mut end_position = position;
            for (i, &dir) in direction.iter().enumerate().take(N) {
                if i < 3 {
                    end_position[i] += dir * arrow_length;
                }
            }

            // Color based on field value proximity to zero (brighter near isosurface)
            // Use inverse of absolute field value so arrows are brightest near zero
            // More aggressive scaling to make more arrows bright
            let field_proximity = 1.0 / (1.0 + abs_value * 1.5); // Higher values = closer to zero
            let field_proximity = field_proximity.powf(0.5); // Square root for more aggressive scaling

            // Simple gray to white interpolation based on proximity
            let gray_value = (128.0 + field_proximity * 127.0) as u8; // 128-255 range
            let color = format!("#{:02x}{:02x}{:02x}", gray_value, gray_value, gray_value);

            arrows.push(json!({
                "start": position,
                "end": end_position,
                "direction": direction,
                "magnitude": grad_magnitude,
                "fieldValue": value,
                "color": color,
                "length": arrow_length,
                "opacity": field_proximity  // Pass proximity for opacity control
            }));
        }

        Ok(json!({
            "type": "arrows",
            "data": arrows,
            "arrowScale": self.arrow_scale,
            "densityFactor": self.density_factor
        }))
    }

    fn compute_gradients<B: Backend, const N: usize>(
        &self,
        field: &dyn ScalarField<N, B>,
        samples: &FieldSamples,
    ) -> Result<Vec<Vec<f32>>, Box<dyn Error>> {
        let device = B::Device::default();
        let epsilon = 1e-4;

        let mut gradients = Vec::new();

        // Compute numerical gradients using finite differences
        for (i, _) in samples.values.iter().enumerate() {
            let base_pos = self.extract_position::<N>(&samples.positions, i);

            let mut gradient = Vec::new();

            for dim in 0..N {
                // Forward difference
                let mut pos_forward = base_pos;
                let mut pos_backward = base_pos;

                if dim < 3 {
                    pos_forward[dim] += epsilon;
                    pos_backward[dim] -= epsilon;
                }

                // Evaluate field at forward and backward positions
                let forward_coords: Vec<f32> = [
                    pos_forward[0],
                    pos_forward[1],
                    pos_forward.get(2).copied().unwrap_or(0.0),
                ]
                .iter()
                .take(N)
                .copied()
                .collect();
                let backward_coords: Vec<f32> = [
                    pos_backward[0],
                    pos_backward[1],
                    pos_backward.get(2).copied().unwrap_or(0.0),
                ]
                .iter()
                .take(N)
                .copied()
                .collect();

                let tensor_forward =
                    Tensor::<B, 1>::from_data(forward_coords.as_slice(), &device).reshape([1, N]);
                let tensor_backward =
                    Tensor::<B, 1>::from_data(backward_coords.as_slice(), &device).reshape([1, N]);

                let val_forward = field
                    .evaluate(tensor_forward)
                    .to_data()
                    .iter::<f32>()
                    .next()
                    .unwrap_or(0.0);
                let val_backward = field
                    .evaluate(tensor_backward)
                    .to_data()
                    .iter::<f32>()
                    .next()
                    .unwrap_or(0.0);

                let partial_derivative = (val_forward - val_backward) / (2.0 * epsilon);
                gradient.push(partial_derivative);
            }

            gradients.push(gradient);
        }

        Ok(gradients)
    }

    fn extract_position<const N: usize>(&self, positions: &[f32], index: usize) -> [f32; 3] {
        match N {
            2 => [positions[index * 2], positions[index * 2 + 1], 0.0],
            3 => [
                positions[index * 3],
                positions[index * 3 + 1],
                positions[index * 3 + 2],
            ],
            _ => [0.0, 0.0, 0.0],
        }
    }
}

/// Professional color mapping with multiple schemes
struct ColorMapper {
    scheme: ColorScheme,
}

impl ColorMapper {
    fn new(scheme: &ColorScheme) -> Self {
        Self {
            scheme: scheme.clone(),
        }
    }

    fn map_value(&self, value: f32, min_val: f32, max_val: f32) -> (u8, u8, u8) {
        let normalized = if max_val > min_val {
            ((value - min_val) / (max_val - min_val)).clamp(0.0, 1.0)
        } else {
            0.5
        };

        match &self.scheme {
            ColorScheme::Viridis => viridis_colormap(normalized),
            ColorScheme::Thermal => thermal_colormap(normalized),
            ColorScheme::Grayscale => {
                let gray = (normalized * 255.0) as u8;
                (gray, gray, gray)
            }
            ColorScheme::Custom {
                negative,
                zero,
                positive,
            } => self.custom_gradient(value, min_val, max_val, negative, zero, positive),
        }
    }

    fn map_to_hex(&self, value: f32, min_val: f32, max_val: f32) -> Result<String, Box<dyn Error>> {
        let (r, g, b) = self.map_value(value, min_val, max_val);
        Ok(format!("#{:02x}{:02x}{:02x}", r, g, b))
    }

    fn custom_gradient(
        &self,
        value: f32,
        min_val: f32,
        max_val: f32,
        neg_color: &str,
        zero_color: &str,
        pos_color: &str,
    ) -> (u8, u8, u8) {
        // Simplified three-color interpolation
        if value < 0.0 && min_val < 0.0 {
            let t = (value / min_val).clamp(0.0, 1.0);
            interpolate_hex_colors(neg_color, zero_color, t).unwrap_or((128, 128, 128))
        } else if value > 0.0 && max_val > 0.0 {
            let t = (value / max_val).clamp(0.0, 1.0);
            interpolate_hex_colors(zero_color, pos_color, t).unwrap_or((128, 128, 128))
        } else {
            parse_hex_color(zero_color).unwrap_or((128, 128, 128))
        }
    }
}

/// High-quality viridis color mapping
fn viridis_colormap(t: f32) -> (u8, u8, u8) {
    let t = t.clamp(0.0, 1.0);

    // High-fidelity viridis control points
    let control_points = [
        (68, 1, 84),    // Deep purple
        (59, 82, 139),  // Blue
        (33, 145, 140), // Teal
        (94, 201, 98),  // Green
        (253, 231, 37), // Yellow
    ];

    interpolate_control_points(&control_points, t)
}

/// Thermal (blue-red) color mapping
fn thermal_colormap(t: f32) -> (u8, u8, u8) {
    let t = t.clamp(0.0, 1.0);
    let r = (t * 255.0) as u8;
    let b = ((1.0 - t) * 255.0) as u8;
    (r, 0, b)
}

/// Interpolate between color control points
fn interpolate_control_points(points: &[(u8, u8, u8)], t: f32) -> (u8, u8, u8) {
    if points.is_empty() {
        return (0, 0, 0);
    }

    let scaled = t * (points.len() - 1) as f32;
    let index = scaled.floor() as usize;
    let frac = scaled - index as f32;

    if index >= points.len() - 1 {
        return points[points.len() - 1];
    }

    let (r1, g1, b1) = points[index];
    let (r2, g2, b2) = points[index + 1];

    let r = (r1 as f32 + (r2 as f32 - r1 as f32) * frac) as u8;
    let g = (g1 as f32 + (g2 as f32 - g1 as f32) * frac) as u8;
    let b = (b1 as f32 + (b2 as f32 - b1 as f32) * frac) as u8;

    (r, g, b)
}

/// Parse hex color string to RGB
fn parse_hex_color(hex: &str) -> Result<(u8, u8, u8), Box<dyn Error>> {
    if !hex.starts_with('#') || hex.len() != 7 {
        return Err("Invalid hex color format".into());
    }

    let r = u8::from_str_radix(&hex[1..3], 16)?;
    let g = u8::from_str_radix(&hex[3..5], 16)?;
    let b = u8::from_str_radix(&hex[5..7], 16)?;

    Ok((r, g, b))
}

/// Interpolate between two hex colors
fn interpolate_hex_colors(
    color1: &str,
    color2: &str,
    t: f32,
) -> Result<(u8, u8, u8), Box<dyn Error>> {
    let (r1, g1, b1) = parse_hex_color(color1)?;
    let (r2, g2, b2) = parse_hex_color(color2)?;

    let t = t.clamp(0.0, 1.0);
    let r = (r1 as f32 + (r2 as f32 - r1 as f32) * t) as u8;
    let g = (g1 as f32 + (g2 as f32 - g1 as f32) * t) as u8;
    let b = (b1 as f32 + (b2 as f32 - b1 as f32) * t) as u8;

    Ok((r, g, b))
}

/// Validate configuration parameters
fn validate_config<const N: usize>(
    config: &FieldVisualizationConfig,
) -> Result<(), Box<dyn Error>> {
    if N > 3 {
        return Err("Field visualization only supports 2D and 3D fields".into());
    }

    if config.resolution.iter().any(|&r| r < 2) {
        return Err("Resolution must be at least 2 in each dimension".into());
    }

    if config.opacity < 0.0 || config.opacity > 1.0 {
        return Err("Opacity must be between 0.0 and 1.0".into());
    }

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::csg::fields::Field2D;
    use backend_macro::with_backend;

    #[with_backend]
    #[test]
    fn test_field_visualization() {
        let device = device();
        let circle = Field2D::<Backend>::circle(1.0, device);

        let config = FieldVisualizationConfig {
            resolution: [16, 16, 2],
            bounds: [[-2.0, -2.0, 0.0], [2.0, 2.0, 0.0]],
            color_scheme: ColorScheme::Viridis,
            opacity: 0.8,
            mode: RenderingMode::VolumeGradient,
        };

        let result = visualize_field(&circle, &config);
        println!("{:?}", result);

        assert!(result.is_ok());

        let visualization = result.unwrap();
        assert_eq!(visualization.metadata.dimension, 2);
        assert_eq!(visualization.metadata.sample_count, 256); // 16x16
    }

    #[test]
    fn test_viridis_colormap() {
        let (r, g, b) = viridis_colormap(0.0);
        assert_eq!((r, g, b), (68, 1, 84)); // Should be deep purple

        let (r, g, b) = viridis_colormap(1.0);
        assert_eq!((r, g, b), (253, 231, 37)); // Should be yellow
    }

    #[test]
    fn test_config_validation() {
        let config = FieldVisualizationConfig {
            opacity: 1.5,
            ..Default::default()
        };

        let result = validate_config::<2>(&config);
        assert!(result.is_err());
    }
}