// This file is part of faster, the SIMD library for humans.
// Copyright 2017 Adam Niederer <adam.niederer@gmail.com>

// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

use iters::{PackedIterator};
use vecs::{Packed, Packable};

/// A lazy iterator which returns tuples of the elements of its contained
/// iterators.
pub struct PackedZip<T> {
    iters: T
}

/// A lazy mapping iterator which applies its function to a stream of tuples of
/// vectors.
pub struct PackedZipMap<I, F> where I : PackedZippedIterator {
    iter: I,
    func: F,
    defaults: I::Vectors
}

/// A trait which can transform a collection of iterators into a `PackedZip`
pub trait IntoPackedZip : Sized {
    /// Return an iterator which may iterate over `self` in lockstep.
    fn zip(self) -> PackedZip<Self>;
}

/// A macro which takes a number n and an expression, and returns a tuple
/// containing n copies of the expression. Only works for numbers less than or
/// equal to 12.
///
/// ```
/// #[macro_use] extern crate faster;
/// use faster::*;
///
/// # fn main() {
/// assert_eq!(tuplify!(2, 1), (1, 1));
/// assert_eq!(tuplify!(5, "hi"), ("hi", "hi", "hi", "hi", "hi"));
/// assert_eq!(tuplify!(3, i8s::splat(0)), (i8s::splat(0), i8s::splat(0), i8s::splat(0)));
/// # }
/// ```
#[macro_export] macro_rules! tuplify {
    (1, $i:expr) => { ($i) };
    (2, $i:expr) => { ($i, $i) };
    (3, $i:expr) => { ($i, $i, $i) };
    (4, $i:expr) => { ($i, $i, $i, $i) };
    (5, $i:expr) => { ($i, $i, $i, $i, $i) };
    (6, $i:expr) => { ($i, $i, $i, $i, $i, $i) };
    (7, $i:expr) => { ($i, $i, $i, $i, $i, $i, $i) };
    (8, $i:expr) => { ($i, $i, $i, $i, $i, $i, $i, $i) };
    (9, $i:expr) => { ($i, $i, $i, $i, $i, $i, $i, $i, $i) };
    (10, $i:expr) => { ($i, $i, $i, $i, $i, $i, $i, $i, $i, $i) };
    (11, $i:expr) => { ($i, $i, $i, $i, $i, $i, $i, $i, $i, $i, $i) };
    (12, $i:expr) => { ($i, $i, $i, $i, $i, $i, $i, $i, $i, $i, $i, $i) };
}

/// A collection of packed iterators of the same scalar length and vector width,
/// which may be iterated over in lockstep.
pub trait PackedZippedIterator : ExactSizeIterator + Sized {
    type Scalars : Copy + Sized;
    type Vectors : Copy + Sized;

    /// Return the width of this iterator's constitutent vectors.
    fn width(&self) -> usize;

    /// Return the length of this iterator, measured in scalar elements.
    fn scalar_len(&self) -> usize;

    /// Return the current position of this iterator, measured in scalar
    /// elements.
    fn scalar_position(&self) -> usize;

    /// Pack and return a vector containing the next `self.width()` elements
    /// of the iterator, or return None if there aren't enough elements left
    fn next_vectors(&mut self) -> Option<Self::Vectors>;

    /// Pack and return a partially full vector containing upto the next
    /// `self.width()` of the iterator, or None if no elements are left.
    /// Elements which are not filled are instead initialized to default.
    fn next_partials(&mut self, default: Self::Vectors) -> Option<(Self::Vectors, usize)>;

    /// Pack and return a splatted vector containing the next element
    /// of the iterator, or None if no elements are left.
    fn next_splats(&mut self) -> Option<Self::Vectors>;

    #[inline(always)]
    /// Return an iterator which calls `func` on vectors of elements.
    fn simd_map<B, F>(self, defaults: Self::Vectors, func: F) -> PackedZipMap<Self, F>
        where F : FnMut(Self::Vectors) -> B {
        PackedZipMap {
            iter: self,
            func: func,
            defaults: defaults
        }
    }

    #[inline(always)]
    /// Return a vector generated by reducing `func` over accumulator `start`
    /// and the values of this iterator, initializing all vectors to `default`
    /// before populating them with elements of the iterator.
    ///
    /// # Examples
    ///
    /// ```
    /// extern crate faster;
    /// use faster::*;
    ///
    /// # fn main() {
    /// let reduced = (&[2.0f32; 100][..]).simd_iter()
    ///    .simd_reduce(f32s::splat(0.0), f32s::splat(0.0), |acc, v| acc + v);
    /// # }
    /// ```
    ///
    /// In this example, on a machine with 4-element vectors, the argument to
    /// the last call of the closure is
    ///
    /// ```rust,ignore
    /// [ 2.0 | 2.0 | 2.0 | 2.0 ]
    /// ```
    ///
    /// and the result of the reduction is
    ///
    /// ```rust,ignore
    /// [ 50.0 | 50.0 | 50.0 | 50.0 ]
    /// ```
    ///
    /// whereas on a machine with 8-element vectors, the last call is passed
    ///
    /// ```rust,ignore
    /// [ 2.0 | 2.0 | 2.0 | 2.0 | 0.0 | 0.0 | 0.0 | 0.0 ]
    /// ```
    ///
    /// and the result of the reduction is
    ///
    /// ```rust,ignore
    /// [ 26.0 | 26.0 | 26.0 | 26.0 | 24.0 | 24.0 | 24.0 | 24.0 ]
    /// ```
    ///
    /// # Footgun Warning
    ///
    /// The results of `simd_reduce` are not portable, and it is your
    /// responsibility to interepret the result in such a way that the it is
    /// consistent across different architectures. See [`Packed::sum`] and
    /// [`Packed::product`] for built-in functions which may be helpful.
    ///
    /// [`Packed::sum`]: vecs/trait.Packed.html#tymethod.sum
    /// [`Packed::product`]: vecs/trait.Packed.html#tymethod.product
    fn simd_reduce<A, F>(&mut self, start: A, default: Self::Vectors, mut func: F) -> A
        where F : FnMut(A, Self::Vectors) -> A {
        let mut acc: A;
        if let Some(v) = self.next_vectors() {
            acc = func(start, v);
            while let Some(mut v) = self.next_vectors() {
                acc = func(acc, v);
            }
            if let Some((v, _)) = self.next_partials(default) {
                acc = func(acc, v);
            }
            debug_assert!(self.next_partials(default).is_none());
            acc
        } else if let Some((v, _)) = self.next_partials(default) {
            acc = func(start, v);
            debug_assert!(self.next_partials(default).is_none());
            acc
        } else {
            start
        }
    }
}

macro_rules! impl_iter_zip {
    (($($a:tt),*), ($($b:tt),*), ($($n:tt),*)) => (
        impl<$($a),*> IntoPackedZip for ($($a),*) where $($a : PackedIterator),* {

            #[inline(always)]
            fn zip(self) -> PackedZip<Self> {
                PackedZip { iters: self }
            }
        }

        // impl<'a, $($a),*> IntoPackedRefMutIterator<'a> for ($(&'a mut $a),*) where $(&'a mut $a : PackedIterator),* {
        //     type Iter = PackedZip<($(&'a mut $a),*)>;

        //     fn simd_iter_mut(&mut self) -> Self::Iter {
        //         PackedZip { iters: self }
        //     }
        // }

        impl<Z, $($a),*> ExactSizeIterator for PackedZip<($($a),*)>
            where $($a : PackedIterator<Scalar = Z, Item = Z>),*, Z : Packable {

            #[inline(always)]
            fn len(&self) -> usize {
                debug_assert!($(self.iters.$n.len() == self.iters.0.len())&&*);
                self.iters.0.len()
            }
        }

        impl<Z, $($a),*> Iterator for PackedZip<($($a),*)>
            where $($a : PackedIterator<Scalar = Z, Item = Z>),*, Z : Packable {
            type Item = ($(<$a as Iterator>::Item),*);

            fn next(&mut self) -> Option<Self::Item> {
                Some(($(self.iters.$n.next()?),*))
            }
        }

        impl<Z, $($a),*> PackedZippedIterator for PackedZip<($($a),*)>
            where $($a : PackedIterator<Scalar = Z, Item = Z>),*, Z : Packable {
            type Vectors = ($($a::Vector),*);
            type Scalars = ($($a::Scalar),*);

            #[inline(always)]
            fn width(&self) -> usize {
                debug_assert!($(self.iters.$n.width() == self.iters.0.width())&&*);
                self.iters.0.width()
            }

            #[inline(always)]
            fn scalar_len(&self) -> usize {
                debug_assert!($(self.iters.$n.scalar_len() == self.iters.0.scalar_len())&&*);
                self.iters.0.scalar_len()
            }

            #[inline(always)]
            fn scalar_position(&self) -> usize {
                debug_assert!($(self.iters.$n.scalar_position() == self.iters.0.scalar_position())&&*);
                self.iters.0.scalar_position()
            }

            #[inline(always)]
            fn next_vectors(&mut self) -> Option<Self::Vectors> {
                Some(($(self.iters.$n.next_vector()?),*))
            }

            #[inline(always)]
            fn next_splats(&mut self) -> Option<Self::Vectors> {
                Some(($($a::Vector::splat(self.iters.$n.next()?)),*))
            }

            #[inline(always)]
            fn next_partials(&mut self, default: Self::Vectors) -> Option<(Self::Vectors, usize)> {
                let a = ($(self.iters.$n.next_partial(default.$n)),*);
                // Ensure everything is None, Or nothing is None and all vectors
                // are the same size.
                debug_assert!($((!a.$n.is_none() && a.$n.unwrap().1 == a.0.unwrap().1))&&*
                              || $(a.$n.is_none())&&*);
                if !a.0.is_none() {
                    Some((($(a.$n.unwrap().0),*), a.0.unwrap().1))
                } else {
                    None
                }
            }
        }
    );
}

impl<I, F, A> Iterator for PackedZipMap<I, F>
    where I : PackedZippedIterator, F : FnMut(I::Vectors) -> A, A : Packed {
    type Item = A::Scalar;

    #[inline(always)]
    fn next(&mut self) -> Option<Self::Item> {
        self.iter.next_splats().map(&mut self.func).map(A::coalesce)
    }
}

impl<I, F, A> ExactSizeIterator for PackedZipMap<I, F>
    where I : PackedZippedIterator, F : FnMut(I::Vectors) -> A, A : Packed {

    #[inline(always)]
    fn len(&self) -> usize {
        self.iter.len()
    }
}

impl<I, F, A> PackedIterator for PackedZipMap<I, F>
    where I : PackedZippedIterator, F : FnMut(I::Vectors) -> A, A : Packed {
    type Vector = A;
    type Scalar = A::Scalar;

    #[inline(always)]
    fn width(&self) -> usize {
        self.iter.width()
    }

    #[inline(always)]
    fn scalar_len(&self) -> usize {
        self.iter.scalar_len()
    }

    #[inline(always)]
    fn scalar_position(&self) -> usize {
        self.iter.scalar_position()
    }

    #[inline(always)]
    fn next_vector(&mut self) -> Option<Self::Vector> {
        self.iter.next_vectors().map(&mut self.func)
    }

    #[inline(always)]
    fn next_partial(&mut self, default: Self::Vector) -> Option<(Self::Vector, usize)> {
        let (v, n) = self.iter.next_partials(self.defaults)?;
        Some((default.merge_partitioned((&mut self.func)(v), n), n))
    }
}


impl_iter_zip!((A, B),
               (AA, BB),
               (0, 1));
impl_iter_zip!((A, B, C),
               (AA, BB, CC),
               (0, 1, 2));
impl_iter_zip!((A, B, C, D),
               (AA, BB, CC, DD),
               (0, 1, 2, 3));
impl_iter_zip!((A, B, C, D, E),
               (AA, BB, CC, DD, EE),
               (0, 1, 2, 3, 4));
impl_iter_zip!((A, B, C, D, E, F),
               (AA, BB, CC, DD, EE, FF),
               (0, 1, 2, 3, 4, 5));
impl_iter_zip!((A, B, C, D, E, F, G),
               (AA, BB, CC, DD, EE, FF, GG),
               (0, 1, 2, 3, 4, 5, 6));
impl_iter_zip!((A, B, C, D, E, F, G, H),
               (AA, BB, CC, DD, EE, FF, GG, HH),
               (0, 1, 2, 3, 4, 5, 6, 7));
impl_iter_zip!((A, B, C, D, E, F, G, H, I),
               (AA, BB, CC, DD, EE, FF, GG, HH, II),
               (0, 1, 2, 3, 4, 5, 6, 7, 8));
impl_iter_zip!((A, B, C, D, E, F, G, H, I, J),
               (AA, BB, CC, DD, EE, FF, GG, HH, II, JJ),
               (0, 1, 2, 3, 4, 5, 6, 7, 8, 9));
impl_iter_zip!((A, B, C, D, E, F, G, H, I, J, K),
               (AA, BB, CC, DD, EE, FF, GG, HH, II, JJ, KK),
               (0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10));
impl_iter_zip!((A, B, C, D, E, F, G, H, I, J, K, L),
               (AA, BB, CC, DD, EE, FF, GG, HH, II, JJ, KK, LL),
               (0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11));
impl_iter_zip!((A, B, C, D, E, F, G, H, I, J, K, L, M),
               (AA, BB, CC, DD, EE, FF, GG, HH, II, JJ, KK, LL, MM),
               (0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12));