exr 1.74.1

//! How to read a set of resolution levels.

use crate::{
    block::{chunk::TileCoordinates, lines::LineRef, samples::*},
    error::*,
    image::{
        read::{any_channels::*, specific_channels::*},
        recursive::*,
        *,
    },
    math::Vec2,
    meta::{attribute::*, header::Header, *},
};

// Note: In the resulting image, the `FlatSamples` are placed
// directly inside the channels, without `LargestLevel<>` indirection
/// Specify to read only the highest resolution level, skipping all smaller
/// variations. The sample storage can be [`ReadFlatSamples`].
#[derive(Debug, Clone, Eq, PartialEq)]
pub struct ReadLargestLevel<DeepOrFlatSamples> {
    /// The sample reading specification
    pub read_samples: DeepOrFlatSamples,
}

// FIXME rgba levels???

// Read the largest level, directly, without intermediate structs
impl<DeepOrFlatSamples> ReadLargestLevel<DeepOrFlatSamples> {
    /// Read all arbitrary channels in each layer.
    pub fn all_channels(self) -> ReadAnyChannels<DeepOrFlatSamples> {
        ReadAnyChannels {
            read_samples: self.read_samples,
        }
    }

    // Instead of Self, the `FlatSamples` are used directly

    /// Read only layers that contain rgba channels. Skips any other channels in
    /// the layer. The alpha channel will contain the value `1.0` if no
    /// alpha channel can be found in the image.
    ///
    /// Using two closures, define how to store the pixels.
    /// The first closure creates an image, and the second closure inserts a
    /// single pixel. The type of the pixel can be defined by the second
    /// closure; it must be a tuple containing four values, each being
    /// either `f16`, `f32`, `u32` or `Sample`.
    ///
    /// Throws an error for images with deep data or subsampling.
    /// Use `specific_channels` or `all_channels` if you want to read something
    /// other than rgba.
    pub fn rgba_channels<R, G, B, A, Create, Set, Pixels>(
        self,
        create_pixels: Create,
        set_pixel: Set,
    ) -> CollectPixels<
        ReadOptionalChannel<
            ReadRequiredChannel<ReadRequiredChannel<ReadRequiredChannel<NoneMore, R>, G>, B>,
            A,
        >,
        (R, G, B, A),
        Pixels,
        Create,
        Set,
    >
    where
        R: FromNativeSample,
        G: FromNativeSample,
        B: FromNativeSample,
        A: FromNativeSample,
        Create: Fn(Vec2<usize>, &RgbaChannels) -> Pixels,
        Set: Fn(&mut Pixels, Vec2<usize>, (R, G, B, A)),
    {
        self.specific_channels()
            .required("R")
            .required("G")
            .required("B")
            .optional("A", A::from_f32(1.0))
            .collect_pixels(create_pixels, set_pixel)
    }

    /// Read only layers that contain rgb channels. Skips any other channels in
    /// the layer.
    ///
    /// Using two closures, define how to store the pixels.
    /// The first closure creates an image, and the second closure inserts a
    /// single pixel. The type of the pixel can be defined by the second
    /// closure; it must be a tuple containing three values, each being
    /// either `f16`, `f32`, `u32` or `Sample`.
    ///
    /// Throws an error for images with deep data or subsampling.
    /// Use `specific_channels` or `all_channels` if you want to read something
    /// other than rgb.
    pub fn rgb_channels<R, G, B, Create, Set, Pixels>(
        self,
        create_pixels: Create,
        set_pixel: Set,
    ) -> CollectPixels<
        ReadRequiredChannel<ReadRequiredChannel<ReadRequiredChannel<NoneMore, R>, G>, B>,
        (R, G, B),
        Pixels,
        Create,
        Set,
    >
    where
        R: FromNativeSample,
        G: FromNativeSample,
        B: FromNativeSample,
        Create: Fn(Vec2<usize>, &RgbChannels) -> Pixels,
        Set: Fn(&mut Pixels, Vec2<usize>, (R, G, B)),
    {
        self.specific_channels()
            .required("R")
            .required("G")
            .required("B")
            .collect_pixels(create_pixels, set_pixel)
    }

    /// Read only layers that contain the specified channels, skipping any other
    /// channels in the layer. Further specify which channels should be
    /// included by calling `.required("ChannelName")` or `.optional("
    /// ChannelName", default_value)` on the result of this function.
    /// Call `collect_pixels` afterwards to define the pixel container for your
    /// set of channels.
    ///
    /// Throws an error for images with deep data or subsampling.
    pub fn specific_channels(self) -> ReadZeroChannels {
        ReadZeroChannels {}
    }
}

/// Specify to read all contained resolution levels from the image, if any.
#[derive(Debug, Clone, Eq, PartialEq)]
pub struct ReadAllLevels<DeepOrFlatSamples> {
    /// The sample reading specification
    pub read_samples: DeepOrFlatSamples,
}

impl<ReadDeepOrFlatSamples> ReadAllLevels<ReadDeepOrFlatSamples> {
    /// Read all arbitrary channels in each layer.
    pub fn all_channels(self) -> ReadAnyChannels<Self> {
        ReadAnyChannels {
            read_samples: self,
        }
    }

    // TODO specific channels for multiple resolution levels
}

// pub struct ReadLevels<S> {
// read_samples: S,
// }

/// Processes pixel blocks from a file and accumulates them into multiple levels
/// per channel.
#[derive(Debug, Clone, Eq, PartialEq)]
pub struct AllLevelsReader<SamplesReader> {
    levels: Levels<SamplesReader>,
}

/// A template that creates a [`SamplesReader`] once for each resolution level.
pub trait ReadSamplesLevel {
    /// The type of the temporary level reader
    type Reader: SamplesReader;

    /// Create a single reader for a single resolution level
    fn create_samples_level_reader(
        &self,
        header: &Header,
        channel: &ChannelDescription,
        level: Vec2<usize>,
        resolution: Vec2<usize>,
    ) -> Result<Self::Reader>;
}

impl<S: ReadSamplesLevel> ReadSamples for ReadAllLevels<S> {
    type Reader = AllLevelsReader<S::Reader>;

    fn create_sample_reader(
        &self,
        header: &Header,
        channel: &ChannelDescription,
    ) -> Result<Self::Reader> {
        let data_size = header.layer_size / channel.sampling;

        let levels = {
            if let crate::meta::BlockDescription::Tiles(tiles) = &header.blocks {
                match tiles.level_mode {
                    LevelMode::Singular => {
                        Levels::Singular(self.read_samples.create_samples_level_reader(
                            header,
                            channel,
                            Vec2(0, 0),
                            header.layer_size,
                        )?)
                    }

                    LevelMode::MipMap => Levels::Mip {
                        rounding_mode: tiles.rounding_mode,
                        level_data: {
                            let round = tiles.rounding_mode;
                            let maps: Result<LevelMaps<S::Reader>> =
                                mip_map_levels(round, data_size)
                                    .map(|(index, level_size)| {
                                        self.read_samples.create_samples_level_reader(
                                            header,
                                            channel,
                                            Vec2(index, index),
                                            level_size,
                                        )
                                    })
                                    .collect();

                            maps?
                        },
                    },

                    // TODO put this into Levels::new(..) ?
                    LevelMode::RipMap => Levels::Rip {
                        rounding_mode: tiles.rounding_mode,
                        level_data: {
                            let round = tiles.rounding_mode;
                            let level_count_x = compute_level_count(round, data_size.width());
                            let level_count_y = compute_level_count(round, data_size.height());
                            let maps: Result<LevelMaps<S::Reader>> =
                                rip_map_levels(round, data_size)
                                    .map(|(index, level_size)| {
                                        self.read_samples.create_samples_level_reader(
                                            header, channel, index, level_size,
                                        )
                                    })
                                    .collect();

                            RipMaps {
                                map_data: maps?,
                                level_count: Vec2(level_count_x, level_count_y),
                            }
                        },
                    },
                }
            }
            // scan line blocks never have mip maps
            else {
                Levels::Singular(self.read_samples.create_samples_level_reader(
                    header,
                    channel,
                    Vec2(0, 0),
                    data_size,
                )?)
            }
        };

        Ok(AllLevelsReader {
            levels,
        })
    }
}

impl<S: SamplesReader> SamplesReader for AllLevelsReader<S> {
    type Samples = Levels<S::Samples>;

    fn filter_block(&self, _: TileCoordinates) -> bool {
        true
    }

    fn read_line(&mut self, line: LineRef<'_>) -> UnitResult {
        self.levels.get_level_mut(line.location.level)?.read_line(line)
    }

    fn into_samples(self) -> Self::Samples {
        match self.levels {
            Levels::Singular(level) => Levels::Singular(level.into_samples()),
            Levels::Mip {
                rounding_mode,
                level_data,
            } => Levels::Mip {
                rounding_mode,
                level_data: level_data.into_iter().map(|s| s.into_samples()).collect(),
            },

            Levels::Rip {
                rounding_mode,
                level_data,
            } => Levels::Rip {
                rounding_mode,
                level_data: RipMaps {
                    level_count: level_data.level_count,
                    map_data: level_data.map_data.into_iter().map(|s| s.into_samples()).collect(),
                },
            },
        }
    }
}