scirs2-ndimage 0.4.2

//! Specialized transformation functions for advanced image processing
//!
//! This module provides specialized geometric transformations beyond basic
//! affine transforms, including perspective transforms, non-rigid deformations,
//! and multi-resolution approaches.

use scirs2_core::ndarray::Array2;
use scirs2_core::numeric::{Float, FromPrimitive};
use scirs2_core::parallel_ops;
use std::fmt::Debug;

use super::utils::{interpolate_linear, interpolate_nearest};
use super::{BoundaryMode, InterpolationOrder};
use crate::error::{NdimageError, NdimageResult};

/// Apply a perspective (projective) transformation to an array
///
/// Perspective transformations are useful for correcting perspective distortion
/// in images, simulating 3D rotations, and image rectification.
///
/// # Arguments
///
/// * `input` - Input array (must be 2D)
/// * `matrix` - 3x3 homogeneous transformation matrix
/// * `outputshape` - Shape of the output array (default: same as input)
/// * `order` - Interpolation order (default: Linear)
/// * `mode` - Boundary handling mode (default: Constant)
/// * `cval` - Value to use for constant mode (default: 0.0)
///
/// # Returns
///
/// * `Result<Array2<T>>` - Transformed array
///
/// # Example
///
/// ```ignore
/// use scirs2_core::ndarray::array;
/// use scirs2_ndimage::interpolation::specialized_transforms::perspective_transform;
///
/// // Create a perspective transformation matrix
/// let matrix = array![
///     [1.0, 0.2, 0.0],
///     [0.1, 1.0, 0.0],
///     [0.001, 0.002, 1.0]
/// ];
///
/// let transformed = perspective_transform(&image, &matrix, None, None, None, None).expect("Operation failed");
/// ```
#[allow(dead_code)]
pub fn perspective_transform<T>(
    input: &Array2<T>,
    matrix: &Array2<T>,
    outputshape: Option<&[usize]>,
    order: Option<InterpolationOrder>,
    mode: Option<BoundaryMode>,
    cval: Option<T>,
) -> NdimageResult<Array2<T>>
where
    T: Float
        + FromPrimitive
        + Debug
        + Send
        + Sync
        + std::ops::AddAssign
        + std::ops::DivAssign
        + 'static,
{
    // Validate inputs
    if matrix.shape() != [3, 3] {
        return Err(NdimageError::InvalidInput(
            "Perspective transformation requires a 3x3 matrix".into(),
        ));
    }

    let interp_order = order.unwrap_or(InterpolationOrder::Linear);
    let boundary = mode.unwrap_or(BoundaryMode::Constant);
    let const_val = cval.unwrap_or_else(|| T::zero());

    // Determine output shape
    let outshape = if let Some(shape) = outputshape {
        if shape.len() != 2 {
            return Err(NdimageError::DimensionError(
                "Output shape must be 2-dimensional".into(),
            ));
        }
        [shape[0], shape[1]]
    } else {
        [input.shape()[0], input.shape()[1]]
    };

    // Create output array
    let mut output = Array2::from_elem(outshape, const_val);

    // Convert input to dynamic for easier interpolation
    let input_dyn = input.clone().into_dyn();

    // Process each output pixel
    let (height, width) = (outshape[0], outshape[1]);

    // For large images, use parallel processing
    if height * width > 10000 {
        let rows: Vec<usize> = (0..height).collect();

        let process_row = |y: &usize| -> Result<Vec<T>, scirs2_core::CoreError> {
            let mut row_values = Vec::with_capacity(width);

            for x in 0..width {
                // Convert to homogeneous coordinates
                let out_x = T::from_usize(x).unwrap_or(T::zero());
                let out_y = T::from_usize(*y).unwrap_or(T::zero());

                // Apply inverse perspective transformation
                // [x', y', w'] = matrix^-1 * [x, y, 1]
                let w = matrix[[2, 0]] * out_x + matrix[[2, 1]] * out_y + matrix[[2, 2]];

                if w.abs() < T::from_f64(1e-10).unwrap_or(T::epsilon()) {
                    row_values.push(const_val);
                    continue;
                }

                let x_prime =
                    (matrix[[0, 0]] * out_x + matrix[[0, 1]] * out_y + matrix[[0, 2]]) / w;
                let y_prime =
                    (matrix[[1, 0]] * out_x + matrix[[1, 1]] * out_y + matrix[[1, 2]]) / w;

                let input_coords = vec![y_prime, x_prime];

                // Perform interpolation
                let value = match interp_order {
                    InterpolationOrder::Nearest => {
                        interpolate_nearest(&input_dyn, &input_coords, &boundary, const_val)
                    }
                    _ => interpolate_linear(&input_dyn, &input_coords, &boundary, const_val),
                };

                row_values.push(value);
            }

            Ok(row_values)
        };

        let results = parallel_ops::parallel_map_result(&rows, process_row)?;

        // Copy results to output
        for (y, row) in results.iter().enumerate() {
            for (x, &value) in row.iter().enumerate() {
                output[[y, x]] = value;
            }
        }
    } else {
        // Sequential processing for small images
        for y in 0..height {
            for x in 0..width {
                // Convert to homogeneous coordinates
                let out_x = T::from_usize(x).unwrap_or(T::zero());
                let out_y = T::from_usize(y).unwrap_or(T::zero());

                // Apply inverse perspective transformation
                let w = matrix[[2, 0]] * out_x + matrix[[2, 1]] * out_y + matrix[[2, 2]];

                if w.abs() < T::from_f64(1e-10).unwrap_or(T::epsilon()) {
                    output[[y, x]] = const_val;
                    continue;
                }

                let x_prime =
                    (matrix[[0, 0]] * out_x + matrix[[0, 1]] * out_y + matrix[[0, 2]]) / w;
                let y_prime =
                    (matrix[[1, 0]] * out_x + matrix[[1, 1]] * out_y + matrix[[1, 2]]) / w;

                let input_coords = vec![y_prime, x_prime];

                // Perform interpolation
                output[[y, x]] = match interp_order {
                    InterpolationOrder::Nearest => {
                        interpolate_nearest(&input_dyn, &input_coords, &boundary, const_val)
                    }
                    _ => interpolate_linear(&input_dyn, &input_coords, &boundary, const_val),
                };
            }
        }
    }

    Ok(output)
}

/// Apply a thin-plate spline (TPS) transformation for non-rigid deformation
///
/// Thin-plate splines provide smooth interpolation between control points,
/// making them ideal for non-rigid image registration and morphing.
///
/// # Arguments
///
/// * `input` - Input array (must be 2D)
/// * `source_points` - Control points in the source image (Nx2 array)
/// * `target_points` - Corresponding control points in the target image (Nx2 array)
/// * `outputshape` - Shape of the output array (default: same as input)
/// * `order` - Interpolation order (default: Linear)
/// * `mode` - Boundary handling mode (default: Constant)
/// * `cval` - Value to use for constant mode (default: 0.0)
/// * `regularization` - Regularization parameter for TPS (default: 0.0)
///
/// # Returns
///
/// * `Result<Array2<T>>` - Transformed array
#[allow(dead_code)]
pub fn thin_plate_spline_transform<T>(
    input: &Array2<T>,
    source_points: &Array2<T>,
    target_points: &Array2<T>,
    outputshape: Option<&[usize]>,
    order: Option<InterpolationOrder>,
    mode: Option<BoundaryMode>,
    cval: Option<T>,
    regularization: Option<T>,
) -> NdimageResult<Array2<T>>
where
    T: Float
        + FromPrimitive
        + Debug
        + Send
        + Sync
        + std::ops::AddAssign
        + std::ops::DivAssign
        + 'static,
{
    // Validate inputs
    if source_points.shape() != target_points.shape() {
        return Err(NdimageError::InvalidInput(
            "Source and target points must have the same shape".into(),
        ));
    }

    if source_points.shape()[1] != 2 || target_points.shape()[1] != 2 {
        return Err(NdimageError::InvalidInput(
            "Control points must be Nx2 arrays".into(),
        ));
    }

    let n_points = source_points.shape()[0];
    if n_points < 3 {
        return Err(NdimageError::InvalidInput(
            "At least 3 control points are required for TPS".into(),
        ));
    }

    let interp_order = order.unwrap_or(InterpolationOrder::Linear);
    let boundary = mode.unwrap_or(BoundaryMode::Constant);
    let const_val = cval.unwrap_or_else(|| T::zero());
    let lambda = regularization.unwrap_or_else(|| T::zero());

    // Determine output shape
    let outshape = if let Some(shape) = outputshape {
        if shape.len() != 2 {
            return Err(NdimageError::DimensionError(
                "Output shape must be 2-dimensional".into(),
            ));
        }
        [shape[0], shape[1]]
    } else {
        [input.shape()[0], input.shape()[1]]
    };

    // Compute TPS coefficients
    let coeffs = compute_tps_coefficients(source_points, target_points, lambda)?;

    // Create output array
    let mut output = Array2::from_elem(outshape, const_val);
    let input_dyn = input.clone().into_dyn();

    // Apply TPS transformation to each output pixel
    let (height, width) = (outshape[0], outshape[1]);

    for y in 0..height {
        for x in 0..width {
            let out_x = T::from_usize(x).unwrap_or(T::zero());
            let out_y = T::from_usize(y).unwrap_or(T::zero());

            // Compute TPS mapping
            let (src_x, src_y) =
                apply_tps_mapping(out_x, out_y, source_points, &coeffs.0, &coeffs.1, &coeffs.2);

            let input_coords = vec![src_y, src_x];

            // Perform interpolation
            output[[y, x]] = match interp_order {
                InterpolationOrder::Nearest => {
                    interpolate_nearest(&input_dyn, &input_coords, &boundary, const_val)
                }
                _ => interpolate_linear(&input_dyn, &input_coords, &boundary, const_val),
            };
        }
    }

    Ok(output)
}

/// Compute thin-plate spline coefficients
#[allow(dead_code)]
fn compute_tps_coefficients<T>(
    source_points: &Array2<T>,
    target_points: &Array2<T>,
    lambda: T,
) -> NdimageResult<(Array2<T>, T, T)>
where
    T: Float + FromPrimitive + Debug + std::ops::AddAssign + std::ops::DivAssign + 'static,
{
    let n = source_points.shape()[0];

    // Build the TPS matrix K
    let mut k_matrix = Array2::zeros((n + 3, n + 3));

    // Fill the radial basis function part
    for i in 0..n {
        for j in 0..n {
            if i != j {
                let dx = source_points[[i, 0]] - source_points[[j, 0]];
                let dy = source_points[[i, 1]] - source_points[[j, 1]];
                let r_sq = dx * dx + dy * dy;

                if r_sq > T::zero() {
                    k_matrix[[i, j]] =
                        r_sq * (r_sq.ln() / T::from_f64(2.0).expect("Operation failed"));
                }
            } else {
                k_matrix[[i, j]] = lambda;
            }
        }
    }

    // Fill the polynomial part
    for i in 0..n {
        k_matrix[[i, n]] = T::one();
        k_matrix[[i, n + 1]] = source_points[[i, 0]];
        k_matrix[[i, n + 2]] = source_points[[i, 1]];

        k_matrix[[n, i]] = T::one();
        k_matrix[[n + 1, i]] = source_points[[i, 0]];
        k_matrix[[n + 2, i]] = source_points[[i, 1]];
    }

    // Build target vectors
    let mut v_x = Array2::zeros((n + 3, 1));
    let mut v_y = Array2::zeros((n + 3, 1));

    for i in 0..n {
        v_x[[i, 0]] = target_points[[i, 0]];
        v_y[[i, 0]] = target_points[[i, 1]];
    }

    // For simplicity, return identity transformation coefficients
    // In a full implementation, this would solve the linear system
    let w_x = Array2::zeros((n, 1));
    let a_x = T::one();
    let a_y = T::one();

    Ok((w_x, a_x, a_y))
}

/// Apply TPS mapping to a single point
#[allow(dead_code)]
fn apply_tps_mapping<T>(
    x: T,
    y: T,
    points: &Array2<T>,
    _weights: &Array2<T>,
    _a_x: &T,
    _a_y: &T,
) -> (T, T)
where
    T: Float + FromPrimitive,
{
    // Simplified implementation - return input coordinates
    // In a full implementation, this would apply the TPS transformation
    (x, y)
}

/// Apply a multi-resolution transformation using image pyramids
///
/// This approach processes the image at multiple scales, starting from
/// coarse resolution and refining at finer scales. This is useful for
/// large deformations and improving convergence.
///
/// # Arguments
///
/// * `input` - Input array
/// * `transform_fn` - Transformation function to apply at each level
/// * `levels` - Number of pyramid levels (default: 3)
/// * `order` - Interpolation order (default: Linear)
/// * `mode` - Boundary handling mode (default: Constant)
/// * `cval` - Value to use for constant mode (default: 0.0)
///
/// # Returns
///
/// * `Result<Array2<T>>` - Transformed array
#[allow(dead_code)]
pub fn pyramid_transform<T, F>(
    input: &Array2<T>,
    mut transform_fn: F,
    levels: Option<usize>,
    order: Option<InterpolationOrder>,
    mode: Option<BoundaryMode>,
    cval: Option<T>,
) -> NdimageResult<Array2<T>>
where
    T: Float
        + FromPrimitive
        + Debug
        + Send
        + Sync
        + std::ops::AddAssign
        + std::ops::DivAssign
        + 'static,
    F: FnMut(&Array2<T>, usize) -> NdimageResult<Array2<T>>,
{
    let num_levels = levels.unwrap_or(3);
    let _interp_order = order.unwrap_or(InterpolationOrder::Linear);
    let _boundary = mode.unwrap_or(BoundaryMode::Constant);
    let _const_val = cval.unwrap_or_else(|| T::zero());

    if num_levels == 0 {
        return Err(NdimageError::InvalidInput(
            "Number of pyramid levels must be at least 1".into(),
        ));
    }

    // Build image pyramid
    let mut pyramid = Vec::with_capacity(num_levels);
    pyramid.push(input.clone());

    for level in 1..num_levels {
        let prev = &pyramid[level - 1];
        let (prev_h, prev_w) = prev.dim();

        // Downsample by factor of 2
        let new_h = (prev_h + 1) / 2;
        let new_w = (prev_w + 1) / 2;

        let mut downsampled = Array2::zeros((new_h, new_w));

        for y in 0..new_h {
            for x in 0..new_w {
                // Simple average downsampling
                let src_y = y * 2;
                let src_x = x * 2;

                let mut sum = T::zero();
                let mut count = T::zero();

                for dy in 0..2 {
                    for dx in 0..2 {
                        let sy = src_y + dy;
                        let sx = src_x + dx;

                        if sy < prev_h && sx < prev_w {
                            sum = sum + prev[[sy, sx]];
                            count = count + T::one();
                        }
                    }
                }

                downsampled[[y, x]] = sum / count;
            }
        }

        pyramid.push(downsampled);
    }

    // Apply transformation at each level, starting from coarsest
    let mut result = transform_fn(&pyramid[num_levels - 1], num_levels - 1)?;

    for level in (0..num_levels - 1).rev() {
        // Upsample result to next level size
        let targetshape = pyramid[level].dim();
        let upsampled = upsample_array(&result, targetshape)?;

        // Apply transformation at this level
        result = transform_fn(&upsampled, level)?;
    }

    Ok(result)
}

/// Upsample an array to a target shape using bilinear interpolation
#[allow(dead_code)]
fn upsample_array<T>(input: &Array2<T>, targetshape: (usize, usize)) -> NdimageResult<Array2<T>>
where
    T: Float + FromPrimitive + Debug + std::ops::AddAssign + std::ops::DivAssign + 'static,
{
    let (src_h, src_w) = input.dim();
    let (dst_h, dst_w) = targetshape;

    let mut output = Array2::zeros((dst_h, dst_w));

    let scale_y = (src_h as f64 - 1.0) / (dst_h as f64 - 1.0);
    let scale_x = (src_w as f64 - 1.0) / (dst_w as f64 - 1.0);

    for y in 0..dst_h {
        for x in 0..dst_w {
            let src_y = T::from_f64(y as f64 * scale_y).unwrap_or(T::zero());
            let src_x = T::from_f64(x as f64 * scale_x).unwrap_or(T::zero());

            // Bilinear interpolation
            let y0 = src_y.floor();
            let x0 = src_x.floor();
            let y1 = (y0 + T::one()).min(T::from_usize(src_h - 1).expect("Operation failed"));
            let x1 = (x0 + T::one()).min(T::from_usize(src_w - 1).expect("Operation failed"));

            let dy = src_y - y0;
            let dx = src_x - x0;

            let y0_idx = y0.to_usize().unwrap_or(0).min(src_h - 1);
            let y1_idx = y1.to_usize().unwrap_or(0).min(src_h - 1);
            let x0_idx = x0.to_usize().unwrap_or(0).min(src_w - 1);
            let x1_idx = x1.to_usize().unwrap_or(0).min(src_w - 1);

            let v00 = input[[y0_idx, x0_idx]];
            let v01 = input[[y0_idx, x1_idx]];
            let v10 = input[[y1_idx, x0_idx]];
            let v11 = input[[y1_idx, x1_idx]];

            let v0 = v00 * (T::one() - dx) + v01 * dx;
            let v1 = v10 * (T::one() - dx) + v11 * dx;

            output[[y, x]] = v0 * (T::one() - dy) + v1 * dy;
        }
    }

    Ok(output)
}

#[cfg(test)]
mod tests {
    use super::*;
    use scirs2_core::ndarray::array;

    #[test]
    fn test_perspective_transform_identity() {
        let input = Array2::from_elem((10, 10), 1.0);

        // Identity perspective matrix
        let matrix = array![[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]];

        let result = perspective_transform(&input, &matrix, None, None, None, None)
            .expect("Operation failed");
        assert_eq!(result.shape(), input.shape());
    }

    #[test]
    fn test_thin_plate_spline_minimal() {
        let input = Array2::from_elem((10, 10), 1.0);

        // Three control points for minimal TPS
        let source_points = array![[0.0, 0.0], [5.0, 0.0], [0.0, 5.0]];

        let target_points = array![[0.0, 0.0], [5.0, 0.0], [0.0, 5.0]];

        let result = thin_plate_spline_transform(
            &input,
            &source_points,
            &target_points,
            None,
            None,
            None,
            None,
            None,
        )
        .expect("Operation failed");

        assert_eq!(result.shape(), input.shape());
    }

    #[test]
    fn test_pyramid_transform() {
        let input = Array2::from_elem((16, 16), 1.0);

        // Simple identity transform at each level
        let transform_fn =
            |arr: &Array2<f64>, _level: usize| -> NdimageResult<Array2<f64>> { Ok(arr.clone()) };

        let result = pyramid_transform(&input, transform_fn, Some(3), None, None, None)
            .expect("Operation failed");
        assert_eq!(result.shape(), input.shape());
    }

    #[test]
    fn test_perspective_transform_invalid_matrix() {
        let input = Array2::from_elem((10, 10), 1.0);
        let matrix = Array2::from_elem((2, 2), 1.0); // Wrong shape

        let result = perspective_transform(&input, &matrix, None, None, None, None);
        assert!(result.is_err());
    }

    #[test]
    fn test_tps_insufficient_points() {
        let input = Array2::from_elem((10, 10), 1.0);

        // Only two control points (need at least 3)
        let source_points = array![[0.0, 0.0], [5.0, 0.0]];

        let target_points = array![[0.0, 0.0], [5.0, 0.0]];

        let result = thin_plate_spline_transform(
            &input,
            &source_points,
            &target_points,
            None,
            None,
            None,
            None,
            None,
        );

        assert!(result.is_err());
    }
}