crystal_ball 0.3.0

use std::convert::TryFrom;
use std::fs;
use std::fs::File;
use std::io::BufReader;
use std::path::Path;

use gltf::image::Format;
use image::codecs::hdr::HdrDecoder;
use image::{ImageBuffer, ImageFormat, Rgb};

use crate::color::{Color, Texture};
use crate::util::Error;

/// An image of fixed size represented as a row-major [`Vec`] of [`Color`].
/// In addition, this also stores the [`Interpolation`] used for accessing colors between pixels.
///
/// (0, 0) is considered to be the upper left corner.
#[derive(Clone, Debug)]
pub struct Image {
    pub width: u32,
    pub height: u32,
    pub pixels: Vec<Color>,
    pub interpolation: Interpolation,
}

impl Image {
    /// Create a new (black) [`Image`].
    pub fn new(width: u32, height: u32, interpolation: Interpolation) -> Self {
        Image {
            width,
            height,
            pixels: vec![Color::default(); (width * height) as usize],
            interpolation,
        }
    }

    /// Create a new [`Image`] from an already existing [`Vec`] of [`Color`].
    ///
    /// # Panics
    /// Panics if the dimensions don't match the size of `pixels`.
    pub fn from_pixels(
        width: u32,
        height: u32,
        pixels: Vec<Color>,
        interpolation: Interpolation,
    ) -> Self {
        assert_eq!(
            width * height,
            pixels.len() as u32,
            "image dimensions don't match the data size"
        );

        Image {
            width,
            height,
            pixels,
            interpolation,
        }
    }

    /// Create a new [`Image`] from a byte buffer and a given [`ColorFormat`].
    ///
    /// # Panics
    /// Panics if the dimensions don't match the size of `data` (respecting the [`ColorFormat`]).
    pub fn from_raw_with_format(
        width: u32,
        height: u32,
        data: &[u8],
        format: ColorFormat,
        interpolation: Interpolation,
    ) -> Self {
        let offset = format.offset();

        assert_eq!(
            width * height * offset as u32,
            data.len() as u32,
            "image dimensions don't match the data size"
        );

        let pixels = (0..height * width)
            .map(|i| {
                let x = (i % width) as usize;
                let y = (i / width) as usize;

                let start = offset * (x + y * width as usize);

                match format {
                    ColorFormat::Rgb8 | ColorFormat::Rgba8 => Color::new(
                        data[start] as f32 / 255.0,
                        data[start + 1] as f32 / 255.0,
                        data[start + 2] as f32 / 255.0,
                    ),
                    ColorFormat::Rgb16 | ColorFormat::Rgba16 => Color::new(
                        (((data[start + 1] as u16) << 8) | data[start] as u16) as f32 / 65535.0,
                        (((data[start + 3] as u16) << 8) | data[start + 2] as u16) as f32 / 65535.0,
                        (((data[start + 5] as u16) << 8) | data[start + 4] as u16) as f32 / 65535.0,
                    ),
                }
            })
            .collect::<Vec<Color>>();

        Image {
            width,
            height,
            pixels,
            interpolation,
        }
    }

    /// Create a new [`Image`] from a byte buffer.
    ///
    /// This assumes the [`ColorFormat`] to be [`ColorFormat::Rgb8`].
    ///
    /// # Panics
    /// Panics if the dimensions don't match the size of `data`.
    pub fn from_raw(width: u32, height: u32, data: &[u8], interpolation: Interpolation) -> Self {
        Self::from_raw_with_format(width, height, data, ColorFormat::Rgb8, interpolation)
    }

    /// Load an image from the specified path.
    ///
    /// This automatically detects whether an image is HDR and performs the necessary conversions.
    ///
    /// Returns an [`Error`] if the image can not be loaded.
    pub fn from_file<P: AsRef<Path>>(path: P, interpolation: Interpolation) -> Result<Self, Error> {
        match ImageFormat::from_path(&path)? {
            ImageFormat::Hdr => Self::from_hdr_file(path, interpolation),
            _ => Self::from_image_file(path, interpolation),
        }
    }

    /// Load an SDR image from the specified path.
    ///
    /// Returns an [`Error`] if the image can not be loaded.
    fn from_image_file<P: AsRef<Path>>(
        path: P,
        interpolation: Interpolation,
    ) -> Result<Self, Error> {
        let image = image::open(path)?.into_rgb32f();
        let (width, height) = image.dimensions();

        let mut pixels = vec![Color::default(); (width * height) as usize];
        pixels.iter_mut().enumerate().for_each(|(i, color)| {
            let x = i % width as usize;
            let y = i / width as usize;

            let pixel = image.get_pixel(x as u32, y as u32);
            *color = Color::from(pixel.0);
        });

        Ok(Image {
            width: image.width(),
            height: image.height(),
            pixels,
            interpolation,
        })
    }

    /// Load an HDR image from the specified path.
    ///
    /// Returns an [`Error`] if the image can not be loaded.
    fn from_hdr_file<P: AsRef<Path>>(path: P, interpolation: Interpolation) -> Result<Self, Error> {
        let decoder = HdrDecoder::new(BufReader::new(File::open(path)?))?;

        let width = decoder.metadata().width;
        let height = decoder.metadata().height;

        let image = decoder.read_image_hdr()?;

        let mut pixels = vec![Color::default(); (width * height) as usize];
        pixels.iter_mut().enumerate().for_each(|(i, color)| {
            let x = i % width as usize;
            let y = i / width as usize;

            let pixel = image[x + y * width as usize];

            *color = Color::new(pixel.0[0], pixel.0[1], pixel.0[2]).srgb_to_linear();
        });

        Ok(Image {
            width,
            height,
            pixels,
            interpolation,
        })
    }

    /// Create an image from a [`Vec`] of [`ImageTile`].
    pub(crate) fn from_tiles(
        width: u32,
        height: u32,
        tiles: &[ImageTile],
        interpolation: Interpolation,
    ) -> Self {
        let mut pixels = vec![Color::default(); (width * height) as usize];

        for tile in tiles {
            for y in 0..tile.image.height {
                let start = (tile.x + (tile.y + y) * width) as usize;
                let end = start + tile.image.width as usize;

                let tiles_start = (y * tile.image.width) as usize;
                let tiles_end = tiles_start + tile.image.width as usize;

                pixels[start..end].copy_from_slice(&tile.image.pixels[tiles_start..tiles_end])
            }
        }

        Image {
            width,
            height,
            pixels,
            interpolation,
        }
    }

    /// Return the dimensions of the image.
    pub fn dimensions(&self) -> (u32, u32) {
        (self.width, self.height)
    }

    /// Return the color value of the pixel `(x, y)`.
    ///
    /// The image is considered to wrap periodically and is thus virtually infinite.
    pub fn get_pixel(&self, x: f32, y: f32) -> Color {
        match self.interpolation {
            Interpolation::Closest => self.interpolate_closest(x, y),
            Interpolation::Bilinear => self.interpolate_bilinear(x, y),
        }
    }

    /// Return the closest color value at the position `(x, y)`.
    fn interpolate_closest(&self, x: f32, y: f32) -> Color {
        self.get_pixel_periodic(x.round(), y.round())
    }

    /// Return the bilinearly interpolated color value at the position `(x, y)`.
    fn interpolate_bilinear(&self, x: f32, y: f32) -> Color {
        let x1 = x.floor();
        let x2 = x.ceil();
        let y1 = y.floor();
        let y2 = y.ceil();

        let (weight1, weight2) = if x1 == x2 {
            (1.0, 0.0)
        } else {
            ((x2 - x) / (x2 - x1), (x - x1) / (x2 - x1))
        };

        let color1 =
            self.get_pixel_periodic(x1, y1) * weight1 + self.get_pixel_periodic(x2, y1) * weight2;
        let color2 =
            self.get_pixel_periodic(x1, y2) * weight1 + self.get_pixel_periodic(x2, y2) * weight2;

        if y1 == y2 {
            color1
        } else {
            (color1 * (y2 - y) + color2 * (y - y1)) / (y2 - y1)
        }
    }

    /// Return the wrapped color value at the position `(x, y)`.
    fn get_pixel_periodic(&self, x: f32, y: f32) -> Color {
        self.pixels[x.rem_euclid(self.width as f32) as usize
            + self.width as usize * y.rem_euclid(self.height as f32) as usize]
    }

    /// Set the color value of the pixel `(x, y)`.
    ///
    /// # Panics
    /// Panics if `(x, y)` ist out of bounds.
    pub fn set_pixel(&mut self, x: u32, y: u32, color: Color) {
        assert!(y < self.height);
        assert!(x < self.width);

        self.pixels[(x + y * self.width) as usize] = color;
    }

    /// Write the image to the provided path.
    ///
    /// The file format is inferred from the path.
    ///
    /// Returns an [`Error`] if the document can not be written.
    pub fn write<P: AsRef<Path>>(&self, path: P) -> Result<(), Error> {
        let mut image_buffer = ImageBuffer::new(self.width, self.height);
        image_buffer
            .enumerate_pixels_mut()
            .for_each(|(x, y, pixel)| {
                let color = self.pixels[(x + self.width * y) as usize];
                *pixel = Rgb([
                    (color.r * 255.0).round().min(255.0).max(0.0) as u8,
                    (color.g * 255.0).round().min(255.0).max(0.0) as u8,
                    (color.b * 255.0).round().min(255.0).max(0.0) as u8,
                ])
            });

        image_buffer.save(&path)?;

        println!("Saved image to {:?}", fs::canonicalize(path)?);

        Ok(())
    }

    /// Denoise the image using [Open Image Denoise](https://www.openimagedenoise.org/).
    ///
    /// Returns an [`Error`] if the image can not be denoised.
    #[cfg_attr(doc, doc(cfg(feature = "oidn")))]
    #[cfg(feature = "oidn")]
    pub fn denoise(&mut self) -> Result<(), Error> {
        use oidn::{Device, RayTracing};
        use std::time::Instant;

        let width = self.width;
        let height = self.height;

        let mut input_img = vec![0.0f32; (3 * width * height) as usize];

        for y in 0..height {
            for x in 0..width {
                let pixel = self.pixels[(x + y * width) as usize];
                for c in 0..3 {
                    input_img[3 * (y * width + x) as usize + c] = match c {
                        0 => pixel.r,
                        1 => pixel.g,
                        2 => pixel.b,
                        _ => 0.0,
                    };
                }
            }
        }

        let mut filter_output = vec![0.0f32; input_img.len()];

        println!("Denoising...");

        let start_time = Instant::now();

        let device = Device::new();
        let mut filter = RayTracing::new(&device);
        filter
            .srgb(true)
            .image_dimensions(width as usize, height as usize);
        filter.filter(&input_img[..], &mut filter_output[..])?;

        if let Err((e, _)) = device.get_error() {
            return Err(e)?;
        }

        println!(
            "Denoising finished after {:.2?}",
            Instant::now() - start_time
        );

        for i in (0..filter_output.len()).step_by(3) {
            let pixel = i as u32 / 3;

            let x = pixel % width;
            let y = pixel / width;

            let color = Color::new(filter_output[i], filter_output[i + 1], filter_output[i + 2]);

            // TODO: Just use an iterator and collect values
            self.set_pixel(x, y, color)
        }

        Ok(())
    }
}

impl Texture for Image {
    fn get_pixel(&self, u: f64, v: f64) -> Color {
        self.get_pixel(
            (self.width - 1) as f32 * u as f32,
            (self.height - 1) as f32 * v as f32,
        )
    }
}

/// A tile of an [`Image`] storing its position and its chunk of the original image.
#[derive(Clone, Debug)]
pub(crate) struct ImageTile {
    pub x: u32,
    pub y: u32,
    pub image: Image,
}

impl ImageTile {
    /// Create a new [`ImageTile`].
    pub fn new(x: u32, y: u32, image: Image) -> Self {
        ImageTile { x, y, image }
    }
}

/// The interpolation used to access pixels of an [`Image`].
#[derive(Copy, Clone, Debug)]
pub enum Interpolation {
    /// Chooses the [`Color`] of the closest pixel.
    Closest,
    /// Interpolates linearly between the [`Color`]s of the 4 surrounding pixels.
    Bilinear,
}

/// The color format for a given [`Image`].
///
/// This is used for conversion when loading a byte buffer.
#[derive(Copy, Clone, Debug)]
pub enum ColorFormat {
    Rgb8,
    Rgba8,
    Rgb16,
    Rgba16,
}

impl ColorFormat {
    /// Return the color format's number of bytes per pixel.
    pub fn offset(&self) -> usize {
        match self {
            ColorFormat::Rgb8 => 3,
            ColorFormat::Rgba8 => 4,
            ColorFormat::Rgb16 => 6,
            ColorFormat::Rgba16 => 8,
        }
    }
}

impl TryFrom<Format> for ColorFormat {
    type Error = Error;

    fn try_from(format: Format) -> Result<Self, Self::Error> {
        match format {
            Format::R8G8B8 => Ok(ColorFormat::Rgb8),
            Format::R8G8B8A8 => Ok(ColorFormat::Rgba8),
            Format::R16G16B16 => Ok(ColorFormat::Rgb16),
            Format::R16G16B16A16 => Ok(ColorFormat::Rgba16),
            _ => Err(Error::UnsupportedColorFormat(format)),
        }
    }
}