pub struct StoneWall {
    rows: Vec<Vec<u32>>,
}

impl StoneWall {
    pub fn create_empty(n: u32) -> Self {
        let n = n as usize;
        let h = n / 2 + 1;
        let mut rows = vec![vec![0; n]; h];
        rows.get_mut(0).map(|r| r[0] = 1);
        rows.get_mut(1).map(|r| r[1] = 2);
        Self { rows }
    }

    pub fn set_stone(&mut self, row: u32, pos: u32, stone: u32) -> Option<()> {
        self.rows
            .get_mut(row as usize)
            .and_then(|v| v.get_mut(pos as usize))
            .map(|v| *v = stone)
    }

    pub fn output(&mut self) {
        let colors = [[31, 32], [33, 35]];
        //self.rows.sort_by_key(|x| x[0]);

        for (i, row) in self.rows.iter().enumerate() {
            for (j, &stone) in row.iter().enumerate() {
                //print!("{}", colors[i & 1][j & 1]);
                print!(
                    "\x1b[{}m{}",
                    colors[i & 1][j & 1],
                    "◙".repeat(stone as usize)
                );
            }
            println!("\x1b[m");
        }
    }
}

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

impl GapHeights {
    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 add_gap(&mut self, height: u32) {
        self.add_gap(height)
    }

    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 mut colors = [
            "\x1b[31m", "\x1b[32m", "\x1b[33m", "\x1b[34m", "\x1b[35m", "\x1b[36m",
        ]
        .iter()
        .cycle();
        let mut indices = Vec::with_capacity(h as usize);
        for i in 0..n {
            let n = self.heights[i as usize];
            if indices.iter().all(|x| *x != n) {
                indices.push(n);
            }
        }
        //println!("{:?} : {:?}", indices, self.heights);
        for row in indices {
            self.calculate_row(row, &mut stones);
            for &len in stones.iter() {
                print!("{}", colors.next().unwrap());
                for _ in 0..len {
                    print!("◙");
                }
            }
            println!("\x1b[m");
        }
    }

    pub fn iter<'a>(&'a self) -> GapIter<'a> {
        GapIter {
            gap_heights: self,
            i: 0,
        }
    }

    pub fn as_stone_wall(&self, n: u32) -> StoneWall {
        let h = n / 2 + 1;
        let mut rows = Vec::with_capacity(h as usize);
        for i in 0..h {
            let mut row = vec![0; n as usize];
            self.calculate_row(i, &mut row);
            rows.push(row);
        }
        StoneWall { rows }
    }
}

pub struct GapIter<'a> {
    gap_heights: &'a GapHeights,
    i: usize,
}

impl<'a> Iterator for GapIter<'a> {
    type Item = u32;
    fn next(&mut self) -> Option<Self::Item> {
        let i = self.i;
        self.i += 1;
        self.gap_heights.heights.get(i).map(|&x| x)
    }
}