pub struct Wall {
    heights: Vec<u32>,
}

impl Wall {
    pub fn from_heights(heights: Vec<u32>) -> Self {
        Self { heights }
    }

    #[allow(dead_code)]
    fn create_empty(w: u32) -> Self {
        let heights = if w == 0 {
            vec![]
        } else if w == 1 {
            vec![0]
        } else {
            let mut v = Vec::with_capacity(w as usize);
            v.push(0);
            v.push(1);
            v
        };
        Self { heights }
    }

    pub fn calculate_row(&self, r: u32, stones: &mut [u32]) {
        let mut len = 1;
        let mut i = 0;
        for &height in self.heights.iter().chain([r].iter()) {
            if height == r {
                stones[i] = len;
                i += 1;
                len = 0;
            }
            len += 1;
        }
    }

    pub fn output(&self, n: u32, h: u32) {
        let mut stones = vec![0; n as usize];
        let mut toggle = 0;
        let colors = [
            "\x1b[31m", "\x1b[32m", "\x1b[33m", "\x1b[34m", "\x1b[35m", "\x1b[36m",
        ];
        for row in 0..h {
            self.calculate_row(row, &mut stones);
            for &len in stones.iter() {
                print!("{}", colors[toggle]);
                toggle = (toggle + 1) % colors.len();
                for _ in 0..len {
                    print!("◙");
                }
            }
            println!("\x1b[m");
        }
    }
}

/// Solve for a given N and return the resulting wall
pub struct Solver<T: num::PrimInt> {
    pub n: u32,
    /// calculated height [might not be correct!]
    pub h: u32,
    /// width
    pub w: u32,
    /// Use to store already used blocks as a bitmask
    solve_stack: Vec<SaveState<T>>,
    permutations: Vec<Vec<u32>>,
}

// used during row creation, will get deprecated
static mut COUNT: u32 = 0;

/// Save the current state for each row
#[derive(Clone, Eq, Copy)]
pub struct SaveState<T: num::PrimInt> {
    /// Mask of all currently used stones
    pub bitmask: T,
    /// sum of all placed stones
    pub sum: u32,
    /// the row index
    pub row: u32,
}

impl<T: num::PrimInt> std::cmp::Ord for SaveState<T> {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        self.sum.cmp(&other.sum)
    }
}
impl<T: num::PrimInt> std::cmp::PartialOrd for SaveState<T> {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        Some(self.sum.cmp(&other.sum))
    }
}
impl<T: num::PrimInt> std::cmp::PartialEq for SaveState<T> {
    fn eq(&self, other: &Self) -> bool {
        self.sum.eq(&other.sum)
    }
}

impl<T: num::PrimInt> SaveState<T> {
    pub fn new() -> Self {
        Self {
            bitmask: T::zero(),
            sum: 0,
            row: unsafe {
                COUNT += 1;
                COUNT
            },
        }
    }

    #[allow(dead_code)]
    pub fn set_bit(&mut self, stone: u32) {
        self.bitmask =
            self.bitmask | T::from(1 << stone).expect("Stone placing index out of bounds");
    }

    #[allow(dead_code)]
    pub fn get_bit(&self, stone: u32) -> T {
        self.bitmask & T::from(1 << stone).expect("Requested stone index out of bounds")
    }

    #[allow(dead_code)]
    pub fn bridges_gap(&self, gap: u32) -> bool {
        self.get_bit(gap - self.sum) == T::zero()
    }

    #[allow(dead_code)]
    pub fn bridge_gap(mut self, gap: u32) -> Self {
        self.set_bit(gap - self.sum);
        self.sum = gap;
        self
    }
}

impl<T: num::PrimInt> Solver<T> {
    pub fn new(n: u32) -> Self {
        let h = n / 2 + 1;
        let w = h * (n - 1);
        let mut heap = (1..=n).collect::<Vec<u32>>();
        let heap = permutohedron::Heap::new(&mut heap);
        let n_f = permutohedron::factorial(n as usize);
        let mut permutations = Vec::with_capacity(n_f);
        for data in heap {
            permutations.push(data.clone());
        }
        Self {
            n,
            h,
            w,
            solve_stack: vec![SaveState::new(); h as usize],
            permutations,
        }
    }

    pub fn solve(&mut self) {
        //for (n, i) in self.permutations.iter().enumerate() {
        //    println!("perm {}: {:?}", n, i);
        //}
        self.check_perm(&[3, 16, 23]);
        self.permute(
            permutohedron::factorial(self.n as usize),
            0,
            ((0..self.h).collect::<Vec<u32>>()).as_ref(),
        );
    }

    fn permute(&mut self, up: usize, index: usize, numbers: &[u32]) {
        if index as usize == numbers.len() {
            //println!("checking {:?}", numbers);
            if self.check_perm(&numbers) {
                println!("success");
            }
            return;
        }
        let mut new_num = Vec::from(numbers);
        for i in numbers[index as usize] as usize..up {
            self.permute(up, index + 1, &new_num);
            new_num[index] += 1;
            for n in (index + 1)..numbers.len() {
                new_num[n] = (i + 1 + n - index) as u32;
            }
        }
    }

    fn check_perm(&mut self, nums: &[u32]) -> bool {
        let mut sums = vec![self.w; self.w as usize];
        for (i, num) in nums.iter().enumerate() {
            let mut sum = 0;
            for n in self.permutations[*num as usize][..self.h as usize].iter() {
                sum += *n as usize;
                if sums[sum - 1] != self.w {
                    return false;
                }
                sums[sum - 1] = i as u32;
            }
        }
        Wall::from_heights(sums.iter().map(|x| *x as u32).collect()).output(self.w, self.h);
        true
    }

    #[allow(dead_code)]
    fn check_gaps(&mut self, wall: &Wall) -> bool {
        for (r, i) in wall.heights.iter().zip(1..self.w) {
            let stone = i - self.solve_stack[*r as usize].sum;
            self.solve_stack[*r as usize].set_bit(stone);
        }
        for state in self.solve_stack.as_ref() as &[SaveState<T>] {
            if state.bitmask != T::from(1 << self.n).unwrap() {
                return false;
            }
        }
        true
    }
}