pub mod io;
pub mod metadata;
pub mod ops;
pub mod pixel;

use crate::core::pbrt::{Float, lerp};
use crate::geometry::{Bounds2f, Point2f, Point2fi, Point2i};
use crate::utils::color::{ColorEncoding, ColorEncodingTrait, LINEAR};
use crate::utils::containers::Array2D;
use crate::utils::math::square;
use core::hash;
use half::f16;
use pixel::PixelStorage;
use rayon::prelude::*;
use smallvec::{SmallVec, smallvec};
use std::ops::{Deref, DerefMut};

pub use metadata::{ImageChannelDesc, ImageChannelValues, ImageMetadata, WrapMode, WrapMode2D};

#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum PixelFormat {
    U8,
    F16,
    F32,
}

impl PixelFormat {
    pub fn is_8bit(&self) -> bool {
        matches!(self, PixelFormat::U8)
    }

    pub fn is_16bit(&self) -> bool {
        matches!(self, PixelFormat::F16)
    }

    pub fn is_32bit(&self) -> bool {
        matches!(self, PixelFormat::F32)
    }

    pub fn texel_bytes(&self) -> usize {
        match self {
            PixelFormat::U8 => 1,
            PixelFormat::F16 => 2,
            PixelFormat::F32 => 4,
        }
    }
}

impl std::fmt::Display for PixelFormat {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            PixelFormat::U8 => write!(f, "U256"),
            PixelFormat::F16 => write!(f, "Half"),
            PixelFormat::F32 => write!(f, "Float"),
        }
    }
}

#[derive(Debug, Clone)]
pub enum PixelData {
    U8(Vec<u8>),
    F16(Vec<f16>),
    F32(Vec<f32>),
}

#[derive(Debug, Clone)]
pub struct Image {
    pub format: PixelFormat,
    pub resolution: Point2i,
    pub channel_names: Vec<String>,
    pub encoding: ColorEncoding,
    pub pixels: PixelData,
}

#[derive(Debug)]
pub struct ImageAndMetadata {
    pub image: Image,
    pub metadata: ImageMetadata,
}

impl Image {
    fn from_vector(
        format: PixelFormat,
        resolution: Point2i,
        channel_names: Vec<String>,
        encoding: ColorEncoding,
    ) -> Self {
        let size = (resolution.x() * resolution.y()) as usize * channel_names.len();

        let pixels = match format {
            PixelFormat::U8 => PixelData::U8(vec![0; size]),
            PixelFormat::F16 => PixelData::F16(vec![f16::ZERO; size]),
            PixelFormat::F32 => PixelData::F32(vec![0.0; size]),
        };

        Self {
            format,
            resolution,
            channel_names,
            encoding,
            pixels,
        }
    }

    pub fn new(
        format: PixelFormat,
        resolution: Point2i,
        channel_names: &[&str],
        encoding: ColorEncoding,
    ) -> Self {
        let owned_names = channel_names.iter().map(|s| s.to_string()).collect();
        Self::from_vector(format, resolution, owned_names, encoding)
    }

    pub fn format(&self) -> PixelFormat {
        self.format
    }
    pub fn resolution(&self) -> Point2i {
        self.resolution
    }
    pub fn n_channels(&self) -> usize {
        self.channel_names.len()
    }
    pub fn channel_names(&self) -> Vec<&str> {
        self.channel_names.iter().map(|s| s.as_str()).collect()
    }
    pub fn channel_names_from_desc(&self, desc: &ImageChannelDesc) -> Vec<&str> {
        desc.offset
            .iter()
            .map(|&i| self.channel_names[i].as_str())
            .collect()
    }
    pub fn encoding(&self) -> ColorEncoding {
        self.encoding
    }

    pub fn pixel_offset(&self, p: Point2i) -> usize {
        (p.y() as usize * self.resolution.x() as usize + p.x() as usize) * self.n_channels()
    }

    pub fn get_channel(&self, p: Point2i, c: usize) -> Float {
        self.get_channel_with_wrap(p, c, WrapMode::Clamp.into())
    }

    pub fn get_channel_with_wrap(&self, p: Point2i, c: usize, wrap: WrapMode2D) -> Float {
        let mut pp = p;
        if !self.remap_pixel_coords(&mut pp, wrap) {
            return 0.0;
        }

        let idx = self.pixel_offset(pp) + c;
        match &self.pixels {
            PixelData::U8(d) => u8::to_linear(d[idx], self.encoding),
            PixelData::F16(d) => f16::to_linear(d[idx], self.encoding),
            PixelData::F32(d) => f32::to_linear(d[idx], self.encoding),
        }
    }

    pub fn get_channels(&self, p: Point2i, wrap: WrapMode2D) -> ImageChannelValues {
        let mut pp = p;

        if !self.remap_pixel_coords(&mut pp, wrap) {
            return ImageChannelValues(smallvec![0.0; self.n_channels()]);
        }

        let start_idx = self.pixel_offset(pp);
        let n_channels = self.n_channels();

        let mut values: SmallVec<[Float; 4]> = SmallVec::with_capacity(n_channels);

        match &self.pixels {
            PixelData::U8(data) => {
                let slice = &data[start_idx..start_idx + n_channels];
                for &v in slice {
                    values.push(u8::to_linear(v, self.encoding));
                }
            }
            PixelData::F16(data) => {
                let slice = &data[start_idx..start_idx + n_channels];
                for &v in slice {
                    values.push(f16::to_linear(v, self.encoding));
                }
            }
            PixelData::F32(data) => {
                let slice = &data[start_idx..start_idx + n_channels];
                for &v in slice {
                    values.push(f32::to_linear(v, self.encoding));
                }
            }
        }

        ImageChannelValues(values)
    }

    pub fn get_channels_desc(
        &self,
        p: Point2i,
        desc: &ImageChannelDesc,
        wrap: WrapMode2D,
    ) -> ImageChannelValues {
        let mut pp = p;
        if !self.remap_pixel_coords(&mut pp, wrap) {
            return ImageChannelValues(smallvec![0.0; desc.offset.len()]);
        }

        let pixel_offset = self.pixel_offset(pp);

        let mut values: SmallVec<[Float; 4]> = SmallVec::with_capacity(desc.offset.len());

        match &self.pixels {
            PixelData::U8(data) => {
                for &channel_idx in &desc.offset {
                    let val = data[pixel_offset + channel_idx];
                    values.push(u8::to_linear(val, self.encoding));
                }
            }
            PixelData::F16(data) => {
                for &channel_idx in &desc.offset {
                    let val = data[pixel_offset + channel_idx];
                    values.push(f16::to_linear(val, self.encoding));
                }
            }
            PixelData::F32(data) => {
                for &channel_idx in &desc.offset {
                    let val = data[pixel_offset + channel_idx];
                    values.push(f32::to_linear(val, self.encoding));
                }
            }
        }

        ImageChannelValues(values)
    }

    pub fn get_channels_default(&self, p: Point2i) -> ImageChannelValues {
        self.get_channels(p, WrapMode::Clamp.into())
    }

    pub fn all_channels_desc(&self) -> ImageChannelDesc {
        ImageChannelDesc {
            offset: (0..self.n_channels()).collect(),
        }
    }

    pub fn get_channel_desc(
        &self,
        requested_channels: &[&str],
    ) -> Result<ImageChannelDesc, String> {
        let mut offset = Vec::with_capacity(requested_channels.len());

        for &req in requested_channels.iter() {
            match self.channel_names.iter().position(|n| n == req) {
                Some(idx) => {
                    offset.push(idx);
                }
                None => {
                    return Err(format!(
                        "Image is missing requested channel '{}'. Available channels: {:?}",
                        req, self.channel_names
                    ));
                }
            }
        }

        Ok(ImageChannelDesc { offset })
    }

    pub fn set_channel(&mut self, p: Point2i, c: usize, value: Float) {
        let val_no_nan = if value.is_nan() { 0.0 } else { value };
        let offset = self.pixel_offset(p) + c;
        match &mut self.pixels {
            PixelData::U8(data) => {
                let linear = [val_no_nan];
                self.encoding
                    .from_linear_slice(&linear, &mut data[offset..offset + 1]);
            }
            PixelData::F16(data) => data[offset] = f16::from_f32(val_no_nan),
            PixelData::F32(data) => data[offset] = val_no_nan,
        }
    }

    pub fn set_channels(
        &mut self,
        p: Point2i,
        desc: &ImageChannelDesc,
        values: &ImageChannelValues,
    ) {
        assert_eq!(desc.size(), values.len());
        for i in 0..desc.size() {
            self.set_channel(p, desc.offset[i], values[i]);
        }
    }

    pub fn set_channels_all(&mut self, p: Point2i, values: &ImageChannelValues) {
        self.set_channels(p, &self.all_channels_desc(), values)
    }

    fn remap_pixel_coords(&self, p: &mut Point2i, wrap_mode: WrapMode2D) -> bool {
        for i in 0..2 {
            if p[i] >= 0 && p[i] < self.resolution[i] {
                continue;
            }
            match wrap_mode.uv[i] {
                WrapMode::Black => return false,
                WrapMode::Clamp => p[i] = p[i].clamp(0, self.resolution[i] - 1),
                WrapMode::Repeat => p[i] = p[i].rem_euclid(self.resolution[i]),
                WrapMode::OctahedralSphere => {
                    p[i] = p[i].clamp(0, self.resolution[i] - 1);
                }
            }
        }
        true
    }

    pub fn bilerp_channel(&self, p: Point2f, c: usize) -> Float {
        self.bilerp_channel_with_wrap(p, c, WrapMode::Clamp.into())
    }

    pub fn bilerp_channel_with_wrap(&self, p: Point2f, c: usize, wrap_mode: WrapMode2D) -> Float {
        let x = p.x() * self.resolution.x() as Float - 0.5;
        let y = p.y() * self.resolution.y() as Float - 0.5;
        let xi = x.floor() as i32;
        let yi = y.floor() as i32;
        let dx = x - xi as Float;
        let dy = y - yi as Float;
        let v00 = self.get_channel_with_wrap(Point2i::new(xi, yi), c, wrap_mode);
        let v10 = self.get_channel_with_wrap(Point2i::new(xi + 1, yi), c, wrap_mode);
        let v01 = self.get_channel_with_wrap(Point2i::new(xi, yi + 1), c, wrap_mode);
        let v11 = self.get_channel_with_wrap(Point2i::new(xi + 1, yi + 1), c, wrap_mode);
        lerp(dy, lerp(dx, v00, v10), lerp(dx, v01, v11))
    }

    pub fn lookup_nearest_channel_with_wrap(
        &self,
        p: Point2f,
        c: usize,
        wrap_mode: WrapMode2D,
    ) -> Float {
        let pi = Point2i::new(
            p.x() as i32 * self.resolution.x(),
            p.y() as i32 * self.resolution.y(),
        );

        self.get_channel_with_wrap(pi, c, wrap_mode)
    }

    pub fn lookup_nearest_channel(&self, p: Point2f, c: usize) -> Float {
        self.lookup_nearest_channel_with_wrap(p, c, WrapMode::Clamp.into())
    }

    pub fn get_sampling_distribution<F>(&self, dxd_a: F, domain: Bounds2f) -> Array2D<Float>
    where
        F: Fn(Point2f) -> Float + Sync + Send,
    {
        let width = self.resolution.x();
        let height = self.resolution.y();

        let mut dist = Array2D::new_with_dims(width as usize, height as usize);

        dist.values
            .par_chunks_mut(width as usize)
            .enumerate()
            .for_each(|(y, row)| {
                let y = y as i32;

                for (x, out_val) in row.iter_mut().enumerate() {
                    let x = x as i32;

                    let value = self.get_channels_default(Point2i::new(x, y)).average();

                    let u = (x as Float + 0.5) / width as Float;
                    let v = (y as Float + 0.5) / height as Float;
                    let p = domain.lerp(Point2f::new(u, v));
                    *out_val = value * dxd_a(p);
                }
            });

        dist
    }

    pub fn get_sampling_distribution_uniform(&self) -> Array2D<Float> {
        let default_domain = Bounds2f::from_points(Point2f::new(0.0, 0.0), Point2f::new(1.0, 1.0));

        self.get_sampling_distribution(|_| 1.0, default_domain)
    }

    pub fn mse(
        &self,
        desc: ImageChannelDesc,
        ref_img: &Image,
        generate_mse_image: bool,
    ) -> (ImageChannelValues, Option<Image>) {
        let mut sum_se: Vec<f64> = vec![0.; desc.size()];
        let names_ref = self.channel_names_from_desc(&desc);
        let ref_desc = ref_img
            .get_channel_desc(&self.channel_names_from_desc(&desc))
            .expect("Channels not found in image");
        assert_eq!(self.resolution(), ref_img.resolution());

        let width = self.resolution.x() as usize;
        let height = self.resolution.y() as usize;
        let n_channels = desc.offset.len();
        let mut mse_pixels = if generate_mse_image {
            vec![0.0f32; width * height * n_channels]
        } else {
            Vec::new()
        };

        for y in 0..self.resolution().y() {
            for x in 0..self.resolution().x() {
                let v = self.get_channels_desc(Point2i::new(x, y), &desc, WrapMode::Clamp.into());
                let v_ref =
                    self.get_channels_desc(Point2i::new(x, y), &ref_desc, WrapMode::Clamp.into());
                for c in 0..desc.size() {
                    let se = square(v[c] as f64 - v_ref[c] as f64);
                    if se.is_infinite() {
                        continue;
                    }
                    sum_se[c] += se;
                    if generate_mse_image {
                        let idx = (y as usize * width + x as usize) * n_channels + c;
                        mse_pixels[idx] = se as f32;
                    }
                }
            }
        }

        let pixel_count = (self.resolution.x() * self.resolution.y()) as f64;
        let mse_values: SmallVec<[Float; 4]> =
            sum_se.iter().map(|&s| (s / pixel_count) as Float).collect();

        let mse_image = if generate_mse_image {
            Some(Image::new(
                PixelFormat::F32,
                self.resolution,
                &names_ref,
                LINEAR,
            ))
        } else {
            None
        };

        (ImageChannelValues(mse_values), mse_image)
    }
}