fits-io 0.2.0

use crate::hdu::hdu::HDU;
use crate::header::{BayerPattern, ImageType};
use crate::image::{Group, Image};
use image::{ImageBuffer, Luma, Primitive};
use std::error::Error;
use std::fmt;

/// A stream of `(x, y, value)` triples, with `value` normalised to `0.0..=1.0`.
#[cfg(feature = "tokio")]
pub type NormalisedImageStream<'a> = futures::stream::BoxStream<'a, (u32, u32, f64)>;

/// An HDU whose data section is an image, or a stack of them.
pub trait ImageHDU: HDU + fmt::Debug + Send + Sync {
    /// How many images this HDU holds.
    fn image_count(&self) -> usize;
    /// The width of every image here, from NAXIS1.
    fn images_width(&self) -> u32;
    /// The height of every image here, from NAXIS2.
    fn images_height(&self) -> u32;
    /// The colour filter layout over the sensor, or `None` if it was monochrome.
    fn images_bayer_pattern(&self) -> Option<BayerPattern>;
    /// Whether these are light, dark, flat or bias frames.
    fn images_type(&self) -> Option<&ImageType>;
    /// How long the exposure lasted.
    fn images_exposure_time(&self) -> Option<std::time::Duration>;
    /// Reads one image, or `None` past the last one.
    fn read_image(&self, index: usize) -> Result<Option<Image>, Box<dyn Error + Send + Sync>>;

    /// How many groups this HDU holds, under the random-groups convention.
    ///
    /// Zero for an ordinary image HDU, which is nearly all of them; see
    /// [`Group`](crate::image::Group).
    fn group_count(&self) -> usize {
        0
    }

    /// Reads one group, or `None` past the last one.
    fn read_group(&self, index: usize) -> Result<Option<Group>, Box<dyn Error + Send + Sync>>;

    /// Replaces the images with 8-bit ones, taking their size from the first.
    fn set_images_u8(
        &mut self,
        images: &[&ImageBuffer<Luma<u8>, Vec<u8>>],
    ) -> Result<(), Box<dyn Error + Send + Sync>> {
        get_raw_data_from_image(self, images, Self::set_raw_images_u8)
    }
    /// Replaces the images with signed 16-bit ones.
    fn set_images_i16(
        &mut self,
        images: &[&ImageBuffer<Luma<i16>, Vec<i16>>],
    ) -> Result<(), Box<dyn Error + Send + Sync>> {
        get_raw_data_from_image(self, images, Self::set_raw_images_i16)
    }
    /// Replaces the images with signed 32-bit ones.
    fn set_images_i32(
        &mut self,
        images: &[&ImageBuffer<Luma<i32>, Vec<i32>>],
    ) -> Result<(), Box<dyn Error + Send + Sync>> {
        get_raw_data_from_image(self, images, Self::set_raw_images_i32)
    }
    /// Replaces the images with single precision floating point ones.
    fn set_images_f32(
        &mut self,
        images: &[&ImageBuffer<Luma<f32>, Vec<f32>>],
    ) -> Result<(), Box<dyn Error + Send + Sync>> {
        get_raw_data_from_image(self, images, Self::set_raw_images_f32)
    }
    /// Replaces the images with double precision floating point ones.
    fn set_images_f64(
        &mut self,
        images: &[&ImageBuffer<Luma<f64>, Vec<f64>>],
    ) -> Result<(), Box<dyn Error + Send + Sync>> {
        get_raw_data_from_image(self, images, Self::set_raw_images_f64)
    }

    /// Removes every image, leaving a header-only HDU.
    fn clear_images(&mut self) -> Result<(), Box<dyn Error + Send + Sync>>;

    /// Replaces the data with an array of `shape`, given as raw 8-bit samples.
    ///
    /// `shape` is the NAXISn cards in order, fastest-varying axis first, so a
    /// stack of three 640 by 480 images is `[640, 480, 3]`. Any number of axes
    /// is allowed; [`set_raw_images_u8`] is the everyday two- and three-axis
    /// form of the same thing.
    ///
    /// The header's BITPIX and NAXISn cards are brought into line with the data.
    ///
    /// [`set_raw_images_u8`]: ImageHDU::set_raw_images_u8
    fn set_raw_array_u8(
        &mut self,
        shape: &[u32],
        values: &[u8],
    ) -> Result<(), Box<dyn Error + Send + Sync>>;

    /// Replaces the data with an array of `shape`, given as raw signed 16-bit
    /// samples.
    fn set_raw_array_i16(
        &mut self,
        shape: &[u32],
        values: &[i16],
    ) -> Result<(), Box<dyn Error + Send + Sync>>;

    /// Replaces the data with an array of `shape`, given as raw signed 32-bit
    /// samples.
    fn set_raw_array_i32(
        &mut self,
        shape: &[u32],
        values: &[i32],
    ) -> Result<(), Box<dyn Error + Send + Sync>>;

    /// Replaces the data with an array of `shape`, given as raw single precision
    /// samples.
    fn set_raw_array_f32(
        &mut self,
        shape: &[u32],
        values: &[f32],
    ) -> Result<(), Box<dyn Error + Send + Sync>>;

    /// Replaces the data with an array of `shape`, given as raw double precision
    /// samples.
    fn set_raw_array_f64(
        &mut self,
        shape: &[u32],
        values: &[f64],
    ) -> Result<(), Box<dyn Error + Send + Sync>>;

    /// Replaces the images with raw 8-bit samples, `width` by `height` each.
    ///
    /// The header's BITPIX and NAXISn cards are brought into line with them.
    fn set_raw_images_u8(
        &mut self,
        width: u32,
        height: u32,
        images: &[&[u8]],
    ) -> Result<(), Box<dyn Error + Send + Sync>> {
        set_planes(self, width, height, images, Self::set_raw_array_u8)
    }
    /// Replaces the images with raw signed 16-bit samples.
    fn set_raw_images_i16(
        &mut self,
        width: u32,
        height: u32,
        images: &[&[i16]],
    ) -> Result<(), Box<dyn Error + Send + Sync>> {
        set_planes(self, width, height, images, Self::set_raw_array_i16)
    }
    /// Replaces the images with raw signed 32-bit samples.
    fn set_raw_images_i32(
        &mut self,
        width: u32,
        height: u32,
        images: &[&[i32]],
    ) -> Result<(), Box<dyn Error + Send + Sync>> {
        set_planes(self, width, height, images, Self::set_raw_array_i32)
    }
    /// Replaces the images with raw single precision samples.
    fn set_raw_images_f32(
        &mut self,
        width: u32,
        height: u32,
        images: &[&[f32]],
    ) -> Result<(), Box<dyn Error + Send + Sync>> {
        set_planes(self, width, height, images, Self::set_raw_array_f32)
    }
    /// Replaces the images with raw double precision samples.
    fn set_raw_images_f64(
        &mut self,
        width: u32,
        height: u32,
        images: &[&[f64]],
    ) -> Result<(), Box<dyn Error + Send + Sync>> {
        set_planes(self, width, height, images, Self::set_raw_array_f64)
    }

    /// Streams one image as `(x, y, value)` triples, normalised to `0.0..=1.0`.
    #[cfg(feature = "tokio")]
    fn stream_normalised_image(
        &self,
        index: usize,
    ) -> Result<Option<NormalisedImageStream<'_>>, Box<dyn Error + Send + Sync>>;
    /// How many bytes one image occupies.
    fn image_data_size(&self) -> u64;

    /// Whether this HDU's image is stored tile-compressed inside a table.
    fn is_compressed(&self) -> bool {
        self.header().is_compressed_image()
    }

    /// Stores this HDU's image tile-compressed, as `fpack` would.
    ///
    /// The image is cut into tiles, each tile is compressed on its own, and the
    /// result is written as a binary table whose header says what image it
    /// stands for. Everything that reads an image here goes on working — the
    /// HDU is still an image as far as this crate is concerned, and it is
    /// written out as a compressed image extension.
    ///
    /// Compressing an already compressed HDU decompresses it first, so that
    /// changing the settings does not compress the tiles twice.
    ///
    /// ```no_run
    /// # #[cfg(feature = "fs")]
    /// # fn example() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
    /// use fits_io::Fits;
    /// use fits_io::fs::FsFits;
    /// use fits_io::hdu::ImageHDU;
    /// use fits_io::image::compression::{Compression, CompressionOptions};
    ///
    /// let mut fits = FsFits::open("observation.fits".as_ref())?;
    ///
    /// fits.primary_hdu_mut()
    ///     .compress(&CompressionOptions::new(Compression::Rice))?;
    ///
    /// let smaller = fits.to_vec()?;
    /// # Ok(())
    /// # }
    /// # fn main() {}
    /// ```
    ///
    /// # Errors
    ///
    /// Returns an error when the image cannot be compressed the way the options
    /// ask — Rice coding a floating point image without quantising it, say —
    /// and when its data cannot be read.
    fn compress(
        &mut self,
        options: &crate::image::compression::CompressionOptions,
    ) -> Result<(), Box<dyn Error + Send + Sync>>;

    /// Stores this HDU's image plainly again, undoing [`ImageHDU::compress`].
    ///
    /// An HDU that was not compressed is left alone.
    ///
    /// # Errors
    ///
    /// Returns an error when the compressed data cannot be read.
    fn decompress(&mut self) -> Result<(), Box<dyn Error + Send + Sync>>;
}

/// Lays a set of equally sized planes out as one array and stores it.
///
/// This is what the `set_raw_images_*` methods are: a shape of two axes for one
/// image and three for a stack of them.
fn set_planes<T: Copy, S: ImageHDU + ?Sized>(
    hdu: &mut S,
    width: u32,
    height: u32,
    images: &[&[T]],
    set: impl FnOnce(&mut S, &[u32], &[T]) -> Result<(), Box<dyn Error + Send + Sync>>,
) -> Result<(), Box<dyn Error + Send + Sync>> {
    if images.is_empty() {
        return hdu.clear_images();
    }

    let pixels = (width as usize)
        .checked_mul(height as usize)
        .ok_or("Image dimensions overflow the address space")?;

    // Checked here rather than left to the shape, so that a ragged set says
    // which image is the wrong size.
    for (index, image) in images.iter().enumerate() {
        if image.len() != pixels {
            return Err(format!(
                "Image {} has {} pixels, but a {}x{} image has {}",
                index,
                image.len(),
                width,
                height,
                pixels
            )
            .into());
        }
    }

    let mut shape = vec![width, height];
    if images.len() > 1 {
        shape.push(images.len() as u32);
    }

    let values: Vec<T> = images
        .iter()
        .flat_map(|image| image.iter().copied())
        .collect();

    set(hdu, &shape, &values)
}

fn get_raw_data_from_image<
    'a,
    T: Primitive,
    S: ImageHDU + ?Sized,
    CB: FnOnce(&mut S, u32, u32, &[&[T]]) -> Result<(), Box<dyn Error + Send + Sync>>,
>(
    hdu: &mut S,
    images: &'a [&'a ImageBuffer<Luma<T>, Vec<T>>],
    callback: CB,
) -> Result<(), Box<dyn Error + Send + Sync>> {
    if images.is_empty() {
        hdu.clear_images()?;
        Ok(())
    } else {
        let width = images[0].width();
        let height = images[0].height();

        let data = images
            .iter()
            .map(|image| image.iter().as_slice())
            .collect::<Vec<_>>();

        callback(hdu, width, height, &data)
    }
}