use shared::core::geometry::Point2i;
use shared::core::image::{Image, PixelFormat};
use shared::utils::math::f16_to_f32;
use std::ops::Deref;

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

pub use metadata::*;

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(Clone, Debug, Default)]
pub struct ImageChannelValues(pub SmallVec<[Float; 4]>);

impl ImageChannelValues {
    pub fn average(&self) -> Float {
        if self.0.is_empty() {
            return 0.0;
        }
        let sum: Float = self.0.iter().sum();
        sum / (self.0.len() as Float)
    }

    pub fn max_value(&self) -> Float {
        self.0.iter().fold(Float::MIN, |a, &b| a.max(b))
    }
}

impl From<&[Float]> for ImageChannelValues {
    fn from(slice: &[Float]) -> Self {
        Self(SmallVec::from_slice(slice))
    }
}

impl From<Vec<Float>> for ImageChannelValues {
    fn from(vec: Vec<Float>) -> Self {
        Self(SmallVec::from_vec(vec))
    }
}

impl<const N: usize> From<[Float; N]> for ImageChannelValues {
    fn from(arr: [Float; N]) -> Self {
        Self(SmallVec::from_slice(&arr))
    }
}

impl Deref for ImageChannelValues {
    type Target = SmallVec<[Float; 4]>;
    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl DerefMut for ImageChannelValues {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.0
    }
}

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

pub struct ImageBuffer {
    pub view: Image,
    pub channel_names: Vec<String>,
    _storage: PixelStorage,
}

impl Deref for ImageBuffer {
    type Target = Image;
    fn deref(&self) -> &Self::Target {
        &self.view
    }
}

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

impl ImageBuffer {
    fn resolution(&self) {
        self.view.resolution()
    }

    pub fn new_empty() -> Self {
        Self {
            channel_names: Vec::new(),
            _storage: PixelStorage::U8(Vec::new()),
            view: Image {
                format: PixelFormat::U256,
                resolution: Point2i::new(0, 0),
                n_channels: 0,
                pixels: Pixels::U8(std::ptr::null()),
                encoding: ColorEncoding::default(),
            },
        }
    }

    fn from_vector(
        format: PixelFormat,
        resolution: Point2i,
        channel_names: Vec<String>,
        encoding: ColorEncoding,
    ) -> Self {
        let n_channels = channel_names.len() as i32;
        let (format, pixels_view) = match &storage {
            PixelStorage::U8(vec) => (PixelFormat::U8, ImagePixels::U8(vec.as_ptr())),
            PixelStorage::F16(vec) => (PixelFormat::F16, ImagePixels::F16(vec.as_ptr())),
            PixelStorage::F32(vec) => (PixelFormat::F32, ImagePixels::F32(vec.as_ptr())),
        };

        let view = Image {
            format,
            resolution,
            n_channels,
            encoding,
            pixels: pixels_view,
        };

        Self {
            view,
            _storage: storage,
            channel_names,
        }
    }

    pub fn from_u8(
        data: Vec<u8>,
        resolution: Point2i,
        channel_names: Vec<String>,
        encoding: ColorEncoding,
    ) -> Self {
        let n_channels = channel_names.len() as i32;
        let expected_len = (resolution.x * resolution.y) as usize * n_channels as usize;
        if data.len() != expected_len {
            panic!(
                "ImageBuffer::from_u8: Data length {} does not match resolution {:?} * channels {}",
                data.len(),
                resolution,
                n_channels
            );
        }

        let storage = PixelStorage::U8(data);

        let ptr = match &storage {
            PixelStorage::U8(v) => v.as_ptr(),
            _ => unreachable!(),
        };

        let view = Image {
            format: PixelFormat::U256,
            resolution,
            n_channels,
            pixels: Pixels::U8(ptr),
            encoding: encoding.clone(),
        };

        Self {
            view,
            channel_names,
            _storage: storage,
        }
    }

    pub fn from_u8(
        data: Vec<u8>,
        resolution: Point2i,
        channel_names: Vec<String>,
        encoding: ColorEncoding,
    ) -> Self {
        let n_channels = channel_names.len() as i32;
        let expected_len = (resolution.x * resolution.y) as usize * n_channels as usize;

        if data.len() != expected_len {
            panic!(
                "ImageBuffer::from_u8: Data length {} does not match resolution {:?} * channels {}",
                data.len(),
                resolution,
                n_channels
            );
        }

        let storage = PixelStorage::U8(data);

        let ptr = match &storage {
            PixelStorage::U8(v) => v.as_ptr(),
            _ => unreachable!(),
        };

        let view = Image {
            format: PixelFormat::U256,
            resolution,
            n_channels,
            pixels: Pixels::U8(ptr),
            encoding: encoding.clone(),
        };

        Self {
            view,
            channel_names,
            _storage: storage,
        }
    }

    pub fn from_f16(data: Vec<half::f16>, resolution: Point2i, channel_names: Vec<String>) -> Self {
        let n_channels = channel_names.len() as i32;
        let expected_len = (resolution.x * resolution.y) as usize * n_channels as usize;

        if data.len() != expected_len {
            panic!("ImageBuffer::from_f16: Data length mismatch");
        }

        let storage = PixelStorage::F16(data);

        let ptr = match &storage {
            PixelStorage::F16(v) => v.as_ptr() as *const u16,
            _ => unreachable!(),
        };

        let view = Image {
            format: PixelFormat::Half,
            resolution,
            n_channels,
            pixels: Pixels::F16(ptr),
            encoding: ColorEncoding::default(),
        };

        Self {
            view,
            channel_names,
            _storage: storage,
        }
    }

    pub fn from_f32(data: Vec<f32>, resolution: Point2i, channel_names: Vec<String>) -> Self {
        let n_channels = channel_names.len() as i32;
        let expected_len = (resolution.x * resolution.y) as usize * n_channels as usize;

        if data.len() != expected_len {
            panic!("ImageBuffer::from_f32: Data length mismatch");
        }

        let storage = PixelStorage::F32(data);

        let ptr = match &storage {
            PixelStorage::F32(v) => v.as_ptr(),
            _ => unreachable!(),
        };

        let view = Image {
            format: PixelFormat::Float,
            resolution,
            n_channels,
            pixels: Pixels::F32(ptr),
            encoding: ColorEncoding::default(),
        };

        Self {
            view,
            channel_names,
            _storage: storage,
        }
    }

    pub fn new(
        format: PixelFormat,
        resolution: Point2i,
        channel_names: &[&str],
        encoding: *const ColorEncoding,
    ) -> Self {
        let n_channels = channel_names.len();
        let pixel_count = (resolution.x * resolution.y) as usize * n_channels;
        let owned_names: Vec<String> = channel_names.iter().map(|s| s.to_string()).collect();

        let storage = match format {
            PixelFormat::U8 => PixelStorage::U8(vec![0; pixel_count]),
            PixelFormat::F16 => PixelStorage::F16(vec![0; pixel_count]),
            PixelFormat::F32 => PixelStorage::F32(vec![0.0; pixel_count]),
        };

        Self::from_vector(storage, resolution, owned_names, encoding)
    }

    pub fn channel_names(&self) -> Vec<&str> {
        self.channel_names.iter().map(|s| s.as_str()).collect()
    }

    pub fn encoding(&self) -> ColorEncoding {
        self.view.encoding
    }

    pub fn set_channel(&self, p: Point2i, c: i32, mut value: Float) {
        if value.is_nan() {
            value = 0.0;
        }

        if !self.view.resolution.inside_exclusive(p) {
            return;
        }

        let offset = self.view.pixel_offset(p) + c as usize;
        match &mut self._storage {
            PixelStorage::U8(data) => {
                if !self.view.encoding.is_null() {
                    data[offset] = self.view.encoding.from_linear_scalar(value);
                } else {
                    let val = (value * 255.0 + 0.5).clamp(0.0, 255.0);
                    data[offset] = val as u8;
                }
            }
            PixelStorage::F16(data) => {
                data[offset] = f16_to_f32(value);
            }
            PixelStorage::F32(data) => {
                data[offset] = value;
            }
        }
    }

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

        let pixel_offset = self.view.pixel_offset(pp);
        let mut values: SmallVec<[Float; 4]> = SmallVec::with_capacity(desc.offset.len());
        match &self.pixels {
            PixelData::U8(data) => {
                for i in 0..self.view.n_channels() {
                    let raw_val = data[pixel_offset + i];
                    let linear = self.view.encoding.to_linear_scalar(raw_val);
                    values.push(linear);
                }
            }
            PixelData::F16(data) => {
                for i in 0..self.view.n_channels() {
                    let raw_val = data[pixel_offset];
                    values.push(f16_to_f32(raw_val));
                }
            }
            PixelData::F32(data) => {
                for i in 0..self.view.n_channels() {
                    let val = data[pixel_offset + i];
                    values.push(val);
                }
            }
        }

        ImageChannelValues(values)
    }

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

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

        let pixel_offset = self.view.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 raw_val = data[pixel_offset + channel_idx as usize];
                    let linear = self.view.encoding.to_linear_scalar(raw_val);
                    values.push(linear);
                }
            }
            PixelData::F16(data) => {
                for &channel_idx in &desc.offset {
                    let raw_val = data[pixel_offset + channel_idx as usize];
                    values.push(f16_to_f32(raw_val));
                }
            }
            PixelData::F32(data) => {
                for &channel_idx in &desc.offset {
                    let val = data[pixel_offset + channel_idx as usize];
                    values.push(val);
                }
            }
        }

        ImageChannelValues(values)
    }

    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 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 all_channels_desc(&self) -> ImageChannelDesc {
        ImageChannelDesc {
            offset: (0..self.n_channels()).collect(),
        }
    }

    pub fn select_channels(&self, desc: &ImageChannelDesc) -> Self {
        let desc_channel_names: Vec<String> = desc
            .offset
            .iter()
            .map(|&i| self.channel_names[i as usize])
            .collect();

        let new_storage = match &self._storage {
            PixelStorage::U8(src_data) => {
                // Allocate destination buffer
                let mut dst_data = vec![0u8; pixel_count * dst_n_channels];

                // Iterate over every pixel (Flat loop is faster than nested x,y)
                for i in 0..pixel_count {
                    let src_pixel_start = i * src_n_channels;
                    let dst_pixel_start = i * dst_n_channels;

                    // Copy specific channels based on desc
                    for (out_idx, &in_channel_offset) in desc.offset.iter().enumerate() {
                        let val = src_data[src_pixel_start + in_channel_offset as usize];
                        dst_data[dst_pixel_start + out_idx] = val;
                    }
                }
                PixelStorage::U8(dst_data)
            }
            PixelStorage::F16(src_data) => {
                let mut dst_data = vec![half::f16::ZERO; pixel_count * dst_n_channels];

                for i in 0..pixel_count {
                    let src_pixel_start = i * src_n_channels;
                    let dst_pixel_start = i * dst_n_channels;

                    for (out_idx, &in_channel_offset) in desc.offset.iter().enumerate() {
                        let val = src_data[src_pixel_start + in_channel_offset as usize];
                        dst_data[dst_pixel_start + out_idx] = val;
                    }
                }
                PixelStorage::F16(dst_data)
            }
            PixelStorage::F32(src_data) => {
                let mut dst_data = vec![0.0; pixel_count * dst_n_channels];

                for i in 0..pixel_count {
                    let src_pixel_start = i * src_n_channels;
                    let dst_pixel_start = i * dst_n_channels;

                    for (out_idx, &in_channel_offset) in desc.offset.iter().enumerate() {
                        let val = src_data[src_pixel_start + in_channel_offset as usize];
                        dst_data[dst_pixel_start + out_idx] = val;
                    }
                }
                PixelStorage::F32(dst_data)
            }
        };

        let image = Self::new(
            self.format,
            self.resolution,
            desc_channel_names,
            self.encoding(),
        );
    }

    pub fn set_channels(
        &mut self,
        p: Point2i,
        desc: &ImageChannelDesc,
        mut values: &ImageChannelValues,
    ) {
        assert_eq!(desc.size(), values.len());
        for i in 0..desc.size() {
            self.view.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)
    }

    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_channel_with_desc(Point2i::new(x, y), &desc, WrapMode::Clamp.into());
                let v_ref = self.get_channel_with_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)
    }

    pub fn update_view_pointers(&mut self) {
        self.view.pixels = match &self._storage {
            PixelStorage::U8(vec) => Pixels::U8(vec.as_ptr()),
            PixelStorage::F16(vec) => Pixels::F16(vec.as_ptr() as *const u16),
            PixelStorage::F32(vec) => Pixels::F32(vec.as_ptr()),
        };
    }
}