crater/analysis/rays/

animation.rs

//! Ray animation functionality - requires serde-threejs feature

use super::types::*;
use burn::prelude::*;
use serde_json::json;
use std::path::Path;

/// Configuration for ray animation generation
#[derive(Debug, Clone)]
pub struct RayAnimationConfig {
    /// Duration of the animation in milliseconds
    pub duration_ms: u32,
    /// Frame rate for the animation
    pub fps: u32,
    /// Color scheme for rays (hex colors)
    pub ray_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,
    /// Ray line width
    pub ray_width: f32,
    /// Size of hit point dots
    pub hit_dot_size: f32,
    /// Show successful hits as dots instead of lines
    pub show_hits_as_dots: bool,
    /// Maximum number of rays to animate (for performance)
    pub max_rays: Option<usize>,
    /// Distance to show failed rays (instead of infinite extension)
    pub failed_ray_distance: f32,
}

impl Default for RayAnimationConfig {
    fn default() -> Self {
        Self {
            duration_ms: 3000,
            fps: 60,
            ray_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(),
            ray_width: 2.0,
            hit_dot_size: 0.02,
            show_hits_as_dots: true,
            max_rays: Some(100_000),  // Limit for performance
            failed_ray_distance: 2.0, // Show failed rays extending this distance
        }
    }
}

/// Generate a Three.js-based animation from ray cast results
pub fn generate_ray_animation<B: Backend, const N: usize>(
    result: &RayCastResult<B, N>,
    config: &RayAnimationConfig,
    output_path: &Path,
) -> Result<(), Box<dyn std::error::Error>> {
    // Extract ray data
    let origins = result.rays().origins().to_data();
    let extensions = result.extensions().to_data();
    let distances = result.distances().to_data();
    let region_indices = result.region_indices().to_data();

    let ray_count = result.rays().ray_count();
    let actual_ray_count = if let Some(max_rays) = config.max_rays {
        ray_count.min(max_rays)
    } else {
        ray_count
    };

    // Convert tensor data to vectors
    let origins_vec: Vec<f32> = origins.iter::<f32>().collect();
    let extensions_vec: Vec<f32> = extensions.iter::<f32>().collect();
    let distances_vec: Vec<f32> = distances.iter::<f32>().collect();
    let region_indices_vec: Vec<f32> = region_indices.iter::<f32>().collect();

    let mut rays_data = Vec::new();
    for i in 0..actual_ray_count {
        let origin = if N == 2 {
            [origins_vec[i * N], origins_vec[i * N + 1], 0.0]
        } else {
            [
                origins_vec[i * N],
                origins_vec[i * N + 1],
                origins_vec[i * N + 2],
            ]
        };

        let raw_extension = if N == 2 {
            [extensions_vec[i * N], extensions_vec[i * N + 1], 0.0]
        } else {
            [
                extensions_vec[i * N],
                extensions_vec[i * N + 1],
                extensions_vec[i * N + 2],
            ]
        };

        // Check if extension coordinates are reasonable (not RAY_FAIL_VALUE)
        let extension_is_valid = raw_extension
            .iter()
            .all(|&x| x.is_finite() && x.abs() < 1000.0);

        let extension = if extension_is_valid {
            raw_extension
        } else {
            // Extension is invalid (probably RAY_FAIL_VALUE), generate a reasonable direction
            let angle = (i as f32) * 0.618034; // Golden angle for good distribution
            [
                origin[0] + angle.cos() * config.failed_ray_distance,
                origin[1] + angle.sin() * config.failed_ray_distance,
                origin[2] + (angle * 0.5).sin() * config.failed_ray_distance,
            ]
        };

        let distance = distances_vec[i];
        let region_idx = region_indices_vec[i] as usize;

        // Determine if this is a successful hit or failed ray
        let is_hit = extension_is_valid && distance.is_finite() && distance < 1000.0;
        let final_distance = if is_hit {
            distance
        } else {
            config.failed_ray_distance
        };

        // Use a different color for failed rays
        let color = if is_hit {
            config.ray_colors[region_idx % config.ray_colors.len()].clone()
        } else {
            "#666666".to_string() // Gray for failed rays
        };

        rays_data.push(json!({
            "origin": origin,
            "extension": extension,
            "distance": final_distance,
            "region": region_idx,
            "color": color,
            "isHit": is_hit
        }));
    }

    // Generate HTML with Three.js animation
    let html_content = generate_threejs_html(&rays_data, config)?;

    // Write to file
    std::fs::write(output_path, html_content)?;

    Ok(())
}

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

    let html = format!(
        r#"
<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>Ray Casting Animation</title>
    <script src="https://cdnjs.cloudflare.com/ajax/libs/three.js/r128/three.min.js"></script>
    <script src="https://cdn.jsdelivr.net/npm/three@0.128.0/examples/js/controls/OrbitControls.js"></script>
    <style>
        body {{
            margin: 0;
            padding: 0;
            background-color: {background_color};
            font-family: Arial, sans-serif;
            overflow: hidden;
        }}
        #container {{
            width: 100vw;
            height: 100vh;
            position: relative;
        }}
        #controls {{
            position: absolute;
            top: 10px;
            left: 10px;
            z-index: 1000;
            background: rgba(0, 0, 0, 0.7);
            padding: 10px;
            border-radius: 5px;
            color: white;
        }}
        button {{
            margin: 5px;
            padding: 8px 16px;
            background: #007acc;
            color: white;
            border: none;
            border-radius: 3px;
            cursor: pointer;
        }}
        button:hover {{
            background: #005a99;
        }}
        button:disabled {{
            background: #444;
            cursor: not-allowed;
        }}
        #info {{
            position: absolute;
            bottom: 10px;
            left: 10px;
            color: white;
            background: rgba(0, 0, 0, 0.7);
            padding: 10px;
            border-radius: 5px;
            font-size: 12px;
        }}
    </style>
</head>
<body>
    <div id="container">
        <div id="controls">
            <button id="playBtn" onclick="playAnimation()">Play</button>
            <button id="pauseBtn" onclick="pauseAnimation()" disabled>Pause</button>
            <button id="replayBtn" onclick="replayAnimation()" style="display:none;">Replay</button>
            <button onclick="resetAnimation()">Reset</button>
            <label>Speed: <input type="range" id="speedSlider" min="0.1" max="3" value="1" step="0.1" onchange="updateSpeed()"></label>
        </div>
        <div id="info">
            <div>Rays: {ray_count}</div>
            <div>Duration: {duration_ms}ms</div>
            <div>Use mouse to orbit, zoom, and pan</div>
        </div>
    </div>

    <script>
        // Configuration
        const config = {{
            duration: {duration_ms},
            fps: {fps},
            cameraPosition: {camera_position:?},
            cameraTarget: {camera_target:?},
            rayWidth: {ray_width},
            hitDotSize: {hit_dot_size},
            showHitsAsDots: {show_hits_as_dots}
        }};
        
        // Ray data
        const raysData = {rays_json};
        
        // Three.js setup
        let scene, camera, renderer, controls;
        let rayObjects = []; // Can contain both lines and dots
        let animationSpeed = 1.0;
        let isPlaying = false;
        let isCompleted = false;
        let startTime = 0;
        let pausedTime = 0;
        let lastProgress = -1; // Track last progress to avoid unnecessary updates
        
        function init() {{
            // Scene
            scene = new THREE.Scene();
            scene.background = new THREE.Color('{background_color}');
            
            // Camera
            camera = new THREE.PerspectiveCamera(75, window.innerWidth / window.innerHeight, 0.1, 1000);
            camera.position.set(config.cameraPosition[0], config.cameraPosition[1], config.cameraPosition[2]);
            camera.lookAt(config.cameraTarget[0], config.cameraTarget[1], config.cameraTarget[2]);
            
            // Renderer
            renderer = new THREE.WebGLRenderer({{ antialias: true }});
            renderer.setSize(window.innerWidth, window.innerHeight);
            document.getElementById('container').appendChild(renderer.domElement);
            
            // Controls
            controls = new THREE.OrbitControls(camera, renderer.domElement);
            controls.target.set(config.cameraTarget[0], config.cameraTarget[1], config.cameraTarget[2]);
            controls.enableDamping = true;
            controls.dampingFactor = 0.05;
            
            // Create ray objects (lines and/or dots)
            createRayObjects();
            
            // Start animation loop
            animate();
        }}
        
        function createRayObjects() {{
            // Group rays by color for instanced rendering
            const raysByColor = new Map();
            
            raysData.forEach((rayData, index) => {{
                // Validate coordinates to prevent extreme values
                const validateCoord = (coord) => {{
                    return isFinite(coord) && Math.abs(coord) < 1000 ? coord : 0;
                }};
                
                const origin = new THREE.Vector3(
                    validateCoord(rayData.origin[0]),
                    validateCoord(rayData.origin[1]),
                    validateCoord(rayData.origin[2])
                );
                const extension = new THREE.Vector3(
                    validateCoord(rayData.extension[0]),
                    validateCoord(rayData.extension[1]),
                    validateCoord(rayData.extension[2])
                );
                
                const color = rayData.color || '#ffffff';
                if (!raysByColor.has(color)) {{
                    raysByColor.set(color, []);
                }}
                raysByColor.get(color).push({{
                    origin: origin,
                    extension: extension,
                    rayData: rayData
                }});
            }});
            
            // Create instanced meshes for each color group
            raysByColor.forEach((rays, color) => {{
                const count = rays.length;
                const geometry = new THREE.SphereGeometry(config.hitDotSize, 4, 3); // Minimal segments for performance
                const material = new THREE.MeshBasicMaterial({{ color: color }});
                
                const instancedMesh = new THREE.InstancedMesh(geometry, material, count);
                const matrix = new THREE.Matrix4();
                
                // Store ray data in userData for animation
                instancedMesh.userData = {{
                    rays: rays,
                    type: 'instanced_dots'
                }};
                
                // Initialize positions at origins
                rays.forEach((ray, i) => {{
                    matrix.setPosition(ray.origin);
                    instancedMesh.setMatrixAt(i, matrix);
                }});
                instancedMesh.instanceMatrix.needsUpdate = true;
                
                scene.add(instancedMesh);
                rayObjects.push(instancedMesh);
            }});
        }}
        
        function updateRayAnimation(progress) {{
            // Clamp progress between 0 and 1
            progress = Math.max(0, Math.min(1, progress));
            
            // Skip update if progress hasn't changed significantly
            const progressDelta = Math.abs(progress - lastProgress);
            if (progressDelta < 0.001 && progress !== 0 && progress !== 1) {{
                return;
            }}
            lastProgress = progress;
            
            const matrix = new THREE.Matrix4();
            const tempPos = new THREE.Vector3();
            
            rayObjects.forEach(obj => {{
                if (obj.userData.type === 'instanced_dots') {{
                    // Handle instanced mesh animation
                    const rays = obj.userData.rays;
                    rays.forEach((ray, i) => {{
                        tempPos.lerpVectors(ray.origin, ray.extension, progress);
                        matrix.setPosition(tempPos);
                        obj.setMatrixAt(i, matrix);
                    }});
                    obj.instanceMatrix.needsUpdate = true;
                }} else {{
                    // Handle individual dots (fallback)
                    const {{ origin, extension }} = obj.userData;
                    tempPos.lerpVectors(origin, extension, progress);
                    obj.position.copy(tempPos);
                }}
            }});
        }}
        
        function animate() {{
            requestAnimationFrame(animate);
            
            if (isPlaying && !isCompleted) {{
                const currentTime = Date.now();
                const elapsedTime = (currentTime - startTime + pausedTime) * animationSpeed;
                const progress = elapsedTime / config.duration;
                
                if (progress >= 1.0) {{
                    // Animation completed
                    updateRayAnimation(1.0);
                    isPlaying = false;
                    isCompleted = true;
                    updateControls();
                }} else {{
                    updateRayAnimation(progress);
                }}
            }}
            
            controls.update();
            renderer.render(scene, camera);
        }}
        
        function updateControls() {{
            const playBtn = document.getElementById('playBtn');
            const pauseBtn = document.getElementById('pauseBtn');
            const replayBtn = document.getElementById('replayBtn');
            
            if (isCompleted) {{
                playBtn.style.display = 'none';
                pauseBtn.style.display = 'none';
                replayBtn.style.display = 'inline-block';
            }} else if (isPlaying) {{
                playBtn.disabled = true;
                pauseBtn.disabled = false;
            }} else {{
                playBtn.disabled = false;
                pauseBtn.disabled = true;
            }}
        }}
        
        function playAnimation() {{
            if (!isPlaying && !isCompleted) {{
                isPlaying = true;
                startTime = Date.now();
                updateControls();
            }}
        }}
        
        function pauseAnimation() {{
            if (isPlaying) {{
                pausedTime += (Date.now() - startTime);
                isPlaying = false;
                updateControls();
            }}
        }}
        
        function replayAnimation() {{
            isCompleted = false;
            isPlaying = true;
            startTime = Date.now();
            pausedTime = 0;
            updateRayAnimation(0);
            updateControls();
        }}
        
        function resetAnimation() {{
            isPlaying = false;
            isCompleted = false;
            startTime = 0;
            pausedTime = 0;
            updateRayAnimation(0);
            updateControls();
        }}
        
        function updateSpeed() {{
            const slider = document.getElementById('speedSlider');
            animationSpeed = parseFloat(slider.value);
        }}
        
        // Handle window resize
        window.addEventListener('resize', () => {{
            camera.aspect = window.innerWidth / window.innerHeight;
            camera.updateProjectionMatrix();
            renderer.setSize(window.innerWidth, window.innerHeight);
        }});
        
        // Initialize
        init();
        updateControls();
        
        // Auto-play animation after a short delay
        setTimeout(() => {{
            playAnimation();
        }}, 500);
    </script>
</body>
</html>
    "#,
        background_color = config.background_color,
        ray_count = rays_data.len(),
        duration_ms = config.duration_ms,
        fps = config.fps,
        camera_position = config.camera_position,
        camera_target = config.camera_target,
        ray_width = config.ray_width,
        hit_dot_size = config.hit_dot_size,
        show_hits_as_dots = if config.show_hits_as_dots {
            "true"
        } else {
            "false"
        },
        rays_json = rays_json
    );

    Ok(html)
}

/// Convenience function to create an animation with default settings
pub fn create_default_animation<B: Backend, const N: usize>(
    result: &RayCastResult<B, N>,
    output_path: &Path,
) -> Result<(), Box<dyn std::error::Error>> {
    let config = RayAnimationConfig::default();
    generate_ray_animation(result, &config, output_path)
}

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

    #[test]
    fn test_animation_config_default() {
        let config = RayAnimationConfig::default();
        assert_eq!(config.duration_ms, 3000);
        assert_eq!(config.fps, 60);
        assert!(!config.ray_colors.is_empty());
    }
}