mappers_warp 0.1.0-alpha.2

use ndarray::{Array2, ArrayView2, FoldWhile, Zip, s};

use crate::{Warper, WarperError};

impl Warper {
    #[inline]
    fn validate_source_raster_shape(
        &self,
        source_raster: &ArrayView2<f64>,
    ) -> Result<(), WarperError> {
        if source_raster.dim().0 != self.source_shape.0
            || source_raster.dim().1 != self.source_shape.1
        {
            return Err(WarperError::InvalidRasterDimensions);
        }

        Ok(())
    }

    /// This variant does not check anything at all
    pub fn warp_unchecked<'a, A: Into<ArrayView2<'a, f64>>>(
        &self,
        source_raster: A,
    ) -> Array2<f64> {
        let source_raster: ArrayView2<f64> = source_raster.into();

        self.internals.map(|intr| {
            let values = source_raster.slice(s![
                (intr.anchor_idx.1 - 1)..(intr.anchor_idx.1 + 3),
                (intr.anchor_idx.0 - 1)..(intr.anchor_idx.0 + 3)
            ]);

            let mut weight_accum = 0.0;
            let mut result_accum = 0.0;

            for j in 0..4 {
                let mut inner_weight_accum = 0.0;
                let mut inner_result_accum = 0.0;

                for i in 0..4 {
                    let value = values[[j, i]];
                    let x_weight = intr.x_weights[i];

                    inner_weight_accum += x_weight;
                    inner_result_accum += x_weight * value;
                }

                let y_weight = intr.y_weights[j];

                weight_accum += inner_weight_accum * y_weight;
                result_accum += inner_result_accum * y_weight;
            }

            result_accum / weight_accum
        })
    }

    /// From `GdalWarp` documentation: for bilinear, cubic, cubicspline and lanczos, for each target pixel, the coordinate of its center
    /// is projected back to source coordinates and a corresponding source pixel is identified. If this source pixel is invalid,
    /// the target pixel is considered as nodata. Given that those resampling kernels have a non-null kernel radius,
    /// this source pixel is just one among other several source pixels, and it might be possible that there are invalid
    /// values in those other contributing source pixels. The weights used to take into account those invalid values
    /// will be set to zero to ignore them.
    ///
    /// In short: this variant computes the filter value as if NaN points weren't there
    /// or if the result is not finite
    pub fn warp_ignore_nodata<'a, A: Into<ArrayView2<'a, f64>>>(
        &self,
        source_raster: A,
    ) -> Result<Array2<f64>, WarperError> {
        let source_raster: ArrayView2<f64> = source_raster.into();

        self.validate_source_raster_shape(&source_raster)?;

        let mut target_raster = Array2::from_elem(self.internals.raw_dim(), f64::NEG_INFINITY);

        Zip::from(&mut target_raster)
            .and(&self.internals)
            .fold_while(Ok(()), |_, v, intr| {
                let values = source_raster.slice(s![
                    (intr.anchor_idx.1 - 1)..(intr.anchor_idx.1 + 3),
                    (intr.anchor_idx.0 - 1)..(intr.anchor_idx.0 + 3)
                ]);

                let mut weight_accum = 0.0;
                let mut result_accum = 0.0;

                for j in 0..4 {
                    let mut inner_weight_accum = 0.0;
                    let mut inner_result_accum = 0.0;

                    for i in 0..4 {
                        let value = values[[j, i]];

                        if !value.is_nan() {
                            let x_weight = intr.x_weights[i];
                            inner_weight_accum += x_weight;
                            inner_result_accum += x_weight * value;
                        }
                    }

                    let y_weight = intr.y_weights[j];

                    weight_accum += inner_weight_accum * y_weight;
                    result_accum += inner_result_accum * y_weight;
                }

                if (weight_accum).abs() < f64::EPSILON {
                    *v = f64::NAN;
                    return FoldWhile::Continue(Ok(()));
                }

                let result = result_accum / weight_accum;

                if result.is_finite() {
                    *v = result;
                    FoldWhile::Continue(Ok(()))
                } else {
                    FoldWhile::Done(Err(WarperError::WarpingError))
                }
            })
            .into_inner()?;

        Ok(target_raster)
    }

    /// This variant throws an error if there's any NaN in the (used) input data or result is not finite
    ///
    /// Note: the input NaN check checks only values that are actually used in computations,
    /// so you can pass an input raster with NaNs and get `Ok` result as long as those values
    /// do not appear in any computation kernel.
    pub fn warp_reject_nodata<'a, A: Into<ArrayView2<'a, f64>>>(
        &self,
        source_raster: A,
    ) -> Result<Array2<f64>, WarperError> {
        let source_raster: ArrayView2<f64> = source_raster.into();

        self.validate_source_raster_shape(&source_raster)?;

        let mut target_raster = Array2::from_elem(self.internals.raw_dim(), f64::NEG_INFINITY);

        Zip::from(&mut target_raster)
            .and(&self.internals)
            .fold_while(Ok(()), |_, v, intr| {
                let values = source_raster.slice(s![
                    (intr.anchor_idx.1 - 1)..(intr.anchor_idx.1 + 3),
                    (intr.anchor_idx.0 - 1)..(intr.anchor_idx.0 + 3)
                ]);

                let mut weight_accum = 0.0;
                let mut result_accum = 0.0;

                for j in 0..4 {
                    let mut inner_weight_accum = 0.0;
                    let mut inner_result_accum = 0.0;

                    for i in 0..4 {
                        let value = values[[j, i]];

                        if value.is_nan() {
                            return FoldWhile::Done(Err(WarperError::NanError));
                        }
                        let x_weight = intr.x_weights[i];
                        inner_weight_accum += x_weight;
                        inner_result_accum += x_weight * value;
                    }

                    let y_weight = intr.y_weights[j];

                    weight_accum += inner_weight_accum * y_weight;
                    result_accum += inner_result_accum * y_weight;
                }

                let result = result_accum / weight_accum;

                if result.is_finite() {
                    *v = result;
                    FoldWhile::Continue(Ok(()))
                } else {
                    FoldWhile::Done(Err(WarperError::WarpingError))
                }
            })
            .into_inner()?;

        Ok(target_raster)
    }

    /// This variant returns NaN for each target pixel which would include NaN in the filter
    /// or if result is not finite
    pub fn warp_discard_nodata<'a, A: Into<ArrayView2<'a, f64>>>(
        &self,
        source_raster: A,
    ) -> Result<Array2<f64>, WarperError> {
        let source_raster: ArrayView2<f64> = source_raster.into();

        self.validate_source_raster_shape(&source_raster)?;

        let mut target_raster = Array2::from_elem(self.internals.raw_dim(), f64::NEG_INFINITY);

        Zip::from(&mut target_raster)
            .and(&self.internals)
            .fold_while(Ok(()), |_, v, intr| {
                let values = source_raster.slice(s![
                    (intr.anchor_idx.1 - 1)..(intr.anchor_idx.1 + 3),
                    (intr.anchor_idx.0 - 1)..(intr.anchor_idx.0 + 3)
                ]);

                let mut weight_accum = 0.0;
                let mut result_accum = 0.0;

                for j in 0..4 {
                    let mut inner_weight_accum = 0.0;
                    let mut inner_result_accum = 0.0;

                    for i in 0..4 {
                        let value = values[[j, i]];

                        if value.is_nan() {
                            *v = f64::NAN;
                            return FoldWhile::Continue(Ok(()));
                        }
                        let x_weight = intr.x_weights[i];
                        inner_weight_accum += x_weight;
                        inner_result_accum += x_weight * value;
                    }

                    let y_weight = intr.y_weights[j];

                    weight_accum += inner_weight_accum * y_weight;
                    result_accum += inner_result_accum * y_weight;
                }

                let result = result_accum / weight_accum;

                if result.is_finite() {
                    *v = result;
                    FoldWhile::Continue(Ok(()))
                } else {
                    FoldWhile::Done(Err(WarperError::WarpingError))
                }
            })
            .into_inner()?;

        Ok(target_raster)
    }
}

/// Parallel equivalent of `warp_unchecked`.
#[cfg(feature = "multithreading")]
#[cfg_attr(docsrs, doc(cfg(feature = "multithreading")))]
impl Warper {
    pub fn warp_unchecked_parallel<'a, A: Into<ArrayView2<'a, f64>>>(
        &self,
        source_raster: A,
    ) -> Array2<f64> {
        let source_raster: ArrayView2<f64> = source_raster.into();

        Zip::from(&self.internals).par_map_collect(|intr| {
            let values = source_raster.slice(s![
                (intr.anchor_idx.1 - 1)..(intr.anchor_idx.1 + 3),
                (intr.anchor_idx.0 - 1)..(intr.anchor_idx.0 + 3)
            ]);

            let mut weight_accum = 0.0;
            let mut result_accum = 0.0;

            for j in 0..4 {
                let mut inner_weight_accum = 0.0;
                let mut inner_result_accum = 0.0;

                for i in 0..4 {
                    let value = values[[j, i]];
                    let x_weight = intr.x_weights[i];

                    inner_weight_accum += x_weight;
                    inner_result_accum += x_weight * value;
                }

                let y_weight = intr.y_weights[j];

                weight_accum += inner_weight_accum * y_weight;
                result_accum += inner_result_accum * y_weight;
            }

            result_accum / weight_accum
        })
    }

    /// This implementation catches warp operation producing NaNs (that is nans resulting from computation error
    /// not those resulting from nodata)
    #[cfg(feature = "multithreading")]
    #[cfg_attr(docsrs, doc(cfg(feature = "multithreading")))]
    pub fn warp_ignore_nodata_parallel<'a, A: Into<ArrayView2<'a, f64>>>(
        &self,
        source_raster: A,
    ) -> Result<Array2<f64>, WarperError> {
        use ndarray::parallel::prelude::{IntoParallelIterator, ParallelIterator};

        let source_raster: ArrayView2<f64> = source_raster.into();

        self.validate_source_raster_shape(&source_raster)?;

        let mut target_raster = Array2::from_elem(self.internals.raw_dim(), f64::NEG_INFINITY);

        Zip::from(&mut target_raster)
            .and(&self.internals)
            .into_par_iter()
            .try_for_each(|(v, intr)| {
                let values = source_raster.slice(s![
                    (intr.anchor_idx.1 - 1)..(intr.anchor_idx.1 + 3),
                    (intr.anchor_idx.0 - 1)..(intr.anchor_idx.0 + 3)
                ]);

                let mut weight_accum = 0.0;
                let mut result_accum = 0.0;

                for j in 0..4 {
                    let mut inner_weight_accum = 0.0;
                    let mut inner_result_accum = 0.0;

                    for i in 0..4 {
                        let value = values[[j, i]];

                        if !value.is_nan() {
                            let x_weight = intr.x_weights[i];
                            inner_weight_accum += x_weight;
                            inner_result_accum += x_weight * value;
                        }
                    }

                    let y_weight = intr.y_weights[j];

                    weight_accum += inner_weight_accum * y_weight;
                    result_accum += inner_result_accum * y_weight;
                }

                if (weight_accum).abs() < f64::EPSILON {
                    *v = f64::NAN;
                    return Ok(());
                }

                let result = result_accum / weight_accum;

                if result.is_finite() {
                    *v = result;
                    Ok(())
                } else {
                    Err(WarperError::WarpingError)
                }
            })?;

        Ok(target_raster)
    }

    /// According to benchmarks, this is the fastest checked parallel variant
    #[cfg(feature = "multithreading")]
    #[cfg_attr(docsrs, doc(cfg(feature = "multithreading")))]
    pub fn warp_reject_nodata_parallel<'a, A: Into<ArrayView2<'a, f64>>>(
        &self,
        source_raster: A,
    ) -> Result<Array2<f64>, WarperError> {
        use ndarray::parallel::prelude::{IntoParallelIterator, ParallelIterator};

        let source_raster: ArrayView2<f64> = source_raster.into();

        self.validate_source_raster_shape(&source_raster)?;

        Zip::from(&source_raster)
            .into_par_iter()
            .try_for_each(|(v,)| {
                if v.is_nan() {
                    Err(WarperError::NanError)
                } else {
                    Ok(())
                }
            })?;

        let target_raster = self.warp_unchecked_parallel(source_raster);

        Zip::from(&target_raster)
            .into_par_iter()
            .try_for_each(|(v,)| {
                if v.is_finite() {
                    Ok(())
                } else {
                    Err(WarperError::WarpingError)
                }
            })?;

        Ok(target_raster)
    }

    /// This implementation catches warp operation producing NaNs (that is nans resulting from computation error
    /// not those resulting from nodata)
    #[cfg(feature = "multithreading")]
    #[cfg_attr(docsrs, doc(cfg(feature = "multithreading")))]
    pub fn warp_discard_nodata_parallel<'a, A: Into<ArrayView2<'a, f64>>>(
        &self,
        source_raster: A,
    ) -> Result<Array2<f64>, WarperError> {
        use ndarray::parallel::prelude::{IntoParallelIterator, ParallelIterator};

        let source_raster: ArrayView2<f64> = source_raster.into();

        self.validate_source_raster_shape(&source_raster)?;

        let mut target_raster = Array2::from_elem(self.internals.raw_dim(), f64::NEG_INFINITY);

        Zip::from(&mut target_raster)
            .and(&self.internals)
            .into_par_iter()
            .try_for_each(|(v, intr)| {
                let values = source_raster.slice(s![
                    (intr.anchor_idx.1 - 1)..(intr.anchor_idx.1 + 3),
                    (intr.anchor_idx.0 - 1)..(intr.anchor_idx.0 + 3)
                ]);

                let mut weight_accum = 0.0;
                let mut result_accum = 0.0;

                for j in 0..4 {
                    let mut inner_weight_accum = 0.0;
                    let mut inner_result_accum = 0.0;

                    for i in 0..4 {
                        let value = values[[j, i]];

                        if value.is_nan() {
                            *v = f64::NAN;
                            return Ok(());
                        }
                        let x_weight = intr.x_weights[i];
                        inner_weight_accum += x_weight;
                        inner_result_accum += x_weight * value;
                    }

                    let y_weight = intr.y_weights[j];

                    weight_accum += inner_weight_accum * y_weight;
                    result_accum += inner_result_accum * y_weight;
                }

                let result = result_accum / weight_accum;

                if result.is_finite() {
                    *v = result;
                    Ok(())
                } else {
                    Err(WarperError::WarpingError)
                }
            })?;

        Ok(target_raster)
    }
}