exr 0.7.0

//! Read and write a truly minimal RGBA image.
//! This module loads only images with RGBA channels if they all have the same data type (either f16, f32, or u32).
//! Returns `Error::Invalid` if not at least one image part with RGB channels can be found in the file.
//!
//! This module should rather be seen as an example
//! and only be used if you are confident that your images are really RGBA.
//! Use `exr::image::simple` if you need custom channels or specialized error handling.


use std::path::Path;
use std::fs::File;
use std::io::{Read, Seek, BufReader, Write, BufWriter};
use crate::math::Vec2;
use crate::error::{Result, Error, UnitResult};
use crate::meta::attributes::{PixelType, Channel, Text};
use std::convert::TryInto;
use crate::meta::{Header, ImageAttributes, LayerAttributes, MetaData};
use half::f16;
use crate::image::{ReadOptions, OnReadProgress, WriteOptions, OnWriteProgress};


/// A simple RGBA with one 32-bit float per pixel for each channel.
/// Stores all samples in a flattened vector.
#[derive(Debug, Clone, PartialEq)]
pub struct Image {

    /// A typical flattened RGBA sample array.
    /// If `has_alpha_channel` is false, this contains only RGB values.
    ///
    /// Stores in order red, green, blue, then alpha components.
    /// All lines of the image are appended one after another, __bottom to top__.
    pub data: Pixels,

    /// The dimensions of this Image, width times height.
    pub resolution: Vec2<usize>,

    /// Specifies if the `data` vector contains 3 or 4 values per pixel.
    pub has_alpha_channel: bool,

    /// Specifies if this image is in a linear color space.
    pub is_linear: bool,

    /// The attributes of the exr image.
    pub image_attributes: ImageAttributes,

    /// The attributes of the exr layer.
    pub layer_attributes: LayerAttributes,
}

/// A one dimensional array of samples.
///
/// Stores in order red, green, blue, then alpha components.
/// All lines of the image are appended one after another, top to bottom.
#[derive(Clone, PartialEq)]
pub enum Pixels {

    /// 16-bit floating point number samples in a flattened array.
    F16(Vec<f16>),

    /// 32-bit floating point number samples in a flattened array.
    F32(Vec<f32>),

    /// 32-bit unsigned integer samples in a flattened array.
    U32(Vec<u32>),
}

impl Image {

    /// Is 4 if this is an RGBA image. Is 3 if this is an RGB image.
    #[inline]
    pub fn channel_count(&self) -> usize {
        if self.has_alpha_channel { 4 } else { 3 }
    }

    /// Computes the one-dimensional index from a two-dimensional pixel location.
    /// The index points to the red component of the pixel.
    /// The green and blue pixel values can be found directly after it.
    /// Does not consider data window offset or display window offset.
    #[inline]
    pub fn vector_index_of_first_pixel_component(&self, pixel: Vec2<usize>) -> usize {
        debug_assert!(
            pixel.0 < self.resolution.0 && pixel.1 < self.resolution.1,
            "coordinate out of range"
        );

        let flat = pixel.1 * self.resolution.0 + pixel.0;
        flat * self.channel_count()
    }

    /// Read the exr image from a file.
    /// Use `read_from_unbuffered` instead, if you do not have a file.
    /// Returns `Error::Invalid` if not at least one image part with RGB channels can be found in the file.
    #[inline]
    #[must_use]
    pub fn read_from_file(path: impl AsRef<Path>, options: ReadOptions<impl OnReadProgress>) -> Result<Self> {
        Self::read_from_unbuffered(File::open(path)?, options)
    }

    /// Buffer the reader and then read the exr image from it.
    /// Use `read_from_buffered` instead, if your reader is an in-memory reader.
    /// Use `read_from_file` instead, if you have a file path.
    ///
    /// Returns `Error::Invalid` if not at least one image part with RGB channels can be found in the file.
    ///
    /// _Note: If you encounter a reader that is not send or not seek,
    /// open an issue on the github repository._
    #[inline]
    #[must_use]
    pub fn read_from_unbuffered(read: impl Read + Seek + Send, options: ReadOptions<impl OnReadProgress>) -> Result<Self> {
        Self::read_from_buffered(BufReader::new(read), options)
    }

    /// Read the exr image from a reader.
    /// Use `read_from_file` instead, if you have a file path.
    /// Use `read_from_unbuffered` instead, if this is not an in-memory reader.
    ///
    /// Returns `Error::Invalid` if not at least one image part with RGB channels can be found in the file.
    ///
    /// _Note: If you encounter a reader that is not send or not seek,
    /// open an issue on the github repository._
    #[inline]
    #[must_use]
    pub fn read_from_buffered(read: impl Read + Seek + Send, options: ReadOptions<impl OnReadProgress>) -> Result<Self> {
        crate::image::read_filtered_lines_from_buffered(
            read,

            Self::extract,

            // only keep the one header we selected earlier
            |image, header, tile| {
                tile.location.level_index == Vec2(1,1) // also skip multiresolution shenanigans
                    && header.own_attributes.name == image.layer_attributes.name // header names were checked to be unique earlier
            },

            |image, meta, line| {
                debug_assert_eq!(meta[line.location.layer].own_attributes.name, image.layer_attributes.name, "irrelevant header should be filtered out"); // TODO this should be an error right?

                let channel_index = 3 - line.location.channel; // convert ABGR index to RGBA index
                let line_position = line.location.position;
                let Vec2(width, height) = image.resolution;
                let channel_count = image.channel_count();

                let get_index_of_sample = move |sample_index| {
                    let location = line_position + Vec2(sample_index, 0);
                    debug_assert!(location.0 < width && location.1 < height, "coordinate out of range: {:?}", location);

                    let flat = location.1 * width + location.0;
                    let r_index = flat * channel_count;
                    r_index + channel_index
                };

                match &mut image.data {
                    Pixels::F16(vec) => for (sample_index, sample) in line.read_samples().enumerate() { // TODO any pixel_type?
                        vec[get_index_of_sample(sample_index)] = sample?;
                    },

                    Pixels::F32(vec) => for (sample_index, sample) in line.read_samples().enumerate() { // TODO any pixel_type?
                        vec[get_index_of_sample(sample_index)] = sample?;
                    },

                    Pixels::U32(vec) => for (sample_index, sample) in line.read_samples().enumerate() { // TODO any pixel_type?
                        vec[get_index_of_sample(sample_index)] = sample?;
                    },
                };

                Ok(())
            },

            options
        )
    }


    /// Allocate the memory for an image that could contain the described data.
    fn allocate(data_size: Vec2<usize>, linear: bool, alpha: bool, pixel_type: PixelType, image: &ImageAttributes, layer: &LayerAttributes) -> Self {
        let components = if alpha { 4 } else { 3 };
        let samples = components * data_size.area();

        Self {
            resolution: data_size,
            has_alpha_channel: alpha,

            data: match pixel_type {
                PixelType::F16 => Pixels::F16(vec![f16::from_f32(0.0); samples]),
                PixelType::F32 => Pixels::F32(vec![0.0; samples]),
                PixelType::U32 => Pixels::U32(vec![0; samples]),
            },

            is_linear: linear,

            layer_attributes: layer.clone(),
            image_attributes: image.clone(),
        }
    }

    /// Try to find a header matching the RGBA requirements.
    fn extract(headers: &[Header]) -> Result<Self> {
        let first_header_name = headers.first()
            .and_then(|header| header.own_attributes.name.as_ref());

        for (header_index, header) in headers.iter().enumerate() {
            // the following check is required because filtering works by name in this RGBA implementation
            if header_index != 0 && header.own_attributes.name.as_ref() == first_header_name {
                return Err(Error::invalid("duplicate header name"))
            }

            let channels = &header.channels.list;

            // channels are always sorted alphabetically
            let is_rgba = channels.len() == 4
                && channels[0].name == "A".try_into().unwrap()
                && channels[1].name == "B".try_into().unwrap()
                && channels[2].name == "G".try_into().unwrap()
                && channels[3].name == "R".try_into().unwrap();

            // channels are always sorted alphabetically
            let is_rgb = channels.len() == 3
                && channels[0].name == "B".try_into().unwrap()
                && channels[1].name == "G".try_into().unwrap()
                && channels[2].name == "R".try_into().unwrap();

            if !is_rgba && !is_rgb { continue; }

            let first_channel: &Channel = &channels[0];
            let pixel_type_mismatch = channels[1..].iter()
                .any(|channel|
                    channel.pixel_type != first_channel.pixel_type

                        && channel.is_linear == first_channel.is_linear
                );

            if pixel_type_mismatch { continue; }

            return Ok(Self::allocate(
                header.data_size, first_channel.is_linear, is_rgba, first_channel.pixel_type,
                &header.shared_attributes, &header.own_attributes,
            ))
        }

        Err(Error::invalid("no valid RGB or RGBA image part"))
    }

    /// Write the exr image to a file.
    /// Use `write_to_unbuffered` instead if you do not have a file.
    /// If an error occurs, attempts to delete the partially written file.
    #[must_use]
    pub fn write_to_file(&self, path: impl AsRef<Path>, options: WriteOptions<impl OnWriteProgress>) -> UnitResult {
        crate::io::attempt_delete_file_on_write_error(path, |write|
            self.write_to_unbuffered(write, options)

        )
    }

    /// Buffer the writer and then write the exr image to it.
    /// Use `read_from_buffered` instead, if your reader is an in-memory writer.
    /// Use `read_from_file` instead, if you have a file path.
    /// If your writer cannot seek, you can write to an in-memory vector of bytes first, using `write_to_buffered`.
    #[must_use]
    pub fn write_to_unbuffered(&self, write: impl Write + Seek, options: WriteOptions<impl OnWriteProgress>) -> UnitResult {
        self.write_to_buffered(BufWriter::new(write), options)
    }

    /// Write the exr image to a writer.
    /// Use `read_from_file` instead, if you have a file path.
    /// Use `read_from_unbuffered` instead, if this is not an in-memory writer.
    /// If your writer cannot seek, you can write to an in-memory vector of bytes first.
    #[must_use]
    pub fn write_to_buffered(&self, write: impl Write + Seek, options: WriteOptions<impl OnWriteProgress>) -> UnitResult {
        let pixel_type = match self.data {
            Pixels::F16(_) => PixelType::F16,
            Pixels::F32(_) => PixelType::F32,
            Pixels::U32(_) => PixelType::U32,
        };

        let header = Header::new(
            self.layer_attributes.name.clone().unwrap_or(Text::from("RGBA").unwrap()),
            self.resolution,
    if self.has_alpha_channel { smallvec![
                Channel::new("A".try_into().unwrap(), pixel_type, self.is_linear), // TODO make linear a parameter
                Channel::new("B".try_into().unwrap(), pixel_type, self.is_linear),
                Channel::new("G".try_into().unwrap(), pixel_type, self.is_linear),
                Channel::new("R".try_into().unwrap(), pixel_type, self.is_linear),
            ] }

            else { smallvec![
                Channel::new("B".try_into().unwrap(), pixel_type, self.is_linear),
                Channel::new("G".try_into().unwrap(), pixel_type, self.is_linear),
                Channel::new("R".try_into().unwrap(), pixel_type, self.is_linear),
            ] }
        );

        let header = header
            .with_shared_attributes(self.image_attributes.clone())
            .with_attributes(self.layer_attributes.clone());

        crate::image::write_all_lines_to_buffered(
            write,
            MetaData::new(smallvec![ header ]),

            |_meta, line| {
                let channel_index = 3 - line.location.channel; // convert ABGR index to RGBA index
                let line_position = line.location.position;
                let Vec2(width, height) = self.resolution;
                let channel_count = self.channel_count();

                let get_index_of_sample = move |sample_index| {
                    let location = line_position + Vec2(sample_index, 0);
                    debug_assert!(location.0 < width && location.1 < height, "coordinate out of range: {:?}", location);

                    let flat = location.1 * width + location.0;
                    let r_index = flat * channel_count;
                    r_index + channel_index
                };

                match &self.data {
                    Pixels::F16(vec) => line.write_samples(|sample_index|{
                        vec[get_index_of_sample(sample_index)]
                    })?,

                    Pixels::F32(vec) => line.write_samples(|sample_index|{
                        vec[get_index_of_sample(sample_index)]
                    })?,

                    Pixels::U32(vec) => line.write_samples(|sample_index|{
                        vec[get_index_of_sample(sample_index)]
                    })?,
                };

                Ok(())
            },

            options
        )
    }
}



impl std::fmt::Debug for Pixels {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Pixels::F16(ref vec) => write!(formatter, "[F16; {}]", vec.len()),
            Pixels::F32(ref vec) => write!(formatter, "[F32; {}]", vec.len()),
            Pixels::U32(ref vec) => write!(formatter, "[U32; {}]", vec.len()),
        }
    }
}