eorst 1.0.1

//! Processing methods for RasterDataset.
//!
//! This module contains methods for applying worker functions to raster blocks,
//! including parallel processing, reduction operations, and mosaicking.

use crate::core_types::RasterType;
use crate::blocks::RasterBlock;
use crate::metadata::RasterDataBlock;
use crate::types::Dimension as UtilsDimension;
use crate::types::{Offset, ReadWindow, Size};
use crate::gdal_utils::{create_rayon_pool, create_temp_file, file_stem_str, mosaic, mosaic_translate_cleanup, mosaic_translate_cleanup_time_steps};
use crate::parallel_writer::{self, ParallelGeoTiffWriter};
use crate::rasterdataset::RasterDataset;
use anyhow::Context;

use kdam::par_tqdm;
use ndarray::{Array2, Array3, Array4};
use num_traits::NumCast;
use rayon::{
    iter::IndexedParallelIterator,
    prelude::{IntoParallelIterator, ParallelIterator},
};
use std::fs;
use std::path::{Path, PathBuf};

impl<R> RasterDataset<R>
where
    R: RasterType,
{
    // ─── Apply with RasterDataBlock ───

    /// Applies a worker function to each block, receiving a `RasterDataBlock` with metadata.
    ///
    /// The worker receives `&RasterDataBlock<R>` which includes the block data array,
    /// `layer_indices`, `date_indices`, and other metadata. This enables semantic layer
    /// and time selection via the [`Select`](crate::selection::Select) trait.
    ///
    /// Blocks are processed in parallel using a Rayon thread pool. Each block's result
    /// is written directly to the output GeoTIFF via a [`ParallelGeoTiffWriter`](crate::parallel_writer::ParallelGeoTiffWriter),
    /// eliminating intermediate files and the subprocess-based mosaic/translate phase.
    ///
    /// Returns `Err` if any worker invocation fails.
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// // See apply_reduction below for a fully-runnable doctest.
    /// // apply() follows the same pattern but returns Array4 and writes directly to GeoTIFF:
    /// //
    /// // fn worker(block: &RasterDataBlock<i16>) -> anyhow::Result<Array4<i16>> {
    /// //     Ok(block.data.mapv(|v| v as i16))  // pass-through example
    /// // }
    /// //
    /// // rds.apply::<i16>(worker, 1, &out_path)?;
    /// ```
    pub fn apply<U>(
        &self,
        worker: fn(&RasterDataBlock<R>) -> anyhow::Result<Array4<U>>,
        n_cpus: usize,
        out_file: &Path,
    ) -> anyhow::Result<()>
    where
        U: RasterType,
    {
        let pool = create_rayon_pool(n_cpus);

        let n_times = self.metadata.shape.times;
        let n_layers = self.metadata.shape.layers;
        let total_rows = self.metadata.shape.rows;
        let total_cols = self.metadata.shape.cols;
        let epsg_code = self.metadata.epsg_code;
        let geo_transform = self.metadata.geo_transform;
        let n_bands = n_times * n_layers;
        let na_value: U = NumCast::from(self.metadata.na_value).unwrap_or(U::zero());

        pool.install(|| -> anyhow::Result<()> {
            // Phase 1: Pre-create the output GeoTIFF
            parallel_writer::create_output_geotiff::<U>(
                out_file,
                &geo_transform,
                epsg_code,
                total_cols,
                total_rows,
                n_bands,
                na_value,
            )?;

            let writer = ParallelGeoTiffWriter::new(
                out_file.to_path_buf(),
                geo_transform,
                epsg_code,
                total_cols,
                total_rows,
                n_bands,
            );

            // Phase 2: Process blocks in parallel, writing directly to output
            par_tqdm!(self.blocks
                .to_owned()
                .into_par_iter())
                .map(|raster_block| -> anyhow::Result<()> {
                    let bid = raster_block.block_index;
                    let block_data = self.read_block::<R>(bid);
                    let raster_data_block = RasterDataBlock {
                        data: block_data,
                        metadata: self.metadata.clone(),
                        no_data: NumCast::from(0).unwrap(),
                    };
                    let result = worker(&raster_data_block)?;

                    // Flatten 4D (times, layers, rows, cols) → 3D (n_bands, rows, cols)
                    let rows = result.shape()[2];
                    let cols = result.shape()[3];
                    let flat = result
                        .into_shape_with_order((n_bands, rows, cols))
                        .context("Failed to reshape 4D result to 3D for writing")?;

                    // Trim overlap and adjust window
                    let overlap = raster_block.overlap;
                    let trimmed = crate::array_ops::trimm_array3_asymmetric(&flat, &overlap);
                    let trimmed_rows = trimmed.shape()[1];
                    let trimmed_cols = trimmed.shape()[2];

                    let write_window = ReadWindow {
                        offset: Offset {
                            rows: raster_block.read_window.offset.rows + overlap.top as isize,
                            cols: raster_block.read_window.offset.cols + overlap.left as isize,
                        },
                        size: Size {
                            rows: trimmed_rows as isize,
                            cols: trimmed_cols as isize,
                        },
                    };

                    parallel_writer::write_block(&writer, trimmed, write_window)?;

                    Ok(())
                })
                .collect::<anyhow::Result<Vec<_>>>()?;

            Ok(())
        })
    }

    /// Applies a worker function to each block and outputs a Cloud Optimized GeoTIFF.
    ///
    /// This is the COG variant of [`apply`](Self::apply). It follows the same block-processing
    /// pattern but produces a proper COG with overviews and IFD reordering.
    ///
    /// The output process:
    /// 1. Write blocks to an intermediate GeoTIFF (same as `apply`)
    /// 2. Build overviews in-place via GDAL's `build_overviews`
    /// 3. Translate to final COG via `gdal_translate -of COG` for proper IFD ordering
    /// 4. Clean up intermediate file
    ///
    /// # Arguments
    /// * `worker` - Worker function receiving `&RasterDataBlock<R>`, returning `Array4<U>`
    /// * `n_cpus` - Number of CPU threads for parallel processing
    /// * `out_file` - Output COG file path
    /// * `config` - Output configuration controlling format, compression, overviews
    pub fn apply_cog<U>(
        &self,
        worker: fn(&RasterDataBlock<R>) -> anyhow::Result<Array4<U>>,
        n_cpus: usize,
        out_file: &Path,
        config: &crate::types::OutputConfig,
    ) -> anyhow::Result<()>
    where
        U: RasterType,
    {
        use crate::types::OutputFormat;

        match config.format {
            OutputFormat::GeoTiff => {
                // Simple path: no overviews, no COG translation
                self.apply::<U>(worker, n_cpus, out_file)
            }
            OutputFormat::GeoTiffOverviews => {
                // Write GeoTIFF, then build overviews in-place
                self.apply::<U>(worker, n_cpus, out_file)?;

                let writer = ParallelGeoTiffWriter::new(
                    out_file.to_path_buf(),
                    self.metadata.geo_transform,
                    self.metadata.epsg_code,
                    self.metadata.shape.cols,
                    self.metadata.shape.rows,
                    self.metadata.shape.times * self.metadata.shape.layers,
                );
                writer.build_overviews(&config.overview_resampling, &config.overview_levels)?;
                Ok(())
            }
            OutputFormat::COG => {
                // Write to intermediate, then COG driver creates overviews + reorders IFDs
                let intermediate = PathBuf::from(create_temp_file("tif"));

                self.apply::<U>(worker, n_cpus, &intermediate)?;

                crate::gdal_utils::translate_to_cog(
                    &intermediate, out_file,
                    &config.compression,
                    &config.overview_resampling,
                )?;
                fs::remove_file(&intermediate).ok();

                Ok(())
            }
        }
    }

    /// Applies a worker function to pairs of blocks from two datasets.
    ///
    /// The worker receives `&RasterDataBlock<R>` and `&RasterDataBlock<U>` which include
    /// `layer_indices`, `date_indices`, and other metadata. This enables semantic selection
    /// via the [`Select`](crate::selection::Select) trait on both datasets.
    ///
    /// Useful for masking, zonal statistics, or combining two rasters (e.g. applying
    /// a cloud mask, or computing NDVI with a land cover classification).
    ///
    /// Blocks are processed in parallel using a Rayon thread pool. Each block pair's result
    /// is written directly to the output GeoTIFF.
    ///
    /// Returns `Err` if any worker invocation fails.
    ///
    /// ```rust,ignore
    /// // See apply() for the primary usage pattern. apply_with_mask() takes
    /// // a second dataset whose blocks are paired with the primary dataset's
    /// // blocks, and a worker that receives both blocks simultaneously.
    /// // Example worker signature:
    /// //
    /// // fn masked_ndvi(
    /// //     raster: &RasterDataBlock<i16>,
    /// //     mask: &RasterDataBlock<i16>,
    /// // ) -> anyhow::Result<Array4<i16>> {
    /// //     // mask.data is used to nullify invalid pixels
    /// //     ...
    /// // }
    /// //
    /// // rds.apply_with_mask::<i16, i16>(&mask_rds, masked_ndvi, 8, &out_path)?;
    /// ```
    pub fn apply_with_mask<U, V>(
        &self,
        other: &RasterDataset<U>,
        worker: fn(&RasterDataBlock<R>, &RasterDataBlock<U>) -> anyhow::Result<Array4<V>>,
        n_cpus: usize,
        out_file: &Path,
    ) -> anyhow::Result<()>
    where
        U: RasterType,
        V: RasterType,
    {
        let pool = create_rayon_pool(n_cpus);

        let n_times = self.metadata.shape.times;
        let n_layers = self.metadata.shape.layers;
        let total_rows = self.metadata.shape.rows;
        let total_cols = self.metadata.shape.cols;
        let epsg_code = self.metadata.epsg_code;
        let geo_transform = self.metadata.geo_transform;
        let n_bands = n_times * n_layers;
        let na_value: V = NumCast::from(self.metadata.na_value).unwrap_or(V::zero());

        pool.install(|| -> anyhow::Result<()> {
            // Phase 1: Pre-create the output GeoTIFF
            parallel_writer::create_output_geotiff::<V>(
                out_file,
                &geo_transform,
                epsg_code,
                total_cols,
                total_rows,
                n_bands,
                na_value,
            )?;

            let writer = ParallelGeoTiffWriter::new(
                out_file.to_path_buf(),
                geo_transform,
                epsg_code,
                total_cols,
                total_rows,
                n_bands,
            );

            // Phase 2: Process block pairs in parallel, writing directly to output
            par_tqdm!(self.blocks
                .to_owned()
                .into_par_iter()
                .zip(other.blocks.to_owned()))
                .map(|(raster_block_data, _raster_block_mask)| -> anyhow::Result<()> {
                    let bid = raster_block_data.block_index;
                    let block_data = self.read_block::<R>(bid);
                    let mask_data = other.read_block::<U>(bid);

                    let rdb_data = RasterDataBlock {
                        data: block_data,
                        metadata: self.metadata.clone(),
                        no_data: NumCast::from(0).unwrap(),
                    };
                    let rdb_mask = RasterDataBlock {
                        data: mask_data,
                        metadata: other.metadata.clone(),
                        no_data: NumCast::from(0).unwrap(),
                    };

                    let result = worker(&rdb_data, &rdb_mask)?;

                    // Flatten 4D (times, layers, rows, cols) → 3D (n_bands, rows, cols)
                    let rows = result.shape()[2];
                    let cols = result.shape()[3];
                    let flat = result
                        .into_shape_with_order((n_bands, rows, cols))
                        .context("Failed to reshape 4D result to 3D for writing")?;

                    // Trim overlap and adjust window
                    let overlap = raster_block_data.overlap;
                    let trimmed = crate::array_ops::trimm_array3_asymmetric(&flat, &overlap);
                    let trimmed_rows = trimmed.shape()[1];
                    let trimmed_cols = trimmed.shape()[2];

                    let write_window = ReadWindow {
                        offset: Offset {
                            rows: raster_block_data.read_window.offset.rows + overlap.top as isize,
                            cols: raster_block_data.read_window.offset.cols + overlap.left as isize,
                        },
                        size: Size {
                            rows: trimmed_rows as isize,
                            cols: trimmed_cols as isize,
                        },
                    };

                    parallel_writer::write_block(&writer, trimmed, write_window)?;

                    Ok(())
                })
                .collect::<anyhow::Result<Vec<_>>>()?;

            Ok(())
        })
    }

    /// Applies a worker function to pairs of blocks and outputs a Cloud Optimized GeoTIFF.
    ///
    /// COG variant of [`apply_with_mask`](Self::apply_with_mask).
    pub fn apply_with_mask_cog<U, V>(
        &self,
        other: &RasterDataset<U>,
        worker: fn(&RasterDataBlock<R>, &RasterDataBlock<U>) -> anyhow::Result<Array4<V>>,
        n_cpus: usize,
        out_file: &Path,
        config: &crate::types::OutputConfig,
    ) -> anyhow::Result<()>
    where
        U: RasterType,
        V: RasterType,
    {
        use crate::types::OutputFormat;

        match config.format {
            OutputFormat::GeoTiff => {
                self.apply_with_mask::<U, V>(other, worker, n_cpus, out_file)
            }
            OutputFormat::GeoTiffOverviews => {
                self.apply_with_mask::<U, V>(other, worker, n_cpus, out_file)?;

                let writer = ParallelGeoTiffWriter::new(
                    out_file.to_path_buf(),
                    self.metadata.geo_transform,
                    self.metadata.epsg_code,
                    self.metadata.shape.cols,
                    self.metadata.shape.rows,
                    self.metadata.shape.times * self.metadata.shape.layers,
                );
                writer.build_overviews(&config.overview_resampling, &config.overview_levels)?;
                Ok(())
            }
            OutputFormat::COG => {
                let intermediate = PathBuf::from(create_temp_file("tif"));

                self.apply_with_mask::<U, V>(other, worker, n_cpus, &intermediate)?;

                crate::gdal_utils::translate_to_cog(
                    &intermediate, out_file,
                    &config.compression,
                    &config.overview_resampling,
                )?;
                fs::remove_file(&intermediate).ok();

                Ok(())
            }
        }
    }

    // ─── Mosaic ───

    /// Creates a mosaic from the raster blocks.
    pub fn apply_mosaic<T>(&self, n_cpus: usize)
    where
        T: RasterType,
    {
        unsafe { std::env::set_var("GDAL_NUM_THREADS", "16") };
        let pool = create_rayon_pool(n_cpus);
        let tmp_file = PathBuf::from(create_temp_file("vrt"));
        let handle = pool.install(|| {
            par_tqdm!(self.blocks.to_owned().into_par_iter()).map(
                |raster_block: RasterBlock<R>| -> Vec<PathBuf> {
                    let bid = raster_block.block_index;
                    let file_stem = file_stem_str(&tmp_file);

                    let block_data = self.read_block::<T>(bid);

                    raster_block.write_time_step_blocks(
                        &block_data,
                        &tmp_file,
                        file_stem,
                        bid,
                    )
                },
            )
        });
        let collected: Vec<Vec<PathBuf>> = handle.collect();

        let out_file = PathBuf::from("mosaic.tif");
        let n_times = self.metadata.shape.times;
        log::info!("Saving mosaic to {:?}", out_file);
        mosaic_translate_cleanup_time_steps(
            &collected,
            &out_file,
            self.metadata.epsg_code,
            n_times,
        );
    }

    /// Deprecated: use [`apply_mosaic`](Self::apply_mosaic) instead.
    #[deprecated(since = "0.3.2", note = "Use apply_mosaic() instead. mosaic() will be removed in a future release.")]
    #[allow(deprecated)]
    pub fn mosaic<T>(&self, n_cpus: usize)
    where
        T: RasterType,
    {
        self.apply_mosaic::<T>(n_cpus);
    }

    // ─── Reduce ───

    /// Applies a worker function that reduces the data dimensionality over each block.
    ///
    /// The worker receives `&RasterDataBlock<R>` which includes `layer_indices`,
    /// `date_indices`, and other metadata. This enables semantic selection via the
    /// [`Select`](crate::selection::Select) trait.
    ///
    /// Common uses: mean over time (timesteps → 1), mean over layers (layers → 1), or
    /// any other reduction operation. The result is written directly to a GeoTIFF
    /// via a parallel writer — no intermediate files.
    ///
    /// # Example
    ///
    /// ```rust
    /// extern crate eorst;
    /// use std::path::PathBuf;
    /// use eorst::{
    ///     types::{BlockSize, Dimension}, DataSourceBuilder, RasterDatasetBuilder, RasterDataBlock,
    ///     RasterDataset, Select,
    /// };
    /// use ndarray::{Array3, Axis};
    ///
    /// fn mean_over_time(block: &RasterDataBlock<i16>, dim: Dimension) -> Array3<i16> {
    ///     // Reduce along a dimension, returning Array3 (e.g. mean over time)
    ///     let as_f32: ndarray::Array4<f32> = block.data.mapv(|v| v as f32);
    ///     let mean = as_f32.mean_axis(Axis(0)).unwrap();
    ///     mean.mapv(|v| v as i16)
    /// }
    ///
    /// let manifest_dir = env!("CARGO_MANIFEST_DIR");
    /// let data_source = DataSourceBuilder::from_file(
    ///     &PathBuf::from(manifest_dir).join("data/cemsre_t55jfm_20200614_sub_abam5.tif"),
    /// )
    /// .bands(vec![3, 4])
    /// .set_names(vec!["red", "nir"])
    /// .build();
    ///
    /// let rds: RasterDataset<i16> = RasterDatasetBuilder::from_source(&data_source)
    ///     .block_size(BlockSize { cols: 256, rows: 256 })
    ///     .build();
    ///
    /// let out_path = std::env::temp_dir().join("apply_reduction_mean.tif");
    /// rds.apply_reduction::<i16>(mean_over_time, Dimension::Layer, 1, &out_path, i16::MIN);
    /// ```
    pub fn apply_reduction<U>(
        &self,
        worker: fn(&RasterDataBlock<R>, UtilsDimension) -> Array3<U>,
        dimension: UtilsDimension,
        n_cpus: usize,
        out_file: &Path,
        na_value: U,
    ) where
        U: RasterType,
    {
        let ext = out_file.extension().and_then(std::ffi::OsStr::to_str).unwrap();

        let pool = create_rayon_pool(n_cpus);

        let tmp_file = if ext != "vrt" {
            PathBuf::from(create_temp_file("vrt"))
        } else {
            out_file.to_path_buf()
        };

        let handle = pool.install(|| {
            par_tqdm!(self.blocks
                .to_owned()
                .into_par_iter())
                .map(|raster_block: RasterBlock<R>| -> PathBuf {
                    let id = raster_block.block_index;
                    let file_stem = file_stem_str(&tmp_file);
                    let block_fn = tmp_file.with_file_name(format!("{}_{}.tif", file_stem, id));
                    let block_data = self.read_block::<R>(id);
                    let rdb = RasterDataBlock {
                        data: block_data,
                        metadata: self.metadata.clone(),
                        no_data: NumCast::from(0).unwrap(),
                    };
                    let result = worker(&rdb, dimension);
                    self.write_window3(id, result, &block_fn, na_value);
                    block_fn
                })
        });

        let collected: Vec<PathBuf> = handle.collect();

        mosaic(&collected, &tmp_file, self.metadata.epsg_code, None, None).expect("Could not mosaic to vrt");
        if ext != "vrt" {
            mosaic_translate_cleanup(&collected, &tmp_file, out_file, self.metadata.epsg_code);
        } else {
            fs::rename(tmp_file, out_file).unwrap();
        }
    }

    /// Deprecated: use [`apply_reduction`](Self::apply_reduction) instead.
    #[deprecated(since = "0.3.2", note = "Use apply_reduction() instead. reduce() will be removed in a future release.")]
    #[allow(deprecated)]
    pub fn reduce<U>(
        &self,
        worker: fn(&Array4<R>, UtilsDimension) -> Array3<U>,
        dimension: UtilsDimension,
        n_cpus: usize,
        out_file: &Path,
        na_value: U,
    ) where
        U: RasterType,
    {
        let ext = out_file.extension().and_then(std::ffi::OsStr::to_str).unwrap();

        let pool = create_rayon_pool(n_cpus);

        let tmp_file = if ext != "vrt" {
            PathBuf::from(create_temp_file("vrt"))
        } else {
            out_file.to_path_buf()
        };

        let handle = pool.install(|| {
            self.blocks
                .to_owned()
                .into_par_iter()
                .map(|raster_block: RasterBlock<R>| -> PathBuf {
                    let id = raster_block.block_index;
                    let file_stem = file_stem_str(&tmp_file);
                    let block_fn = tmp_file.with_file_name(format!("{}_{}.tif", file_stem, id));
                    let block_data = self.read_block::<R>(id);
                    let result = worker(&block_data, dimension);
                    self.write_window3(id, result, &block_fn, na_value);
                    block_fn
                })
        });

        let collected: Vec<PathBuf> = handle.collect();

        mosaic(&collected, &tmp_file, self.metadata.epsg_code, None, None).expect("Could not mosaic to vrt");
        if ext != "vrt" {
            mosaic_translate_cleanup(&collected, &tmp_file, out_file, self.metadata.epsg_code);
        } else {
            fs::rename(tmp_file, out_file).unwrap();
        }
    }

    /// Reduces the raster dataset with another dataset as a mask.
    ///
    /// The worker receives `&RasterDataBlock<R>` and `&RasterDataBlock<U>` which include
    /// `layer_indices`, `date_indices`, and other metadata.
    ///
    /// This version uses a parallel writer that writes blocks directly to the output GeoTIFF,
    /// eliminating intermediate files and the subprocess-based mosaic/translate phase.
    pub fn apply_reduction_with_mask<U, V>(
        &self,
        other: &RasterDataset<U>,
        worker: fn(&RasterDataBlock<R>, &RasterDataBlock<U>, UtilsDimension) -> Array3<V>,
        dimension: UtilsDimension,
        n_cpus: usize,
        out_file: &Path,
        na_value: V,
    ) where
        V: RasterType,
        U: RasterType,
    {
        log::info!("Starting apply_reduction_with_mask. Using {:?} cpus.", n_cpus);
        let pool = create_rayon_pool(n_cpus);

        let total_rows = self.metadata.shape.rows;
        let total_cols = self.metadata.shape.cols;
        let epsg_code = self.metadata.epsg_code;
        let geo_transform = self.metadata.geo_transform;
        let n_bands: usize = 1; // reduction always produces 1 output band

        pool.install(|| -> anyhow::Result<()> {
            // Phase 1: Pre-create the output GeoTIFF
            parallel_writer::create_output_geotiff::<V>(
                out_file,
                &geo_transform,
                epsg_code,
                total_cols,
                total_rows,
                n_bands,
                na_value,
            )?;

            let writer = ParallelGeoTiffWriter::new(
                out_file.to_path_buf(),
                geo_transform,
                epsg_code,
                total_cols,
                total_rows,
                n_bands,
            );

            // Phase 2: Process block pairs in parallel, writing directly to output
            self.blocks
                .to_owned()
                .into_par_iter()
                .zip(other.blocks.to_owned().into_par_iter())
                .map(|(raster_block_data, _raster_block_mask)| -> anyhow::Result<()> {
                    let id = raster_block_data.block_index;
                    let block_data = self.read_block::<R>(id);
                    let mask = other.read_block::<U>(id);
                    let rdb_data = RasterDataBlock {
                        data: block_data,
                        metadata: self.metadata.clone(),
                        no_data: NumCast::from(0).unwrap(),
                    };
                    let rdb_mask = RasterDataBlock {
                        data: mask,
                        metadata: other.metadata.clone(),
                        no_data: NumCast::from(0).unwrap(),
                    };
                    let result = worker(&rdb_data, &rdb_mask, dimension);

                    // Trim overlap and write at the block's output window
                    let overlap_sz = raster_block_data.overlap_size;
                    let trimmed = crate::array_ops::trimm_array3(&result, overlap_sz);
                    let trimmed_rows = trimmed.shape()[1];
                    let trimmed_cols = trimmed.shape()[2];

                    let write_window = ReadWindow {
                        offset: Offset {
                            rows: raster_block_data.read_window.offset.rows + overlap_sz as isize,
                            cols: raster_block_data.read_window.offset.cols + overlap_sz as isize,
                        },
                        size: Size {
                            rows: trimmed_rows as isize,
                            cols: trimmed_cols as isize,
                        },
                    };

                    parallel_writer::write_block(&writer, trimmed, write_window)?;

                    Ok(())
                })
                .collect::<anyhow::Result<Vec<_>>>()?;

            Ok(())
        })
        .expect("Parallel reduction failed");
    }

    /// Reduces the raster dataset with another dataset as a mask, outputting a Cloud Optimized GeoTIFF.
    ///
    /// COG variant of [`apply_reduction_with_mask`](Self::apply_reduction_with_mask).
    pub fn apply_reduction_with_mask_cog<U, V>(
        &self,
        other: &RasterDataset<U>,
        worker: fn(&RasterDataBlock<R>, &RasterDataBlock<U>, UtilsDimension) -> Array3<V>,
        dimension: UtilsDimension,
        n_cpus: usize,
        out_file: &Path,
        na_value: V,
        config: &crate::types::OutputConfig,
    ) where
        V: RasterType,
        U: RasterType,
    {
        use crate::types::OutputFormat;

        match config.format {
            OutputFormat::GeoTiff => {
                self.apply_reduction_with_mask::<U, V>(other, worker, dimension, n_cpus, out_file, na_value)
            }
            OutputFormat::GeoTiffOverviews => {
                self.apply_reduction_with_mask::<U, V>(other, worker, dimension, n_cpus, out_file, na_value);

                let writer = ParallelGeoTiffWriter::new(
                    out_file.to_path_buf(),
                    self.metadata.geo_transform,
                    self.metadata.epsg_code,
                    self.metadata.shape.cols,
                    self.metadata.shape.rows,
                    1,
                );
                writer.build_overviews(&config.overview_resampling, &config.overview_levels).ok();
            }
            OutputFormat::COG => {
                let intermediate = PathBuf::from(create_temp_file("tif"));

                self.apply_reduction_with_mask::<U, V>(other, worker, dimension, n_cpus, &intermediate, na_value);

                crate::gdal_utils::translate_to_cog(
                    &intermediate, out_file,
                    &config.compression,
                    &config.overview_resampling,
                ).ok();
                fs::remove_file(&intermediate).ok();
            }
        }
    }

    /// Deprecated: use [`apply_reduction_with_mask`](Self::apply_reduction_with_mask) instead.
    #[deprecated(since = "0.3.2", note = "Use apply_reduction_with_mask() instead. reduce_with_mask() will be removed in a future release.")]
    #[allow(deprecated)]
    pub fn reduce_with_mask<U, V>(
        &self,
        other: &RasterDataset<U>,
        worker: fn(&Array4<R>, &Array4<U>, UtilsDimension) -> Array3<V>,
        dimension: UtilsDimension,
        n_cpus: usize,
        out_file: &Path,
        na_value: V,
    ) where
        V: RasterType,
        U: RasterType,
    {
        log::info!("Starting reduce with mask. Using {:?} cpus.", n_cpus);
        let pool = create_rayon_pool(n_cpus);

        let tmp_file = PathBuf::from(create_temp_file("vrt"));

        let handle = pool.install(|| {
            self.blocks
                .to_owned()
                .into_par_iter()
                .zip(other.blocks.to_owned().into_par_iter())
                .map(|(raster_block_data, _raster_block_mask)| -> PathBuf {
                    let id = raster_block_data.block_index;
                    let file_stem = file_stem_str(&tmp_file);
                    let block_fn = tmp_file.with_file_name(format!("{}_{}.tif", file_stem, id));
                    let block_data = self.read_block::<R>(id);
                    let mask = other.read_block::<U>(id);
                    let result = worker(&block_data, &mask, dimension);
                    self.write_window3(id, result, &block_fn, na_value);
                    block_fn
                })
        });

        let collected: Vec<PathBuf> = handle.collect();
        mosaic_translate_cleanup(&collected, &tmp_file, out_file, self.metadata.epsg_code);
    }

    /// Applies a worker function that reduces rows to pixel values.
    ///
    /// Note: this method reads data as type `T` which may differ from the dataset's
    /// stored type `R`, so the worker receives `&Array4<T>` rather than `&RasterDataBlock<T>`.
    pub fn apply_reduction_row_pixel<T>(
        &self,
        worker: fn(&Array4<T>, UtilsDimension) -> Array2<T>,
        na_value: T,
        dimension: UtilsDimension,
        n_cpus: usize,
        out_file: &Path,
    ) where
        T: RasterType,
    {
        let pool = create_rayon_pool(n_cpus);
        let tmp_file = PathBuf::from(create_temp_file("vrt"));

        let handle = pool.install(|| {
            self.blocks
                .to_owned()
                .into_par_iter()
                .map(|raster_block: RasterBlock<R>| -> PathBuf {
                    let id = raster_block.block_index;
                    let file_stem = file_stem_str(&tmp_file);
                    let block_fn = tmp_file.with_file_name(format!("{}_{}.tif", file_stem, id));
                    let block_data = self.read_block::<T>(id);
                    let result = worker(&block_data, dimension);
                    raster_block.write_samples_feature(&result, &block_fn, na_value);
                    block_fn
                })
        });

        let collected: Vec<PathBuf> = handle.collect();
        mosaic_translate_cleanup(&collected, &tmp_file, out_file, self.metadata.epsg_code);
    }

    /// Deprecated: use [`apply_reduction_row_pixel`](Self::apply_reduction_row_pixel) instead.
    #[deprecated(since = "0.3.2", note = "Use apply_reduction_row_pixel() instead. reduce_row_pixel() will be removed in a future release.")]
    #[allow(deprecated)]
    pub fn reduce_row_pixel<T>(
        &self,
        worker: fn(&Array4<T>, UtilsDimension) -> Array2<T>,
        na_value: T,
        dimension: UtilsDimension,
        n_cpus: usize,
        out_file: &Path,
    ) where
        T: RasterType,
    {
        self.apply_reduction_row_pixel(worker, na_value, dimension, n_cpus, out_file);
    }

    /// Reduces rows with a mask dataset.
    ///
    /// Note: this method reads data as types `R` and `V` matching the datasets,
    /// so the worker receives raw `&Array4` references.
    pub fn apply_reduction_row_pixel_with_mask<V, U>(
        &self,
        other: &RasterDataset<V>,
        worker: fn(&Array4<R>, &Array4<V>, UtilsDimension) -> Array2<U>,
        na_value: U,
        dimension: UtilsDimension,
        n_cpus: usize,
        out_file: &Path,
    ) where
        U: RasterType,
        V: RasterType,
    {
        let pool = create_rayon_pool(n_cpus);

        let tmp_file = PathBuf::from(create_temp_file("vrt"));

        let handle = pool.install(|| {
            self.blocks
                .to_owned()
                .into_par_iter()
                .zip(other.blocks.to_owned().into_par_iter())
                .map(|(raster_block_data, _raster_block_mask)| -> PathBuf {
                    let id = raster_block_data.block_index;
                    let file_stem = file_stem_str(&tmp_file);
                    let block_fn = tmp_file.with_file_name(format!(
                        "{}_{}.tif",
                        file_stem, raster_block_data.block_index
                    ));
                    let block_data = self.read_block::<R>(id);
                    let mask = other.read_block::<V>(id);
                    let result = worker(&block_data, &mask, dimension);
                    raster_block_data.write_samples_feature(&result, &block_fn, na_value);
                    block_fn
                })
        });

        let collected: Vec<PathBuf> = handle.collect();
        mosaic_translate_cleanup(&collected, &tmp_file, out_file, self.metadata.epsg_code);
    }

    /// Deprecated: use [`apply_reduction_row_pixel_with_mask`](Self::apply_reduction_row_pixel_with_mask) instead.
    #[deprecated(since = "0.3.2", note = "Use apply_reduction_row_pixel_with_mask() instead. reduce_row_pixel_with_mask() will be removed in a future release.")]
    #[allow(deprecated)]
    pub fn reduce_row_pixel_with_mask<V, U>(
        &self,
        other: &RasterDataset<V>,
        worker: fn(&Array4<R>, &Array4<V>, UtilsDimension) -> Array2<U>,
        na_value: U,
        dimension: UtilsDimension,
        n_cpus: usize,
        out_file: &Path,
    ) where
        U: RasterType,
        V: RasterType,
    {
        self.apply_reduction_row_pixel_with_mask(other, worker, na_value, dimension, n_cpus, out_file);
    }
}