use std::marker::PhantomData;

/// Heap's algorithm for generating permutations, recursive version.
///
/// The recursive algorithm supports slices of any size (even though
/// only a small number of elements is practical), and is generally
/// much faster than the iterative version.
pub fn heap_recursive<T, F>(xs: &mut [T], mut f: F) where F: FnMut(&mut [T])
{
    heap_unrolled_(xs.len(), xs, &mut f);
}

/// Unrolled version of heap's algorithm due to Sedgewick
fn heap_unrolled_<T>(n: usize, xs: &mut [T], f: &mut FnMut(&mut [T])) {
    match n {
        0 | 1 => f(xs),
        2 => {
            // [1, 2], [2, 1]
            f(xs);
            xs.swap(0, 1);
            f(xs);
        }
        3 => {
            // [1, 2, 3], [2, 1, 3], [3, 1, 2], [1, 3, 2], [2, 3, 1], [3, 2, 1]
            f(xs);
            xs.swap(0, 1);
            f(xs);
            xs.swap(0, 2);
            f(xs);
            xs.swap(0, 1);
            f(xs);
            xs.swap(0, 2);
            f(xs);
            xs.swap(0, 1);
            f(xs);
        }
        n => for i in 0..n {
            heap_unrolled_(n - 1, xs, f);
            let j = if n % 2 == 0 { i } else { 0 };
            xs.swap(j, n - 1);
        }
    }
}

/// Maximum number of elements we can generate permutations for using
/// Heap's algorithm (iterative version).
pub const MAXHEAP: usize = 16;

/// Heap's algorithm for generating permutations.
///
/// An iterative method of generating all permutations of a sequence.
///
/// Note that for *n* elements there are *n!* (*n* factorial) permutations.
pub struct Heap<'a, Data: 'a + ?Sized, T: 'a> {
    data: &'a mut Data,
    c: [u8; MAXHEAP],
    n: usize,
    // we can store up to 20! in 64 bits.
    index: u64,
    _element: PhantomData<&'a mut T>
}

impl<'a, T, Data: ?Sized> Heap<'a, Data, T>
    where Data: AsMut<[T]>
{
    /// Create a new `Heap`.
    pub fn new(data: &'a mut Data) -> Self {
        assert!(data.as_mut().len() <= MAXHEAP);
        Heap {
            data: data,
            c: [0; MAXHEAP],
            n: 0,
            index: 0,
            _element: PhantomData,
        }
    }

    /// Return a reference to the inner data
    pub fn get(&self) -> &Data {
        self.data
    }

    /// Return a mutable reference to the inner data
    pub fn get_mut(&mut self) -> &mut Data {
        self.data
    }

    /// Reset the permutations walker, without changing the data. It allows
    /// generating permutations again with the current state as starting
    /// point.
    pub fn reset(&mut self) {
        self.n = 0;
        for c in &mut self.c[..] { *c = 0; }
        self.index = 0;
    }

    /// Step the internal data into the next permutation and return
    /// a reference to it. Return `None` when all permutations
    /// have been visited.
    ///
    /// Note that for *n* elements there are *n!* (*n* factorial) permutations.
    pub fn next_permutation(&mut self) -> Option<&mut Data> {
        if self.index == 0 {
            self.index += 1;
            Some(self.data)
        } else {
            while self.n < self.data.as_mut().len() {
                let nb = self.n as u8;
                let nu = self.n;
                let c = &mut self.c;
                if c[nu] < nb {
                    // `n` acts like the current length - 1 of the slice we are permuting
                    // `c[n]` acts like `i` in the recursive algorithm
                    let j = if (nu + 1) % 2 == 0 { c[nu] as usize } else { 0 };
                    self.data.as_mut().swap(j, nu);
                    c[nu] += 1;
                    self.n = 0;
                    return Some(self.data);
                } else {
                    c[nu] = 0;
                    self.n += 1;
                }
            }
            None
        }
    }
}

/// Iterate the permutations
///
/// **Note:** You can also generate the permutations lazily by using
/// `.next_permutation()`.
impl<'a, Data: ?Sized, T> Iterator for Heap<'a, Data, T>
    where Data: AsMut<[T]> + ToOwned,
{
    type Item = Data::Owned;
    fn next(&mut self) -> Option<Self::Item> {
        match self.next_permutation() {
            None => None,
            Some(perm) => Some(perm.to_owned()),
        }
    }
}

/// Compute *n!* (*n* factorial)
pub fn factorial(n: usize) -> usize {
    let mut prod = 1;
    for x in 1..n + 1 { prod *= x; }
    prod
}

#[test]
fn first_and_reset() {
    let mut data = [1, 2, 3];
    let mut heap = Heap::new(&mut data);
    let mut perm123 = vec![[1, 2, 3], [2, 1, 3], [3, 1, 2], [1, 3, 2], [2, 3, 1], [3, 2, 1]];
    assert_eq!(heap.by_ref().collect::<Vec<_>>(), perm123);

    // test reset
    heap.reset();
    // for the 1,2,3 case this happens to work out to the reverse order
    perm123.reverse();
    assert_eq!(heap.by_ref().collect::<Vec<_>>(), perm123);
}

#[test]
fn permutations_0_to_6() {
    let mut data = [0; 6];
    for n in 0..data.len() {
        let count = factorial(n);
        for (index, elt) in data.iter_mut().enumerate() {
            *elt = index + 1;
        }
        let mut permutations = Heap::new(&mut data[..n]).collect::<Vec<_>>();
        assert_eq!(permutations.len(), count);
        permutations.sort();
        permutations.dedup();
        assert_eq!(permutations.len(), count);
        // Each permutation contains all of 1 to n
        assert!(permutations.iter().all(|perm| perm.len() == n));
        assert!(permutations.iter().all(|perm| (1..n + 1).all(|i| perm.iter().position(|elt| *elt == i).is_some())));
    }
}

#[test]
fn permutations_0_to_6_recursive() {
    let mut data = [0; 6];
    for n in 0..data.len() {
        let count = factorial(n);
        for (index, elt) in data.iter_mut().enumerate() {
            *elt = index + 1;
        }
        let mut permutations = Vec::with_capacity(count);
        heap_recursive(&mut data[..n], |elt| permutations.push(elt.to_owned()));
        println!("{:?}", permutations);
        assert_eq!(permutations.len(), count);
        permutations.sort();
        permutations.dedup();
        assert_eq!(permutations.len(), count);
        // Each permutation contains all of 1 to n
        assert!(permutations.iter().all(|perm| perm.len() == n));
        assert!(permutations.iter().all(|perm| (1..n + 1).all(|i| perm.iter().position(|elt| *elt == i).is_some())));
    }
}