imgal 0.3.1

use ndarray::{ArrayBase, ArrayD, ArrayView, AsArray, Axis, Dimension, IxDyn, Slice, ViewRepr};
use rayon::prelude::*;

use crate::prelude::*;

/// Tile an n-dimensional image using division tiling.
///
/// # Description
///
/// Divides an n-dimensional image into a stack of array views representing
/// tiles created from the input array. Each axis of the input array is divided
/// by `div` into equally sized segments if `div` is a multiple of the length of
/// the axis to be sliced. If `div` is *not* a muliple of the axis length then
/// the remainder is added to the last tile's shape. This produces a total
/// of `divⁿ` tiles for a given array, where `n` is the total number of
/// dimensions. This function is *naive* in that it does not produce tiles
/// intended for image fusing. Instead the tiles are simple array slices of
/// the input data.
///
/// # Arguments
///
/// * `data`: The input n-dimensional image to be tiled.
/// * `div`: The base number of divisions per axis. This value must be `>0`.
/// * `threads`: The requested number of threads to use for parallel execution.
///   If `None` or `Some(1)` sequential execution is used. If `Some(0)`, then
///   the maximum available parallelism is used. Thread counts are clamped to
///   the systems maximum.
///
/// # Returns
///
/// * `Ok(Vec<ArrayView<'a, T, D>>)`: A vector containing views of all tiles in
///   row-major order. The length of the vector will be `divⁿ`, the number of
///   tiles.
/// * `Err(ImgalError)`: If `div == 0`.
pub fn div_tile<'a, T, A, D>(
    data: A,
    div: usize,
    threads: Option<usize>,
) -> Result<Vec<ArrayView<'a, T, D>>, ImgalError>
where
    A: AsArray<'a, T, D>,
    D: Dimension,
    T: 'a + AsNumeric,
{
    if div == 0 {
        return Err(ImgalError::InvalidParameterValueEqual {
            param_name: "div",
            value: 0,
        });
    }
    let data: ArrayBase<ViewRepr<&'a T>, D> = data.into();
    let shape = data.shape().to_vec();
    let n_dims = data.shape().len();
    let tile_positions: Vec<Vec<(isize, isize)>> = shape
        .iter()
        .map(|&v| get_div_start_stop_positions(div, v))
        .collect();
    let n_tiles: usize = tile_positions.iter().map(|v| v.len()).product();
    let tile_view = |t: usize| {
        let mut tile = data.clone();
        let mut remaining = t;
        (0..n_dims).for_each(|a| {
            let stride: usize = tile_positions.iter().skip(a + 1).map(|v| v.len()).product();
            let tile_pos = remaining / stride;
            remaining %= stride;
            let ax_slice = Slice {
                start: tile_positions[a][tile_pos].0,
                end: Some(tile_positions[a][tile_pos].1),
                step: 1,
            };
            tile.slice_axis_inplace(Axis(a), ax_slice);
        });
        tile
    };
    Ok(par!(threads,
    seq_exp: (0..n_tiles).map(&tile_view)
        .collect::<Vec<ArrayView<T, D>>>(),
    par_exp: (0..n_tiles).into_par_iter().map(&tile_view)
        .collect::<Vec<ArrayView<T, D>>>()
    ))
}

/// Untile a tile stack into an n-dimensional image.
///
/// # Description
///
/// Reconstructs (*.i.e.* untiles) an n-dimensional image by assembling a stack
/// of n-dimensional tiles as array views into a single output array of the
/// given `shape`. The input `tile_stack` is assumed to contain tiles resulting
/// from the `div_tile` function or a similar tiling scheme where tiles are
/// stored in row-major order. This function is *naive* in that it does not
/// offer any border fusing strategies.
///
/// # Arguments
///
/// * `tile_stack`: A vector containing views (*i.e.* tiles) to be reassembled
///   into a single n-dimensional image.
/// * `div`: The base number of divisions per axis. This value must be `>0`.
/// * `shape`: The shape of the output array. Its dimensionality must match the
///   dimensionality of the tiles.
///
/// # Returns
///
/// * `Ok(ArrayD<T>)`: An n-dimensional image with the given `shape` containing
///   all tiles in their corresponding positions.
/// * `Err(ImgalError)`: If `tile_stack.is_empty() == true`. If `div == 0`. If
///   `shape.len()` is not equal to the tile shape length. If expected tile
///   shapes do not match given tile shapes. If the number of tiles given does
///   not match the number of tiles expected.
pub fn div_untile<'a, T, D>(
    tile_stack: Vec<ArrayView<'a, T, D>>,
    div: usize,
    shape: &[usize],
) -> Result<ArrayD<T>, ImgalError>
where
    D: Dimension,
    T: 'a + AsNumeric,
{
    if tile_stack.is_empty() {
        return Err(ImgalError::InvalidParameterEmptyArray {
            param_name: "tile_stack",
        });
    }
    if div == 0 {
        return Err(ImgalError::InvalidParameterValueEqual {
            param_name: "div",
            value: 0,
        });
    }
    let n_dims = tile_stack[0].shape().len();
    if shape.len() != n_dims {
        return Err(ImgalError::MismatchedArrayLengths {
            a_arr_name: "tile shape",
            a_arr_len: n_dims,
            b_arr_name: "shape",
            b_arr_len: shape.len(),
        });
    }
    let tile_positions: Vec<Vec<(isize, isize)>> = shape
        .iter()
        .map(|&v| get_div_start_stop_positions(div, v))
        .collect();
    let n_tiles: usize = tile_positions.iter().map(|v| v.len()).product();
    if n_tiles != tile_stack.len() {
        return Err(ImgalError::InvalidArrayLengthExpected {
            arr_name: "tile_stack",
            expected: n_tiles,
            got: tile_stack.len(),
        });
    }
    let tile_shape: Vec<usize> = shape.iter().map(|&v| v / div).collect();
    if tile_shape != tile_stack[0].shape() {
        return Err(ImgalError::MismatchedArrayShapes {
            a_arr_name: "expected tile",
            a_shape: tile_shape,
            b_arr_name: "input tile",
            b_shape: tile_stack[0].shape().to_vec(),
        });
    }
    let mut untile_arr: ArrayD<T> = ArrayD::from_elem(IxDyn(shape), T::default());
    (0..n_tiles).for_each(|t| {
        let tile_view = tile_stack[t].view();
        let mut untile_view = untile_arr.view_mut();
        let mut remaining = t;
        (0..n_dims).for_each(|a| {
            let stride: usize = tile_positions.iter().skip(a + 1).map(|v| v.len()).product();
            let tile_pos = remaining / stride;
            remaining %= stride;
            let ax_slice = Slice {
                start: tile_positions[a][tile_pos].0,
                end: Some(tile_positions[a][tile_pos].1),
                step: 1,
            };
            untile_view.slice_axis_inplace(Axis(a), ax_slice);
        });
        untile_view.assign(&tile_view);
    });
    Ok(untile_arr)
}

/// Compute evenly spaced start and stop positions.
///
/// # Arguments
///
/// * `div`: The base number of divisions per axis. This value must be `>0`.
/// * `axis_len`: The length of the axis to compute start and stop positions.
///
/// # Returns
///
/// * `Vec<(isize, isize)>`: A tuple of start and stop positions,
///   `(start, stop)` along an axis. If `div` is not a multiple of `axis_len`,
///   then the last tile will be larger by the remainder.
fn get_div_start_stop_positions(div: usize, axis_len: usize) -> Vec<(isize, isize)> {
    let mut start_stop_arr: Vec<(isize, isize)> = Vec::with_capacity(div);
    let inc = (axis_len / div) as isize;
    (0..div).fold(0_isize, |acc, _| {
        let start = acc;
        let stop = acc + inc;
        start_stop_arr.push((start, stop));
        stop
    });
    start_stop_arr[div.saturating_sub(1)].1 = axis_len as isize;
    start_stop_arr
}