use crate::{grid::Grid, palette::Palette};

use itertools::multizip;

// Stores data that is located in grids that is used for image generation
pub struct ThinGridData {
    pub width: usize,
    pub height: usize,
    pub data: Vec<f32>,
}


impl Clone for ThinGridData {
    fn clone(&self) -> ThinGridData {
        ThinGridData {
            width: self.width,
            height: self.height,
            data: self.data.clone(),
        }
    }
}

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.clone(),
        }
    }

    #[allow(dead_code)]
    pub fn new_from_grid_vec(in_grids: Vec<Grid>) -> Vec<Self> {
        in_grids.iter().map(|grid|{
            Self::new_from_grid(grid)
        }).collect()
    }

    // from grid.rs (needed in image gen)
    #[allow(dead_code)]
    pub fn data(&self) -> &[f32] {
        &self.data
    }

    // from grid.rs (needed in image gen)
    #[allow(dead_code)]
    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]
    }

    pub fn size_of(&self) -> usize {
        let mut output: usize = 0;
        output += std::mem::size_of_val(&self.width);
        output += std::mem::size_of_val(&self.height);
        for i in self.data.iter() {
            output += std::mem::size_of_val(&i);
        }
        output
    }
}

// Class for storing data that will be used to create images
pub struct ImgData {
    pub grids: Vec<ThinGridData>,
    pub palette: Palette,
    pub iteration: i32,
}

impl Clone for ImgData {
    fn clone(&self) -> ImgData {
        ImgData {
            grids: self.grids.clone(),
            palette: self.palette,
            iteration: self.iteration,
        }
    }
}

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

    pub fn size_of(&self) -> usize {
        let mut output: usize = 0;
        output += std::mem::size_of_val(&self.iteration);
        output += std::mem::size_of_val(&self.palette);
        for grid in self.grids.iter() {
            output += grid.size_of();
        }
        output
    }

    #[inline]
    pub fn save_to_image(&self) {
        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.save(format!("./tmp/out_{}.png", self.iteration).as_str())
            .unwrap();
    }
}