use std::collections::HashSet;
use crate::Day;

pub struct Day09(Vec<Vec<u32>>);

impl Day for Day09 {
    fn init(content: String) -> anyhow::Result<Self> {
        Ok(Self(content.lines().map(|l| l.chars().map(|x| x.to_digit(10).unwrap()).collect::<Vec<_>>()).collect::<Vec<_>>()))
    }

    fn part1(&self) -> anyhow::Result<String> {
        let valid = self.low_points();
        Ok(format!("{}", valid.iter().map(|&p| self.get_point(p) + 1).sum::<u32>()))
    }

    fn part2(&self) -> anyhow::Result<String> {
        let valid = self.low_points();
        let mut sizes = valid.iter().map(|&low| {
            let mut basin_points = HashSet::new();
            let mut basin_edges = HashSet::new();
            basin_points.insert(low);
            basin_edges.extend(self.neighbours(low));
            while !basin_edges.is_empty() {
                let mut new_basin_edges = HashSet::new();
                for &p in &basin_edges {
                    new_basin_edges.extend(self.neighbours(p).filter(|&n| {
                        !basin_points.contains(&n)
                            && !basin_edges.contains(&n)
                            && self.get_point(p) < self.get_point(n)
                            && self.get_point(n) < 9
                    }));
                    basin_points.insert(p);
                }
                basin_edges = new_basin_edges;
            }
            basin_points.len()
        }).collect::<Vec<_>>();

        sizes.sort_unstable();
        let last_three = sizes.iter().rev().take(3);
        Ok(format!("{}", last_three.product::<usize>()))
    }
}

impl Day09 {
    fn neighbours(&self, (x, y): (usize, usize)) -> impl Iterator<Item=(usize, usize)> {
        [
            x.checked_sub(1).map(|r| (r, y)),
            y.checked_sub(1).map(|c| (x, c)),
            (x + 1 < self.dimensions().1).then(|| (x + 1, y)),
            (y + 1 < self.dimensions().0).then(|| (x, y + 1)),
        ].into_iter().flatten()
    }

    fn dimensions(&self) -> (usize, usize) {
        (self.0.len(), self.0[0].len())
    }

    fn check(&self, p: (usize, usize)) -> bool {
        for p2 in self.neighbours(p).collect::<Vec<_>>() {
            if self.get_point(p) >= self.get_point(p2) {
                return false;
            }
        }
        true
    }

    fn low_points(&self) -> Vec<(usize, usize)> {
        let (h, w) = self.dimensions();
        (0..h).flat_map(|y| (0..w).filter(|&x| self.check((x, y))).map(|x| (x, y)).collect::<Vec<_>>()).collect()
    }

    fn get_point(&self, (x, y): (usize, usize)) -> u32 {
        self.0[y][x]
    }
}

#[cfg(test)]
mod tests {
    use crate::{Day09, day_tests};
    use crate::day::Day;

    const INPUT: &str = r"2199943210
3987894921
9856789892
8767896789
9899965678";

    #[test]
    fn part1_test() -> anyhow::Result<()> {
        let d = Day09::init(INPUT.to_string())?;
        assert_eq!("15", d.part1()?);
        Ok(())
    }

    #[test]
    fn part2_test() -> anyhow::Result<()> {
        let d = Day09::init(INPUT.to_string())?;
        assert_eq!("1134", d.part2()?);
        Ok(())
    }

    day_tests!(Day09, "09", "439", "900900");
}