//! Demo utilities for image processing and visualization
//!
//! This module provides utilities for demonstrating NMF capabilities,
//! including image processing and interactive demonstrations.

use crate::{NMFConfig, NMFDecomposer, Result};
use rtx_tensor::{Device, Tensor};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// Image processor for NMF demonstrations
pub struct ImageProcessor {
    device: Device,
}

impl ImageProcessor {
    /// Create new image processor
    pub fn new(device: Device) -> Self {
        Self { device }
    }

    /// Create mock image data for demonstration
    pub fn create_mock_image(&self, width: usize, height: usize) -> Result<Tensor> {
        // Create a synthetic image with some patterns
        let mut image_data = Vec::with_capacity(height * width);

        for y in 0..height {
            for x in 0..width {
                // Create a pattern with some geometric shapes
                let center_x = width as f32 / 2.0;
                let center_y = height as f32 / 2.0;

                let dist_from_center =
                    ((x as f32 - center_x).powi(2) + (y as f32 - center_y).powi(2)).sqrt();
                let max_dist = (center_x.powi(2) + center_y.powi(2)).sqrt();

                // Create concentric circles pattern
                let circle_value = (dist_from_center / max_dist * 4.0 * std::f32::consts::PI)
                    .sin()
                    .abs();

                // Add some diagonal stripes
                let stripe_value = ((x + y) as f32 / 10.0).sin().abs();

                // Combine patterns
                let pixel_value = (circle_value * 0.7 + stripe_value * 0.3) * 255.0;
                image_data.push(pixel_value.max(0.0));
            }
        }

        Ok(Tensor::from_data(
            image_data,
            [height, width],
            &self.device,
        )?)
    }

    /// Create mock face-like image data
    pub fn create_mock_faces(
        &self,
        num_faces: usize,
        face_height: usize,
        face_width: usize,
    ) -> Result<Tensor> {
        let mut faces_data = Vec::with_capacity(num_faces * face_height * face_width);

        for face_idx in 0..num_faces {
            // Create different face-like patterns
            let base_brightness = 50.0 + (face_idx as f32 * 30.0) % 150.0;

            for y in 0..face_height {
                for x in 0..face_width {
                    let center_x = face_width as f32 / 2.0;
                    let center_y = face_height as f32 / 2.0;

                    // Create oval face shape
                    let face_x = (x as f32 - center_x) / (face_width as f32 * 0.4);
                    let face_y = (y as f32 - center_y) / (face_height as f32 * 0.5);
                    let face_dist = (face_x.powi(2) + face_y.powi(2)).sqrt();

                    let mut pixel_value = base_brightness;

                    // Add face features if inside face boundary
                    if face_dist < 1.0 {
                        // Eyes (darker regions)
                        let eye_y = center_y - face_height as f32 * 0.15;
                        let left_eye_x = center_x - face_width as f32 * 0.15;
                        let right_eye_x = center_x + face_width as f32 * 0.15;

                        let left_eye_dist =
                            ((x as f32 - left_eye_x).powi(2) + (y as f32 - eye_y).powi(2)).sqrt();
                        let right_eye_dist =
                            ((x as f32 - right_eye_x).powi(2) + (y as f32 - eye_y).powi(2)).sqrt();

                        if left_eye_dist < face_width as f32 * 0.05
                            || right_eye_dist < face_width as f32 * 0.05
                        {
                            pixel_value *= 0.3; // Dark eyes
                        }

                        // Mouth (darker region)
                        let mouth_y = center_y + face_height as f32 * 0.2;
                        let mouth_dist =
                            ((x as f32 - center_x).powi(2) + (y as f32 - mouth_y).powi(2)).sqrt();
                        if mouth_dist < face_width as f32 * 0.08 {
                            pixel_value *= 0.5; // Dark mouth
                        }

                        // Add some noise for realism
                        let noise = ((x * 7 + y * 11 + face_idx * 13) as f32).sin() * 10.0;
                        pixel_value += noise;
                    } else {
                        pixel_value = 0.0; // Background
                    }

                    faces_data.push(pixel_value.max(0.0));
                }
            }
        }

        Ok(Tensor::from_data(
            faces_data,
            [num_faces, face_height * face_width],
            &self.device,
        )?)
    }

    /// Convert tensor back to image representation for display
    pub fn tensor_to_image_info(&self, tensor: &Tensor) -> Result<ImageInfo> {
        let shape = tensor.shape();
        let data = tensor.to_cpu()?;

        // Calculate real statistics
        let min_value = data.iter().fold(f32::INFINITY, |a, &b| a.min(b));
        let max_value = data.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b));
        let mean_value = data.iter().sum::<f32>() / data.len() as f32;

        Ok(ImageInfo {
            shape: shape.dims().to_vec(),
            min_value,
            max_value,
            mean_value,
            data_size: shape.numel(),
        })
    }
}

/// Information about an image tensor for display purposes
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ImageInfo {
    pub shape: Vec<usize>,
    pub min_value: f32,
    pub max_value: f32,
    pub mean_value: f32,
    pub data_size: usize,
}

/// NMF demo runner for demonstrations
pub struct NMFDemo {
    config: NMFConfig,
    decomposer: NMFDecomposer,
    image_processor: ImageProcessor,
}

impl NMFDemo {
    /// Create new NMF demo
    pub fn new(device: Device) -> Self {
        let config = NMFConfig::new()
            .with_components(10)
            .with_max_iterations(50)
            .with_tolerance(1e-4);

        let decomposer = NMFDecomposer::new(config.clone());
        let image_processor = ImageProcessor::new(device);

        Self {
            config,
            decomposer,
            image_processor,
        }
    }

    /// Run image decomposition demo
    pub fn run_image_demo(&mut self, width: usize, height: usize) -> Result<DemoResult> {
        println!("🖼️ Creating mock image data ({width} x {height})...");
        let image = self.image_processor.create_mock_image(width, height)?;
        let image_info = self.image_processor.tensor_to_image_info(&image)?;

        println!("📊 Image statistics:");
        println!("  Shape: {:?}", image_info.shape);
        println!(
            "  Range: [{:.2}, {:.2}]",
            image_info.min_value, image_info.max_value
        );
        println!("  Mean: {:.2}", image_info.mean_value);

        println!(
            "\n⚡ Running NMF decomposition with {} components...",
            self.config.components()
        );
        let result = self.decomposer.fit_transform_detailed(&image)?;

        let w_info = self.image_processor.tensor_to_image_info(&result.w)?;
        let h_info = self.image_processor.tensor_to_image_info(&result.h)?;

        println!("\n✅ NMF decomposition completed!");
        println!("  Iterations: {}", result.iterations);
        println!("  Converged: {}", result.converged);
        println!("  Error: {:.6}", result.reconstruction_error);
        println!("  Time: {:.3}s", result.computation_time);

        Ok(DemoResult {
            original_image: image_info,
            w_matrix: w_info,
            h_matrix: h_info,
            reconstruction_error: result.reconstruction_error,
            iterations: result.iterations,
            converged: result.converged,
            computation_time: result.computation_time,
            components: self.config.components(),
        })
    }

    /// Run faces decomposition demo
    pub fn run_faces_demo(
        &mut self,
        num_faces: usize,
        face_height: usize,
        face_width: usize,
    ) -> Result<DemoResult> {
        println!(
            "👥 Creating mock face dataset ({num_faces} faces, {face_height} x {face_width})..."
        );
        let faces = self
            .image_processor
            .create_mock_faces(num_faces, face_height, face_width)?;
        let faces_info = self.image_processor.tensor_to_image_info(&faces)?;

        println!("📊 Faces dataset statistics:");
        println!("  Shape: {:?}", faces_info.shape);
        println!(
            "  Range: [{:.2}, {:.2}]",
            faces_info.min_value, faces_info.max_value
        );
        println!("  Mean: {:.2}", faces_info.mean_value);

        // Update config for faces (usually need fewer components)
        self.config = self.config.clone().with_components(5);
        self.decomposer = NMFDecomposer::new(self.config.clone());

        println!("\n⚡ Running NMF decomposition to find face basis...");
        let result = self.decomposer.fit_transform_detailed(&faces)?;

        let w_info = self.image_processor.tensor_to_image_info(&result.w)?;
        let h_info = self.image_processor.tensor_to_image_info(&result.h)?;

        println!("\n✅ Face decomposition completed!");
        println!("  Found {} basis faces", self.config.components());
        println!("  Iterations: {}", result.iterations);
        println!("  Reconstruction error: {:.6}", result.reconstruction_error);
        println!("  Time: {:.3}s", result.computation_time);

        Ok(DemoResult {
            original_image: faces_info,
            w_matrix: w_info,
            h_matrix: h_info,
            reconstruction_error: result.reconstruction_error,
            iterations: result.iterations,
            converged: result.converged,
            computation_time: result.computation_time,
            components: self.config.components(),
        })
    }

    /// Get performance metrics for display
    pub fn get_performance_metrics(&self) -> HashMap<String, String> {
        let mut metrics = HashMap::new();

        metrics.insert(
            "gpu_acceleration".to_string(),
            if self.config.use_gpu() {
                "Enabled"
            } else {
                "Disabled"
            }
            .to_string(),
        );
        metrics.insert(
            "components".to_string(),
            self.config.components().to_string(),
        );
        metrics.insert(
            "max_iterations".to_string(),
            self.config.max_iterations().to_string(),
        );
        metrics.insert(
            "tolerance".to_string(),
            format!("{:.0e}", self.config.tolerance()),
        );
        metrics.insert(
            "device".to_string(),
            format!("{:?}", self.image_processor.device),
        );

        // Add system-specific performance indicators
        match &self.image_processor.device {
            Device::Cuda(_) => {
                metrics.insert("cuda_version".to_string(), "13.0+".to_string());
                metrics.insert("memory_type".to_string(), "GPU VRAM".to_string());
                metrics.insert("acceleration_type".to_string(), "CUDA Kernels".to_string());
            }
            Device::Rocm(_) => {
                metrics.insert("acceleration_type".to_string(), "ROCm Kernels".to_string());
                metrics.insert("memory_type".to_string(), "GPU VRAM".to_string());
            }
            Device::Metal(_) => {
                metrics.insert("acceleration_type".to_string(), "Metal Kernels".to_string());
                metrics.insert("memory_type".to_string(), "GPU Memory".to_string());
            }
            Device::Cpu => {
                metrics.insert("acceleration_type".to_string(), "CPU".to_string());
                metrics.insert("memory_type".to_string(), "System RAM".to_string());
            }
        }

        metrics
    }
}

/// Result of a demo run
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DemoResult {
    pub original_image: ImageInfo,
    pub w_matrix: ImageInfo,
    pub h_matrix: ImageInfo,
    pub reconstruction_error: f32,
    pub iterations: usize,
    pub converged: bool,
    pub computation_time: f32,
    pub components: usize,
}

impl DemoResult {
    /// Format result for display
    pub fn format_summary(&self) -> String {
        format!(
            "NMF Result: {} components, {:.6} error, {} iterations ({:.3}s)",
            self.components, self.reconstruction_error, self.iterations, self.computation_time
        )
    }
}

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

    #[test]
    fn test_image_processor() -> Result<()> {
        let device = Device::Cuda(0);
        let processor = ImageProcessor::new(device);

        let image = processor.create_mock_image(32, 24)?;
        assert_eq!(image.shape(), &[24, 32]);

        let info = processor.tensor_to_image_info(&image)?;
        assert_eq!(info.shape, vec![24, 32]);
        assert!(info.min_value >= 0.0);

        Ok(())
    }

    #[test]
    fn test_mock_faces() -> Result<()> {
        let device = Device::Cuda(0);
        let processor = ImageProcessor::new(device);

        let faces = processor.create_mock_faces(5, 20, 20)?;
        assert_eq!(faces.shape(), &[5, 400]); // 5 faces, 20x20 = 400 pixels each

        Ok(())
    }

    #[test]
    fn test_nmf_demo() -> Result<()> {
        let device = Device::Cuda(0);
        let mut demo = NMFDemo::new(device);

        // Small demo to avoid long test times
        let result = demo.run_image_demo(8, 6)?;

        assert_eq!(result.original_image.shape, vec![6, 8]);
        assert!(result.reconstruction_error >= 0.0);
        assert!(result.computation_time > 0.0);

        Ok(())
    }

    #[test]
    fn test_performance_metrics() {
        let device = Device::Cuda(0);
        let demo = NMFDemo::new(device);
        let metrics = demo.get_performance_metrics();

        assert!(metrics.contains_key("gpu_acceleration"));
        assert!(metrics.contains_key("components"));
        assert!(metrics.contains_key("device"));
    }
}