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

impl Wall {
    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
        };
        let heights = vec![0, 1, 0, 2, 1, 0, 2, 2, 1];
        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<T>,
}

impl<T: num::PrimInt> Solver<T> {
    pub fn new(n: usize) -> Self {
        let h = n / 2 + 1;
        let w = h * (n - 1);
        Self {
            n: (n as u32),
            h: (h as u32),
            w: (w as u32),
            solve_stack: Vec::with_capacity(n),
        }
    }

    pub fn solve(&mut self) -> Wall {
        let wall = Wall::create_empty(self.w);
        wall.heights
            .iter()
            .take(2)
            .zip(0..1)
            .for_each(|(x, i)| self.set_stone(*x as usize, i));
        wall
    }

    fn set_stone(&mut self, row: usize, stone: u32) {
        self.solve_stack[row] =
            self.solve_stack[row] | T::from(1 << stone).expect("Stone placing index out of bounds");
    }

    fn get_stone(&self, row: usize, stone: u32) {
        self.solve_stack[row] & T::from(1 << stone).expect("Requested stone index out of bounds");
    }
}