use std::iter::Sum;
use num::Float;

use ndarray::prelude::*;
use ndarray::{IntoDimension, Zip, Si, S, Data};
use itertools::{zip, Itertools};

#[derive(Debug, Clone, Copy)]
pub enum Boundary {
    Nearest,
    Reflect,
    Wrap,
    Mirror,
}

fn __pad<S, D>(array: &ArrayBase<S, D>, window_size: &D, boundary: Boundary) -> Array<S::Elem, D>
    where S: Data,
          S::Elem: Copy,
          D: Dimension
{
    let window_size = window_size.as_array_view();
    debug_assert_eq!(window_size.len(), array.ndim());

    let mut pad_array: Array<S::Elem, D> = unsafe {
        let mut size = array.raw_dim().clone();
        Zip::from(size.as_array_view_mut()).and(window_size).apply(|mut s, &w| {
            *s += w - 1;
        });
        Array::uninitialized(size)
    };

    {
        let mut pad_array = {
            let dyn_size = pad_array.shape().to_owned();
            pad_array.view_mut().into_shape(dyn_size).unwrap()
        };

        let mut slice: Vec<Si> = zip(window_size, array.shape())
            .map(|(window_size, size)| {
                Si((window_size / 2) as isize,
                   Some((window_size / 2 + size) as isize),
                   1)
            })
            .collect();

        pad_array.slice_mut(slice.as_slice()).assign(&array);

        for (k, (&window_size, &size)) in zip(window_size, array.shape()).enumerate() {
            let bound = (window_size / 2) as isize;

            if bound == 0 {
                continue;
            }

            slice[k] = S;
            let centre = pad_array.slice_mut(slice.as_slice());
            let (mut left, centre) = centre.split_at(Axis(k), window_size / 2);
            let (centre, mut right) = centre.split_at(Axis(k), size);

            let mut assign_to_axis_slice = {
                let mut slice = vec![S; array.ndim()];

                move |to: &mut ArrayViewMut<S::Elem, _>, s| {
                    slice[k] = s;
                    let from = centre.slice(slice.as_slice());
                    to.assign(&from);
                }
            };

            use self::Boundary::*;
            match boundary {
                Nearest => {
                    assign_to_axis_slice(&mut left, Si(0, Some(1), 1));
                    assign_to_axis_slice(&mut right, Si(-1, None, 1));
                }
                Reflect => {
                    assign_to_axis_slice(&mut left, Si(-bound - 1, Some(-1), -1));
                    assign_to_axis_slice(&mut right, Si(1, Some(bound + 1), -1));
                }
                Mirror => {
                    assign_to_axis_slice(&mut left, Si(-bound, None, -1));
                    assign_to_axis_slice(&mut right, Si(0, Some(bound), -1));
                }
                Wrap => {
                    assign_to_axis_slice(&mut left, Si(-bound, None, 1));
                    assign_to_axis_slice(&mut right, Si(0, Some(bound), 1));
                }
            }
        }
    }

    pad_array
}

pub fn pad<A, D, W>(array: ArrayView<A, D>, window_size: W, boundary: Boundary) -> Array<A, D>
    where A: Copy,
          D: Dimension,
          W: IntoDimension<Dim = D>
{
    __pad(&array, &window_size.into_dimension(), boundary)
}

pub fn rolling_mean<A, D, W>(x: ArrayView<A, D>, w: W, b: Boundary) -> Array<A, D>
    where A: Float,
          D: Dimension,
          W: IntoDimension<Dim = D>
{
    let w = w.into_dimension();

    let y = __pad(&x, &w, b);

    let mut x = x.to_owned();
    Zip::from(x.view_mut())
        .and(y.windows(w))
        .apply(|mut x, y| *x = y.scalar_sum() / A::from(y.len()).unwrap());
    x
}

pub fn rel_max<A, D, W>(x: ArrayView<A, D>, w: W, b: Boundary) -> Array<bool, D>
    where A: PartialOrd + Copy,
          D: Dimension,
          W: IntoDimension<Dim = D>
{
    let w = w.into_dimension();

    let y = __pad(&x, &w, b);

    let mut b = unsafe { Array::uninitialized(x.raw_dim()) };

    let mut j = w.clone();
    j.as_array_view_mut().mapv_inplace(|j| j / 2);

    Zip::from(b.view_mut())
        .and(y.windows(w))
        .apply(|mut b, y| *b = y.iter().all(|x| *x <= y[j.clone()]));
    b
}

pub fn locate_nearest_peak<A, D, F, I>(b: ArrayView<bool, D>, i: &[A], f: F) -> Option<D>
    where A: Float + Sum<A>,
          D: Dimension,
          F: Fn(&D::Pattern) -> I,
          I: IntoIterator<Item = A>
{
    b.indexed_iter()
        .filter(|&(_, &b)| b)
        .map(|(j, _)| {
            let d = i.iter()
                .zip(f(&j))
                .map(|(&i, j)| (i - j).powi(2))
                .sum::<A>();
            (j, d)
        })
        .fold1(|(j0, d0), (j1, d1)| if d1 < d0 { (j1, d1) } else { (j0, d0) })
        .map(|(j, _)| j.into_dimension())
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn test() {
        let x = array![1., 2., 3.];

        let y = pad(x.view(), 1, Boundary::Wrap);
        assert_eq!(y, x);

        let y = pad(x.view(), 5, Boundary::Wrap);
        assert_eq!(y, array![2., 3., 1., 2., 3., 1., 2.]);

        let y = pad(x.view(), 5, Boundary::Nearest);
        assert_eq!(y, array![1., 1., 1., 2., 3., 3., 3.]);

        let y = pad(x.view(), 5, Boundary::Reflect);
        assert_eq!(y, array![2., 1., 1., 2., 3., 3., 2.]);

        let y = pad(x.view(), 5, Boundary::Mirror);
        assert_eq!(y, array![3., 2., 1., 2., 3., 2., 1.]);

        let x = array![[1, 2], [3, 4]];

        let y = pad(x.view(), (3, 3), Boundary::Wrap);
        assert_eq!(y,
                   array![[4, 3, 4, 3], [2, 1, 2, 1], [4, 3, 4, 3], [2, 1, 2, 1]]);

        let x = Array1::<usize>::zeros(10);
        assert_eq!(12, pad(x.view(), 3, Boundary::Wrap).len());
    }
}