crater/serde/threejs/

animation.rs

use crate::csg::trajectory::Trajectories;
use burn::prelude::*;
use serde_json::json;
use std::path::Path;

/// Animation type determines coloring strategy and data source
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum AnimationType {
    /// Use custom color/opacity pairs for each trajectory
    CustomColors,
    /// Derive colors from raycast results (hit/miss, region-based)
    RaycastResults,
    /// Use default uniform coloring scheme
    Uniform,
    /// Color by trajectory index with optional grouping
    Indexed { group_size: Option<usize> },
    /// Color by distance/time along trajectory
    Gradient {
        start_color: String,
        end_color: String,
    },
}

/// Configuration for trajectory animation generation
#[derive(Debug, Clone)]
pub struct TrajectoryAnimationConfig {
    /// Duration of the animation in milliseconds
    pub duration_ms: u32,
    /// Frame rate for the animation
    pub fps: u32,
    /// Color scheme for trajectories (hex colors)
    pub colors: Vec<String>,
    /// Camera position for the 3D view
    pub camera_position: [f32; 3],
    /// Camera target (look-at point)
    pub camera_target: [f32; 3],
    /// Background color (hex)
    pub background_color: String,
    /// Size of trajectory dots
    pub dot_size: f32,
    /// Maximum number of trajectories to animate (for performance)
    pub max_trajectories: Option<usize>,
    /// Opacity for failed trajectories
    pub failed_opacity: f32,
    /// Animation type determining coloring strategy
    pub animation_type: AnimationType,
}

/// Color source for trajectory animations
#[derive(Clone)]
pub enum ColorSource<'a, B: Backend, const N: usize> {
    /// Custom color/opacity pairs
    Custom(&'a [(String, f32)]),
    /// Raycast result for hit/miss coloring
    Raycast(&'a crate::analysis::prelude::RayCastResult<B, N>),
    /// Use configuration colors uniformly
    Config,
}

impl Default for TrajectoryAnimationConfig {
    fn default() -> Self {
        Self {
            duration_ms: 3000,
            fps: 60,
            colors: vec![
                "#ff6b6b".to_string(), // Red
                "#4ecdc4".to_string(), // Teal
                "#45b7d1".to_string(), // Blue
                "#96ceb4".to_string(), // Green
                "#feca57".to_string(), // Yellow
                "#ff9ff3".to_string(), // Pink
                "#54a0ff".to_string(), // Light blue
            ],
            camera_position: [10.0, 10.0, 10.0],
            camera_target: [0.0, 0.0, 0.0],
            background_color: "#1a1a1a".to_string(),
            dot_size: 0.02,
            max_trajectories: Some(100_000), // Limit for performance
            failed_opacity: 0.3,
            animation_type: AnimationType::Uniform,
        }
    }
}

impl AnimationType {
    /// Create a gradient animation type with predefined color schemes
    pub fn gradient_preset(preset: &str) -> Self {
        match preset {
            "fire" => Self::Gradient {
                start_color: "#ff4757".to_string(),
                end_color: "#ffa502".to_string(),
            },
            "ocean" => Self::Gradient {
                start_color: "#3742fa".to_string(),
                end_color: "#2ed573".to_string(),
            },
            "sunset" => Self::Gradient {
                start_color: "#ff6b9d".to_string(),
                end_color: "#ffa726".to_string(),
            },
            "forest" => Self::Gradient {
                start_color: "#00d2d3".to_string(),
                end_color: "#54a0ff".to_string(),
            },
            _ => Self::Gradient {
                start_color: "#ff6b6b".to_string(),
                end_color: "#4ecdc4".to_string(),
            },
        }
    }
}

/// Generate trajectory colors based on the specified animation type and color source
pub fn generate_trajectory_colors<B: Backend, const N: usize>(
    trajectories: &Trajectories<B, N>,
    color_source: ColorSource<'_, B, N>,
    config: &TrajectoryAnimationConfig,
) -> Result<Vec<(String, f32)>, Box<dyn std::error::Error>> {
    if trajectories.is_empty() {
        return Ok(Vec::new());
    }

    let first_step = &trajectories[0];
    let trajectory_count = first_step.shape().dims::<2>()[0];

    match (color_source, &config.animation_type) {
        (ColorSource::Custom(colors), AnimationType::CustomColors) => {
            if colors.len() != trajectory_count {
                return Err(format!(
                    "Custom color count ({}) does not match trajectory count ({})",
                    colors.len(),
                    trajectory_count
                )
                .into());
            }
            Ok(colors.to_vec())
        }
        (ColorSource::Raycast(raycast_result), AnimationType::RaycastResults) => {
            compute_raycast_colors(raycast_result, config)
        }
        (_, AnimationType::Uniform) => {
            let mut colors = Vec::with_capacity(trajectory_count);
            for i in 0..trajectory_count {
                let color = config.colors[i % config.colors.len()].clone();
                colors.push((color, 1.0));
            }
            Ok(colors)
        }
        (_, AnimationType::Indexed { group_size }) => {
            let mut colors = Vec::with_capacity(trajectory_count);
            let group_size = group_size.unwrap_or(config.colors.len());
            for i in 0..trajectory_count {
                let color_idx = (i / group_size) % config.colors.len();
                let color = config.colors[color_idx].clone();
                colors.push((color, 1.0));
            }
            Ok(colors)
        }
        (
            _,
            AnimationType::Gradient {
                start_color,
                end_color,
            },
        ) => generate_gradient_colors(trajectory_count, start_color, end_color),
        _ => Err("Incompatible color source and animation type combination".into()),
    }
}

/// Generate gradient colors between two hex colors
fn generate_gradient_colors(
    count: usize,
    start_hex: &str,
    end_hex: &str,
) -> Result<Vec<(String, f32)>, Box<dyn std::error::Error>> {
    let start_rgb = hex_to_rgb(start_hex)?;
    let end_rgb = hex_to_rgb(end_hex)?;

    let mut colors = Vec::with_capacity(count);
    for i in 0..count {
        let t = if count > 1 {
            i as f32 / (count - 1) as f32
        } else {
            0.0
        };
        let r = (start_rgb[0] as f32 * (1.0 - t) + end_rgb[0] as f32 * t) as u8;
        let g = (start_rgb[1] as f32 * (1.0 - t) + end_rgb[1] as f32 * t) as u8;
        let b = (start_rgb[2] as f32 * (1.0 - t) + end_rgb[2] as f32 * t) as u8;
        let hex = format!("#{:02x}{:02x}{:02x}", r, g, b);
        colors.push((hex, 1.0));
    }
    Ok(colors)
}

/// Convert hex color to RGB
fn hex_to_rgb(hex: &str) -> Result<[u8; 3], Box<dyn std::error::Error>> {
    let hex = hex.trim_start_matches('#');
    if hex.len() != 6 {
        return Err("Invalid hex color format".into());
    }

    let r = u8::from_str_radix(&hex[0..2], 16)?;
    let g = u8::from_str_radix(&hex[2..4], 16)?;
    let b = u8::from_str_radix(&hex[4..6], 16)?;
    Ok([r, g, b])
}

/// Core unified trajectory animation generation function
pub fn generate_trajectory_animation<B: Backend, const N: usize>(
    trajectories: &Trajectories<B, N>,
    color_source: ColorSource<'_, B, N>,
    config: &TrajectoryAnimationConfig,
    output_path: &Path,
) -> Result<(), Box<dyn std::error::Error>> {
    if trajectories.is_empty() {
        return Err("No trajectories to animate".into());
    }

    // Generate colors based on animation type and source
    let colors = generate_trajectory_colors(trajectories, color_source, config)?;

    // Generate trajectory data for rendering
    let trajectories_data = convert_trajectories_to_json::<B, N>(trajectories, &colors, config)?;

    // Generate and write HTML
    let html_content = generate_threejs_html(&trajectories_data, config)?;
    std::fs::write(output_path, html_content)?;

    Ok(())
}

/// Convert trajectories to JSON format for Three.js rendering
fn convert_trajectories_to_json<B: Backend, const N: usize>(
    trajectories: &Trajectories<B, N>,
    colors: &[(String, f32)],
    config: &TrajectoryAnimationConfig,
) -> Result<Vec<serde_json::Value>, Box<dyn std::error::Error>> {
    let first_step = &trajectories[0];
    let trajectory_count = first_step.shape().dims::<2>()[0];
    let actual_trajectory_count = config
        .max_trajectories
        .unwrap_or(trajectory_count)
        .min(trajectory_count);

    // Pre-process all trajectory data efficiently
    let step_data_vecs: Vec<Vec<f32>> = trajectories
        .iter()
        .map(|step| {
            let step_data = step.to_data();
            step_data.iter::<f32>().collect()
        })
        .collect();

    // Build trajectory JSON data
    let mut trajectories_data = Vec::with_capacity(actual_trajectory_count);

    for traj_idx in 0..actual_trajectory_count {
        let mut path = Vec::with_capacity(trajectories.len());

        for step_vec in &step_data_vecs {
            let position = if N == 2 {
                [step_vec[traj_idx * N], step_vec[traj_idx * N + 1], 0.0]
            } else {
                [
                    step_vec[traj_idx * N],
                    step_vec[traj_idx * N + 1],
                    step_vec[traj_idx * N + 2],
                ]
            };

            // Validate coordinates
            if position.iter().all(|&x| x.is_finite() && x.abs() < 1000.0) {
                path.push(position);
            }
        }

        if !path.is_empty() {
            let (color, opacity) = &colors[traj_idx.min(colors.len() - 1)];

            trajectories_data.push(json!({
                "path": path,
                "color": color,
                "opacity": opacity
            }));
        }
    }

    Ok(trajectories_data)
}

/// Generate a Three.js-based animation from trajectories with custom coloring (legacy)
pub fn generate_trajectory_animation_with_colors<B: Backend, const N: usize>(
    trajectories: &Trajectories<B, N>,
    colors: &[(String, f32)], // (color_hex, opacity) pairs
    config: &TrajectoryAnimationConfig,
    output_path: &Path,
) -> Result<(), Box<dyn std::error::Error>> {
    let mut custom_config = config.clone();
    custom_config.animation_type = AnimationType::CustomColors;

    generate_trajectory_animation(
        trajectories,
        ColorSource::<B, N>::Custom(colors),
        &custom_config,
        output_path,
    )
}

/// Compute colors from raycast results according to hit/miss status and region indices
pub fn compute_raycast_colors<B: Backend, const N: usize>(
    raycast_result: &crate::analysis::prelude::RayCastResult<B, N>,
    config: &TrajectoryAnimationConfig,
) -> Result<Vec<(String, f32)>, Box<dyn std::error::Error>> {
    use crate::analysis::prelude::RAY_FAIL_VALUE;

    let ray_count = raycast_result.rays().ray_count();
    let distances = raycast_result.distances().to_data();
    let region_indices = raycast_result.region_indices().to_data();
    let extensions = raycast_result.extensions().to_data();

    let distances_vec: Vec<f32> = distances.iter::<f32>().collect();
    let region_indices_vec: Vec<f32> = region_indices.iter::<f32>().collect();
    let extensions_vec: Vec<f32> = extensions.iter::<f32>().collect();

    let mut colors = Vec::with_capacity(ray_count);

    for ray_idx in 0..ray_count {
        let distance = distances_vec[ray_idx];
        let region_idx = region_indices_vec[ray_idx] as usize;

        // Check if extension coordinates are valid (not RAY_FAIL_VALUE)
        let extension_start = ray_idx * N;
        let extension_is_valid = (0..N).all(|i| {
            let coord = extensions_vec[extension_start + i];
            coord.is_finite() && coord.abs() < 1000.0 && coord != RAY_FAIL_VALUE
        });

        let is_hit = extension_is_valid
            && distance.is_finite()
            && distance != RAY_FAIL_VALUE
            && distance < 1000.0;

        let (color, opacity) = if is_hit {
            // Use region-based coloring for hits
            let color = config.colors[region_idx % config.colors.len()].clone();
            (color, 1.0)
        } else {
            // Gray color for failed rays with reduced opacity
            ("#666666".to_string(), config.failed_opacity)
        };

        colors.push((color, opacity));
    }

    Ok(colors)
}

/// Generate a Three.js-based animation from trajectories with raycast context for coloring (legacy)
pub fn generate_trajectory_animation_with_raycast<B: Backend, const N: usize>(
    trajectories: &Trajectories<B, N>,
    raycast_result: &crate::analysis::prelude::RayCastResult<B, N>,
    config: &TrajectoryAnimationConfig,
    output_path: &Path,
) -> Result<(), Box<dyn std::error::Error>> {
    let mut raycast_config = config.clone();
    raycast_config.animation_type = AnimationType::RaycastResults;

    generate_trajectory_animation(
        trajectories,
        ColorSource::<B, N>::Raycast(raycast_result),
        &raycast_config,
        output_path,
    )
}

/// Generate a Three.js-based animation from trajectories with uniform coloring (legacy method)
pub fn generate_trajectory_animation_uniform<B: Backend, const N: usize>(
    trajectories: &Trajectories<B, N>,
    config: &TrajectoryAnimationConfig,
    output_path: &Path,
) -> Result<(), Box<dyn std::error::Error>> {
    let mut uniform_config = config.clone();
    uniform_config.animation_type = AnimationType::Uniform;

    generate_trajectory_animation(
        trajectories,
        ColorSource::<B, N>::Config,
        &uniform_config,
        output_path,
    )
}

fn generate_threejs_html(
    trajectories_data: &[serde_json::Value],
    config: &TrajectoryAnimationConfig,
) -> Result<String, Box<dyn std::error::Error>> {
    let trajectories_json = serde_json::to_string(trajectories_data)?;

    // Load HTML template and replace placeholders
    let template = include_str!("animation_template.html");
    let html = template
        .replace("BACKGROUND_COLOR_PLACEHOLDER", &config.background_color)
        .replace(
            "TRAJECTORY_COUNT_PLACEHOLDER",
            &trajectories_data.len().to_string(),
        )
        .replace("DURATION_MS_PLACEHOLDER", &config.duration_ms.to_string())
        .replace("FPS_PLACEHOLDER", &config.fps.to_string())
        .replace(
            "CAMERA_POSITION_PLACEHOLDER",
            &format!("{:?}", config.camera_position),
        )
        .replace(
            "CAMERA_TARGET_PLACEHOLDER",
            &format!("{:?}", config.camera_target),
        )
        .replace("DOT_SIZE_PLACEHOLDER", &config.dot_size.to_string())
        .replace("TRAJECTORIES_JSON_PLACEHOLDER", &trajectories_json);

    Ok(html)
}

/// Convenience function to create an animation with default uniform coloring
pub fn create_default_trajectory_animation<B: Backend, const N: usize>(
    trajectories: &Trajectories<B, N>,
    output_path: &Path,
) -> Result<(), Box<dyn std::error::Error>> {
    let config = TrajectoryAnimationConfig::default();
    generate_trajectory_animation(
        trajectories,
        ColorSource::<B, N>::Config,
        &config,
        output_path,
    )
}

/// Convenience function to create a gradient animation
pub fn create_gradient_animation<B: Backend, const N: usize>(
    trajectories: &Trajectories<B, N>,
    gradient_preset: &str,
    output_path: &Path,
) -> Result<(), Box<dyn std::error::Error>> {
    let config = TrajectoryAnimationConfig {
        animation_type: AnimationType::gradient_preset(gradient_preset),
        ..Default::default()
    };

    generate_trajectory_animation(
        trajectories,
        ColorSource::<B, N>::Config,
        &config,
        output_path,
    )
}

/// Convenience function to create a raycast-colored animation
pub fn create_raycast_animation<B: Backend, const N: usize>(
    trajectories: &Trajectories<B, N>,
    raycast_result: &crate::analysis::prelude::RayCastResult<B, N>,
    output_path: &Path,
) -> Result<(), Box<dyn std::error::Error>> {
    let config = TrajectoryAnimationConfig {
        animation_type: AnimationType::RaycastResults,
        ..Default::default()
    };

    generate_trajectory_animation(
        trajectories,
        ColorSource::<B, N>::Raycast(raycast_result),
        &config,
        output_path,
    )
}

/// Convenience function to create an indexed/grouped animation
pub fn create_indexed_animation<B: Backend, const N: usize>(
    trajectories: &Trajectories<B, N>,
    group_size: Option<usize>,
    output_path: &Path,
) -> Result<(), Box<dyn std::error::Error>> {
    let config = TrajectoryAnimationConfig {
        animation_type: AnimationType::Indexed { group_size },
        ..Default::default()
    };

    generate_trajectory_animation(
        trajectories,
        ColorSource::<B, N>::Config,
        &config,
        output_path,
    )
}

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

    #[test]
    fn test_animation_config_default() {
        let config = TrajectoryAnimationConfig::default();
        assert_eq!(config.duration_ms, 3000);
        assert_eq!(config.fps, 60);
        assert!(!config.colors.is_empty());
        assert_eq!(config.animation_type, AnimationType::Uniform);
    }

    #[test]
    fn test_animation_type_gradient_presets() {
        let fire = AnimationType::gradient_preset("fire");
        match fire {
            AnimationType::Gradient {
                start_color,
                end_color,
            } => {
                assert_eq!(start_color, "#ff4757");
                assert_eq!(end_color, "#ffa502");
            }
            _ => panic!("Expected gradient type"),
        }

        let unknown = AnimationType::gradient_preset("unknown");
        match unknown {
            AnimationType::Gradient {
                start_color,
                end_color,
            } => {
                assert_eq!(start_color, "#ff6b6b");
                assert_eq!(end_color, "#4ecdc4");
            }
            _ => panic!("Expected gradient type"),
        }
    }

    #[test]
    fn test_hex_to_rgb() {
        assert_eq!(hex_to_rgb("#ff0000").unwrap(), [255, 0, 0]);
        assert_eq!(hex_to_rgb("00ff00").unwrap(), [0, 255, 0]);
        assert_eq!(hex_to_rgb("#0000ff").unwrap(), [0, 0, 255]);
        assert!(hex_to_rgb("invalid").is_err());
    }

    #[test]
    fn test_generate_gradient_colors() {
        let colors = generate_gradient_colors(3, "#ff0000", "#0000ff").unwrap();
        assert_eq!(colors.len(), 3);
        assert_eq!(colors[0].0, "#ff0000");
        assert_eq!(colors[2].0, "#0000ff");
        assert_eq!(colors[0].1, 1.0); // opacity
    }
}