use itertools::Itertools;
use std::collections::{HashSet, VecDeque};

pub fn day9() {
    // let input = r#"
    // 2199943210
    // 3987894921
    // 9856789892
    // 8767896789
    // 9899965678
    //     "#;
    let input = include_str!("../input/day9.txt");

    let grid = input
        .trim()
        .lines()
        .map(|l| {
            l.trim()
                .chars()
                .map(|c| c.to_digit(10).unwrap())
                .collect::<Vec<_>>()
        })
        .collect::<Vec<_>>();

    let grid_slice = &grid.iter().map(|row| row.as_slice()).collect::<Vec<_>>();

    let minima = find_local_minima(grid_slice);
    let sum: u32 = minima.iter().map(|(row, col)| &grid[*row][*col] + 1).sum();
    println!("p1 score sum: {}", sum);

    dbg!(find_basin_size(&(0, 1), grid_slice));
    dbg!(find_basin_size(&(0, 9), grid_slice));

    let basins_product: i32 = minima
        .iter()
        .map(|pos| find_basin_size(pos, grid_slice) as i32)
        .sorted_by_key(|count| -*count)
        .take(3)
        .product();
    println!("p2 basin score product: {}", basins_product);
}

fn find_local_minima(grid: &[&[u32]]) -> Vec<(usize, usize)> {
    let mut res = vec![];
    let rows = grid.len();
    let cols = grid[0].len();
    for row in 0..rows {
        for col in 0..cols {
            let cur = grid[row][col];
            let up = row == 0 || grid[row - 1][col] > cur;
            let down = row == rows - 1 || grid[row + 1][col] > cur;
            let left = col == 0 || grid[row][col - 1] > cur;
            let right = col == cols - 1 || grid[row][col + 1] > cur;
            if up && down && left && right {
                res.push((row, col));
            }
        }
    }
    res
}

fn find_basin_size(pos: &(usize, usize), grid: &[&[u32]]) -> usize {
    let rows = grid.len();
    let cols = grid[0].len();

    let mut queue = VecDeque::from([*pos]);
    let mut visited: HashSet<(usize, usize)> = HashSet::new();

    while let Some(p) = queue.pop_front() {
        let height = grid[p.0][p.1];
        if visited.contains(&p) || height == 9 {
            continue;
        }
        visited.insert(p);
        // println!("({}, {})", p.0, p.1);
        if p.0 > 0 && grid[p.0 - 1][p.1] >= height {
            queue.insert(0, (p.0 - 1, p.1));
        }
        if p.0 < rows - 1 && grid[p.0 + 1][p.1] >= height {
            queue.insert(0, (p.0 + 1, p.1));
        }
        if p.1 > 0 && grid[p.0][p.1 - 1] >= height {
            queue.insert(0, (p.0, p.1 - 1));
        }
        if p.1 < cols - 1 && grid[p.0][p.1 + 1] >= height {
            queue.insert(0, (p.0, p.1 + 1));
        }
    }

    visited.len()
}