ndbioimage 0.2.0

use crate::axes::Axis;
use crate::colors::Color;
use crate::error::Error;
use crate::metadata::Metadata;
use crate::readers::{DynReader, PixelType, Reader};
use crate::stats::MinMax;
use crate::utils::progress::get_bar;
use crate::view::{Number, View};
use indicatif::ProgressBar;
use itertools::iproduct;
use ndarray::{Array0, Array1, Array2, ArrayD, Dimension};
use rayon::prelude::*;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Condvar, Mutex};
use tiffwrite::{Bytes, Colors, Compression, IJTiffFile};

/// options for saving tiff files
#[derive(Debug, Clone)]
pub struct TiffOptions {
    bar: Option<ProgressBar>,
    compression: Compression,
    colors: Option<Vec<Vec<u8>>>,
    overwrite: bool,
}

impl Default for TiffOptions {
    fn default() -> Self {
        Self {
            bar: None,
            compression: Compression::Zstd(10),
            colors: None,
            overwrite: false,
        }
    }
}

impl TiffOptions {
    /// create tiff options with the given parameters
    pub fn new(
        bar: Option<ProgressBar>,
        compression: Option<Compression>,
        colors: Vec<String>,
        overwrite: bool,
    ) -> Result<Self, Error> {
        let mut options = Self {
            bar,
            compression: compression.unwrap_or(Compression::Zstd(10)),
            colors: None,
            overwrite,
        };
        if !colors.is_empty() {
            options.set_colors(&colors)?;
        }
        Ok(options)
    }

    /// show a progress bar while saving tiff
    pub fn enable_bar(&mut self, message: Option<String>) {
        self.bar = Some(get_bar(Some(0), message));
    }

    /// do not show a progress bar while saving tiff
    pub fn disable_bar(&mut self) {
        self.bar = None;
    }

    /// save tiff with zstd compression (default)
    pub fn set_zstd_compression(&mut self) {
        self.compression = Compression::Zstd(10)
    }

    /// save tiff with zstd compression, choose a level between 7..=22
    pub fn set_zstd_compression_level(&mut self, level: i32) {
        self.compression = Compression::Zstd(level)
    }

    /// save tiff with deflate compression
    pub fn set_deflate_compression(&mut self) {
        self.compression = Compression::Deflate
    }

    /// set the color lookup table for the tiff
    pub fn set_colors(&mut self, colors: &[String]) -> Result<(), Error> {
        let colors = colors
            .iter()
            .map(|c| c.parse::<Color>())
            .collect::<Result<Vec<_>, Error>>()?;
        self.colors = Some(colors.into_iter().map(|c| c.to_rgb()).collect());
        Ok(())
    }

    /// set whether an existing tiff file should be overwritten
    pub fn set_overwrite(&mut self, overwrite: bool) {
        self.overwrite = overwrite;
    }
}

impl Drop for TiffOptions {
    fn drop(&mut self) {
        if let Some(bar) = self.bar.take() {
            bar.finish()
        }
    }
}

impl<D> View<D>
where
    D: Dimension,
{
    /// save as tiff with a certain type
    pub fn save_as_tiff_with_type<T, P>(&self, path: P, options: &TiffOptions) -> Result<(), Error>
    where
        P: AsRef<Path>,
        T: Bytes + Number + Send + Sync,
        ArrayD<T>: MinMax<Output = ArrayD<T>>,
        Array1<T>: MinMax<Output = Array0<T>>,
        Array2<T>: MinMax<Output = Array1<T>>,
    {
        let path = path.as_ref().to_path_buf();
        if path.exists() {
            if options.overwrite {
                std::fs::remove_file(&path)?;
            } else {
                return Err(Error::FileAlreadyExists(path.display().to_string()));
            }
        }
        let shape = self.shape();
        let mut tiff = IJTiffFile::new(path)?;
        tiff.set_compression(options.compression);
        let metadata = self.metadata()?;
        tiff.px_size = metadata.pixel_size()?.map(|i| i / 1e3);
        tiff.time_interval = metadata.time_interval()?.map(|i| i / 1e3);
        tiff.delta_z = metadata.delta_z()?.map(|i| i / 1e3);
        tiff.comment = Some(metadata.summary()?);
        if let Some(mut colors) = options.colors.clone() {
            while colors.len() < shape.c {
                colors.push(vec![255, 255, 255]);
            }
            tiff.colors = Colors::Colors(colors);
        }
        let tiff = Arc::new(Mutex::new(tiff));
        if let Some(bar) = options.bar.as_ref() {
            bar.inc_length((shape.c * shape.z * shape.t) as u64);
        }
        iproduct!(0..shape.c, 0..shape.z, 0..shape.t)
            .collect::<Vec<_>>()
            .into_iter()
            .try_for_each(|(c, z, t)| {
                if let Ok(mut tiff) = tiff.lock() {
                    tiff.save(&self.get_frame::<T, _>(c, z, t)?, c, z, t)?;
                    if let Some(bar) = options.bar.as_ref() {
                        bar.inc(1);
                    }
                    Ok(())
                } else {
                    Err(Error::TiffLock)
                }
            })?;
        Ok(())
    }

    /// save as tiff with whatever pixel type the view has
    pub fn save_as_tiff<P>(&self, path: P, options: &TiffOptions) -> Result<(), Error>
    where
        P: AsRef<Path>,
    {
        match self.pixel_type() {
            PixelType::I8 => self.save_as_tiff_with_type::<i8, P>(path, options)?,
            PixelType::U8 => self.save_as_tiff_with_type::<u8, P>(path, options)?,
            PixelType::I16 => self.save_as_tiff_with_type::<i16, P>(path, options)?,
            PixelType::U16 => self.save_as_tiff_with_type::<u16, P>(path, options)?,
            PixelType::I32 => self.save_as_tiff_with_type::<i32, P>(path, options)?,
            PixelType::U32 => self.save_as_tiff_with_type::<u32, P>(path, options)?,
            PixelType::F32 => self.save_as_tiff_with_type::<f32, P>(path, options)?,
            PixelType::F64 => self.save_as_tiff_with_type::<f64, P>(path, options)?,
            PixelType::I64 => self.save_as_tiff_with_type::<i64, P>(path, options)?,
            PixelType::U64 => self.save_as_tiff_with_type::<u64, P>(path, options)?,
            PixelType::I128 => self.save_as_tiff_with_type::<i64, P>(path, options)?,
            PixelType::U128 => self.save_as_tiff_with_type::<u64, P>(path, options)?,
            PixelType::F128 => self.save_as_tiff_with_type::<f64, P>(path, options)?,
        }

        Ok(())
    }
}

/// batch convert multiple files to tiff format
pub fn batch_to_tiff(
    files_in: &[PathBuf],
    files_out: &[PathBuf],
    operations: Option<Vec<(String, String)>>,
    colors: Option<Vec<String>>,
    overwrite: bool,
    bar: bool,
    message: Option<String>,
) -> Result<(), Error> {
    let bar = if bar {
        Some(get_bar(
            Some(0),
            Some(message.unwrap_or("writing tiff files".to_string())),
        ))
    } else {
        None
    };
    let options = TiffOptions::new(bar, None, colors.unwrap_or_default().clone(), overwrite)?;
    let semaphore = Arc::new((Mutex::new(0usize), Condvar::new()));
    files_in
        .iter()
        .zip(files_out)
        .collect::<Vec<_>>()
        .into_par_iter()
        .map(|(file_in, file_out)| {
            let (lock, cvar) = &*semaphore;
            {
                let mut count = lock.lock().unwrap();
                while *count >= 10 {
                    count = cvar.wait(count).unwrap();
                }
                *count += 1;
            }
            let mut view = View::<_, DynReader>::from_path(file_in)?.into_dyn();
            if let Some(operations) = operations.as_ref() {
                for (ax, op) in operations {
                    view = view.operate(ax.parse::<Axis>()?, op.parse()?)?;
                }
            }
            view.save_as_tiff(file_out, &options)?;
            {
                let mut count = lock.lock().unwrap();
                *count -= 1;
                cvar.notify_one();
            }
            Ok(())
        })
        .collect::<Result<Vec<()>, Error>>()?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use crate::error::Error;
    use crate::readers::DynReader;
    use crate::tiffwrite::{TiffOptions, batch_to_tiff, get_bar};
    use crate::view::View;
    use std::fs::create_dir_all;
    use std::path::PathBuf;

    #[cfg(any(
        feature = "czi",
        feature = "tiffseq",
        feature = "tiff",
        feature = "bioformats_java"
    ))]
    #[test]
    fn tiff() -> Result<(), Error> {
        let file = if cfg!(any(feature = "czi", feature = "bioformats_java")) {
            "czi/1xp53-01-AP1.czi"
        } else if cfg!(feature = "tiff") {
            "tiff/20251014_20-Pos_000_000_loc_results_Cy3.tif"
        } else if cfg!(feature = "tiffseq") {
            "tiffseq/20-Pos_005_005"
        } else {
            unreachable!(
                "need to enable one of these features: czi, bioformats_java, tiff, tiffseq"
            );
        };
        let path = std::env::current_dir()?
            .join("tests")
            .join("files")
            .join(file);
        let view: View<_, DynReader> = View::from_path(&path)?;
        println!("{}", view.summary()?);
        let bar = Some(get_bar(Some(0), Some("writing tiff file".to_string())));
        let options = TiffOptions::new(bar, None, Vec::new(), true)?;
        view.save_as_tiff(
            std::env::home_dir().unwrap().join("tmp/movie.tif"),
            &options,
        )?;
        Ok(())
    }

    #[cfg(any(feature = "tiffseq", feature = "tiff", feature = "bioformats_java"))]
    #[test]
    fn tiff_parallel() -> Result<(), Error> {
        let files = [
            #[cfg(any(feature = "tiffseq", feature = "bioformats_java"))]
            "tiffseq/YTL1841B2-2-1_1hr_DMSO_galinduction_1",
            #[cfg(any(feature = "tiff", feature = "bioformats_java"))]
            "tiff/20251014_20-Pos_000_000_max.tif",
            #[cfg(any(feature = "czi", feature = "bioformats_java"))]
            "czi/1xp53-01-AP1.czi",
            #[cfg(any(feature = "czi", feature = "bioformats_java"))]
            "czi/beads_2023_05_04__19_00_22.czi",
            #[cfg(any(feature = "czi", feature = "bioformats_java"))]
            "czi/YTL1849A131_2023_05_04__13_36_36.czi",
            #[cfg(any(feature = "czi", feature = "bioformats_java"))]
            "czi/p53_2x_3-pos_20s_SR-8Y-01_AP-Scene-3-P2.czi",
        ];
        let files_in = files
            .iter()
            .map(|file| {
                Ok(std::env::current_dir()?
                    .join("tests")
                    .join("files")
                    .join(file))
            })
            .collect::<Result<Vec<_>, Error>>()?;
        let files_out = files
            .iter()
            .map(|file| {
                std::env::home_dir()
                    .unwrap()
                    .join("tmp")
                    .join(PathBuf::from(file).with_extension("tif"))
            })
            .collect::<Vec<_>>();
        for file in &files_out {
            create_dir_all(file.parent().unwrap())?;
        }
        batch_to_tiff(&files_in, &files_out, None, None, true, true, None)?;
        Ok(())
    }
}