use aoc2024::matrix;
use std::collections::HashSet;

static AROUND_DELTAS: [(isize, isize); 4] = [(-1, 0), (0, -1), (0, 1), (1, 0)];

type Matrix = matrix::Matrix<u32>;
type Trailhead = matrix::Pos;

fn trailheads(m: &Matrix) -> Vec<Trailhead> {
    let mut result: Vec<Trailhead> = vec![];
    for i in 0..m.limit.0 {
        for j in 0..m.limit.1 {
            let pos = (i, j);
            if m.get(pos) == &0 {
                result.push(pos);
            }
        }
    }
    result
}

fn reachable(m: &Matrix, pos: matrix::Pos) -> HashSet<matrix::Pos> {
    let height = m.get(pos);
    AROUND_DELTAS
        .iter()
        .filter_map(|&d| m.pos_move(pos, d))
        .filter(|&p| *m.get(p) == height + 1)
        .collect()
}

fn score1(m: &Matrix, th: Trailhead) -> u32 {
    let mut heads: HashSet<matrix::Pos> = HashSet::new();
    heads.insert(th);
    let mut height = 0;
    while height < 9 {
        heads = heads.into_iter().flat_map(|th| reachable(m, th)).collect();
        height += 1;
    }
    heads.len() as u32
}

fn p1(input: &str) -> String {
    let m: Matrix = matrix::Matrix::new(input, |c| c.to_digit(10).unwrap());
    let result: u32 = trailheads(&m).iter().map(|&th| score1(&m, th)).sum();

    result.to_string()
}

fn score2(m: &Matrix, th: Trailhead) -> u32 {
    let mut heads: Vec<matrix::Pos> = vec![th];
    let mut height = 0;
    while height < 9 {
        heads = heads.into_iter().flat_map(|th| reachable(m, th)).collect();
        height += 1;
    }
    heads.len() as u32
}

fn p2(input: &str) -> String {
    let m: Matrix = matrix::Matrix::new(input, |c| c.to_digit(10).unwrap());
    let result: u32 = trailheads(&m).iter().map(|&th| score2(&m, th)).sum();

    result.to_string()
}

fn main() {
    aoc2024::run_day("10", Some(p1), Some(p2));
}