//! Image processing utilities for NMF color unmixing
//!
//! Provides image loading, preprocessing, and conversion utilities
//! specifically designed for color blind unmixing applications.

use crate::{NMFError, Result};
use image::{ImageReader, Rgb, RgbImage};
use nalgebra::DMatrix;
use std::path::Path;

/// RGB color representation
#[derive(Debug, Clone, Copy)]
pub struct Color {
    pub r: f32,
    pub g: f32,
    pub b: f32,
}

impl Color {
    pub fn new(r: f32, g: f32, b: f32) -> Self {
        Self { r, g, b }
    }

    pub fn from_u8(r: u8, g: u8, b: u8) -> Self {
        Self {
            r: r as f32 / 255.0,
            g: g as f32 / 255.0,
            b: b as f32 / 255.0,
        }
    }

    pub fn to_u8(&self) -> (u8, u8, u8) {
        (
            (self.r * 255.0).clamp(0.0, 255.0) as u8,
            (self.g * 255.0).clamp(0.0, 255.0) as u8,
            (self.b * 255.0).clamp(0.0, 255.0) as u8,
        )
    }
}

/// Image data for NMF processing
#[derive(Debug, Clone)]
pub struct ImageData {
    pub width: usize,
    pub height: usize,
    pub pixels: Vec<Color>,
}

impl ImageData {
    /// Load image from file path
    pub fn load_from_path<P: AsRef<Path>>(path: P) -> Result<Self> {
        let img = ImageReader::open(path)
            .map_err(|e| NMFError::configuration_error(format!("Failed to open image: {e}")))?
            .decode()
            .map_err(|e| NMFError::configuration_error(format!("Failed to decode image: {e}")))?
            .to_rgb8();

        let (width, height) = (img.width() as usize, img.height() as usize);
        let mut pixels = Vec::with_capacity(width * height);

        for pixel in img.pixels() {
            pixels.push(Color::from_u8(pixel[0], pixel[1], pixel[2]));
        }

        Ok(Self {
            width,
            height,
            pixels,
        })
    }

    /// Convert to RGB matrix for NMF processing
    /// Returns: (red_channel, green_channel, blue_channel) as separate matrices
    pub fn to_rgb_matrices(&self) -> (DMatrix<f32>, DMatrix<f32>, DMatrix<f32>) {
        let mut red_data = Vec::with_capacity(self.width * self.height);
        let mut green_data = Vec::with_capacity(self.width * self.height);
        let mut blue_data = Vec::with_capacity(self.width * self.height);

        for pixel in &self.pixels {
            red_data.push(pixel.r);
            green_data.push(pixel.g);
            blue_data.push(pixel.b);
        }

        let red_matrix = DMatrix::from_vec(self.height, self.width, red_data);
        let green_matrix = DMatrix::from_vec(self.height, self.width, green_data);
        let blue_matrix = DMatrix::from_vec(self.height, self.width, blue_data);

        (red_matrix, green_matrix, blue_matrix)
    }

    /// Convert to single matrix for processing (stacked RGB channels)
    /// Format: [R1, G1, B1, R2, G2, B2, ...] as height×(width*3) matrix
    pub fn to_stacked_matrix(&self) -> DMatrix<f32> {
        let mut data = Vec::with_capacity(self.width * self.height * 3);

        for y in 0..self.height {
            for x in 0..self.width {
                let idx = y * self.width + x;
                let pixel = &self.pixels[idx];
                data.push(pixel.r);
                data.push(pixel.g);
                data.push(pixel.b);
            }
        }

        DMatrix::from_vec(self.height, self.width * 3, data)
    }

    /// Create from RGB matrices (inverse of to_rgb_matrices)
    pub fn from_rgb_matrices(
        red: &DMatrix<f32>,
        green: &DMatrix<f32>,
        blue: &DMatrix<f32>,
    ) -> Result<Self> {
        if red.nrows() != green.nrows()
            || red.nrows() != blue.nrows()
            || red.ncols() != green.ncols()
            || red.ncols() != blue.ncols()
        {
            return Err(NMFError::configuration_error(
                "RGB matrices must have same dimensions",
            ));
        }

        let height = red.nrows();
        let width = red.ncols();
        let mut pixels = Vec::with_capacity(width * height);

        for y in 0..height {
            for x in 0..width {
                pixels.push(Color::new(
                    red[(y, x)].clamp(0.0, 1.0),
                    green[(y, x)].clamp(0.0, 1.0),
                    blue[(y, x)].clamp(0.0, 1.0),
                ));
            }
        }

        Ok(Self {
            width,
            height,
            pixels,
        })
    }

    /// Convert to displayable RGB image
    pub fn to_rgb_image(&self) -> RgbImage {
        let mut img = RgbImage::new(self.width as u32, self.height as u32);

        for (idx, pixel) in self.pixels.iter().enumerate() {
            let x = idx % self.width;
            let y = idx / self.width;
            let (r, g, b) = pixel.to_u8();
            img.put_pixel(x as u32, y as u32, Rgb([r, g, b]));
        }

        img
    }

    /// Get basic statistics about the image
    pub fn statistics(&self) -> ImageStats {
        let mut r_sum = 0.0f32;
        let mut g_sum = 0.0f32;
        let mut b_sum = 0.0f32;
        let mut r_min = 1.0f32;
        let mut g_min = 1.0f32;
        let mut b_min = 1.0f32;
        let mut r_max = 0.0f32;
        let mut g_max = 0.0f32;
        let mut b_max = 0.0f32;

        for pixel in &self.pixels {
            r_sum += pixel.r;
            g_sum += pixel.g;
            b_sum += pixel.b;
            r_min = r_min.min(pixel.r);
            g_min = g_min.min(pixel.g);
            b_min = b_min.min(pixel.b);
            r_max = r_max.max(pixel.r);
            g_max = g_max.max(pixel.g);
            b_max = b_max.max(pixel.b);
        }

        let count = self.pixels.len() as f32;
        ImageStats {
            width: self.width,
            height: self.height,
            pixel_count: self.pixels.len(),
            red_stats: ChannelStats {
                mean: r_sum / count,
                min: r_min,
                max: r_max,
            },
            green_stats: ChannelStats {
                mean: g_sum / count,
                min: g_min,
                max: g_max,
            },
            blue_stats: ChannelStats {
                mean: b_sum / count,
                min: b_min,
                max: b_max,
            },
        }
    }

    /// Convert image to PNG bytes for display
    pub fn to_png_bytes(&self) -> Result<Vec<u8>> {
        let rgb_image = self.to_rgb_image();
        let mut bytes = Vec::new();

        // Use a cursor to write PNG data to memory
        use std::io::Cursor;
        rgb_image
            .write_to(&mut Cursor::new(&mut bytes), image::ImageFormat::Png)
            .map_err(|e| std::io::Error::other(format!("Failed to encode PNG: {e}")))?;

        Ok(bytes)
    }
}

/// Statistics for a single color channel
#[derive(Debug, Clone)]
pub struct ChannelStats {
    pub mean: f32,
    pub min: f32,
    pub max: f32,
}

/// Complete image statistics
#[derive(Debug, Clone)]
pub struct ImageStats {
    pub width: usize,
    pub height: usize,
    pub pixel_count: usize,
    pub red_stats: ChannelStats,
    pub green_stats: ChannelStats,
    pub blue_stats: ChannelStats,
}

impl ImageStats {
    pub fn format_summary(&self) -> String {
        format!(
            "Image: {}×{} ({} pixels)\nChannels: R[{:.3}-{:.3}], G[{:.3}-{:.3}], B[{:.3}-{:.3}]",
            self.width,
            self.height,
            self.pixel_count,
            self.red_stats.min,
            self.red_stats.max,
            self.green_stats.min,
            self.green_stats.max,
            self.blue_stats.min,
            self.blue_stats.max
        )
    }
}

/// Preprocessing utilities for biomedical images
pub struct ImageProcessor;

impl ImageProcessor {
    /// Normalize image to [0, 1] range
    pub fn normalize(image: &mut ImageData) {
        let stats = image.statistics();
        let r_range = stats.red_stats.max - stats.red_stats.min;
        let g_range = stats.green_stats.max - stats.green_stats.min;
        let b_range = stats.blue_stats.max - stats.blue_stats.min;

        for pixel in &mut image.pixels {
            if r_range > 0.0 {
                pixel.r = (pixel.r - stats.red_stats.min) / r_range;
            }
            if g_range > 0.0 {
                pixel.g = (pixel.g - stats.green_stats.min) / g_range;
            }
            if b_range > 0.0 {
                pixel.b = (pixel.b - stats.blue_stats.min) / b_range;
            }
        }
    }

    /// Remove background by subtracting minimum value per channel
    pub fn remove_background(image: &mut ImageData) {
        let stats = image.statistics();

        for pixel in &mut image.pixels {
            pixel.r = (pixel.r - stats.red_stats.min).max(0.0);
            pixel.g = (pixel.g - stats.green_stats.min).max(0.0);
            pixel.b = (pixel.b - stats.blue_stats.min).max(0.0);
        }
    }

    /// Apply gamma correction for better visualization
    pub fn gamma_correction(image: &mut ImageData, gamma: f32) {
        for pixel in &mut image.pixels {
            pixel.r = pixel.r.powf(1.0 / gamma);
            pixel.g = pixel.g.powf(1.0 / gamma);
            pixel.b = pixel.b.powf(1.0 / gamma);
        }
    }
}

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

    #[test]
    fn test_color_conversion() {
        let color = Color::from_u8(128, 64, 192);
        let (r, g, b) = color.to_u8();
        assert_eq!((r, g, b), (128, 64, 192));
    }

    #[test]
    fn test_matrix_conversion() {
        let pixels = vec![
            Color::new(1.0, 0.5, 0.0),
            Color::new(0.0, 1.0, 0.5),
            Color::new(0.5, 0.0, 1.0),
            Color::new(0.25, 0.75, 0.25),
        ];

        let image = ImageData {
            width: 2,
            height: 2,
            pixels,
        };

        let (r, g, b) = image.to_rgb_matrices();
        assert_eq!(r.shape(), (2, 2));
        assert_eq!(g.shape(), (2, 2));
        assert_eq!(b.shape(), (2, 2));

        // Test round trip
        let reconstructed = ImageData::from_rgb_matrices(&r, &g, &b).unwrap();
        assert_eq!(reconstructed.width, 2);
        assert_eq!(reconstructed.height, 2);
        assert_eq!(reconstructed.pixels.len(), 4);
    }
}