use crate::{grid::Grid, palette::Palette};
use image::RgbImage;
use itertools::multizip;

/// Stores data that is located in grids that is used for image generation
#[derive(Clone)]
pub struct ThinGridData {
    width: usize,
    height: usize,
    data: Vec<f32>,
}

impl ThinGridData {
    /// Convert Grid to ThinGridData
    pub fn new_from_grid(in_grid: Grid) -> Self {
        ThinGridData {
            width: in_grid.width,
            height: in_grid.height,
            data: in_grid.data,
        }
    }

    pub fn new_from_grid_vec(in_grids: &[Grid]) -> Vec<Self> {
        in_grids.iter().cloned().map(Self::new_from_grid).collect()
    }

    /// from grid.rs (needed in image gen)
    pub fn quantile(&self, fraction: f32) -> f32 {
        let index = if (fraction - 1.0_f32).abs() < f32::EPSILON {
            self.data.len() - 1
        } else {
            (self.data.len() as f32 * fraction) as usize
        };
        let mut sorted = self.data.clone();
        sorted
            .as_mut_slice()
            .select_nth_unstable_by(index, |a, b| a.partial_cmp(b).unwrap());
        sorted[index]
    }
}

/// Class for storing data that will be used to create images
#[derive(Clone)]
pub struct ImgData {
    grids: Vec<ThinGridData>,
    palette: Palette,
}

impl ImgData {
    pub const fn new(in_grids: Vec<ThinGridData>, in_palette: Palette) -> Self {
        ImgData {
            grids: in_grids,
            palette: in_palette,
        }
    }

    pub fn to_image(&self) -> RgbImage {
        let (width, height) = (self.grids[0].width, self.grids[0].height);
        let mut img = image::RgbImage::new(width as u32, height as u32);

        let max_values: Vec<_> = self
            .grids
            .iter()
            .map(|grid| grid.quantile(0.999) * 1.5)
            .collect();

        for y in 0..height {
            for x in 0..width {
                let i = y * width + x;
                let (mut r, mut g, mut b) = (0.0_f32, 0.0_f32, 0.0_f32);
                for (grid, max_value, color) in
                    multizip((&self.grids, &max_values, &self.palette.colors))
                {
                    let mut t = (grid.data[i] / max_value).clamp(0.0, 1.0);
                    t = t.powf(1.0 / 2.2); // gamma correction
                    r += color.0[0] as f32 * t;
                    g += color.0[1] as f32 * t;
                    b += color.0[2] as f32 * t;
                }
                r = r.clamp(0.0, 255.0);
                g = g.clamp(0.0, 255.0);
                b = b.clamp(0.0, 255.0);
                img.put_pixel(x as u32, y as u32, image::Rgb([r as u8, g as u8, b as u8]));
            }
        }
        img
    }
}