use std::collections::{HashMap, HashSet};

fn parse(input: &str) -> Vec<Vec<u32>> {
    const RADIX: u32 = 10;
    input
        .lines()
        .map(|l| l.chars().map(|c| c.to_digit(RADIX).unwrap()).collect())
        .collect()
}

fn solve_1(input: &str) -> u32 {
    let map = parse(input);

    let h = map.len();
    let w = map[0].len();

    let mut map_info: HashMap<(usize, usize), HashSet<(usize, usize)>> = HashMap::new();

    for (j, row) in map.iter().enumerate() {
        for (i, elevation) in row.iter().enumerate() {
            if *elevation == 9 {
                map_info.insert((i, j), HashSet::from([(i, j)]));
            }
        }
    }

    for e in (0..=8).rev() {
        let mut new_map_info: HashMap<(usize, usize), HashSet<(usize, usize)>> = HashMap::new();

        for ((i, j), trailheads) in map_info.into_iter() {
            let mut positions = vec![];
            if i > 0 {
                positions.push((i - 1, j));
            }
            if i < w - 1 {
                positions.push((i + 1, j));
            }
            if j > 0 {
                positions.push((i, j - 1));
            }
            if j < h - 1 {
                positions.push((i, j + 1));
            }

            for (x, y) in positions.into_iter() {
                if map[y][x] == e {
                    match new_map_info.get(&(x, y)) {
                        None => new_map_info.insert((x, y), trailheads.clone()),
                        Some(&ref t) => new_map_info.insert(
                            (x, y),
                            t.union(&trailheads).map(|ts| ts.to_owned()).collect(),
                        ),
                    };
                }
            }
        }
        map_info = new_map_info;
    }

    map_info
        .values()
        .map(|trailheads| trailheads.len() as u32)
        .sum()
}

fn solve_2(input: &str) -> u32 {
    let map = parse(input);

    let h = map.len();
    let w = map[0].len();

    let mut map_info: HashMap<(usize, usize), u32> = HashMap::new();

    for (j, row) in map.iter().enumerate() {
        for (i, elevation) in row.iter().enumerate() {
            if *elevation == 9 {
                map_info.insert((i, j), 1);
            }
        }
    }

    for e in (0..=8).rev() {
        let mut new_map_info: HashMap<(usize, usize), u32> = HashMap::new();

        for ((i, j), rating) in map_info.into_iter() {
            let mut positions = vec![];
            if i > 0 {
                positions.push((i - 1, j));
            }
            if i < w - 1 {
                positions.push((i + 1, j));
            }
            if j > 0 {
                positions.push((i, j - 1));
            }
            if j < h - 1 {
                positions.push((i, j + 1));
            }

            for (x, y) in positions.into_iter() {
                if map[y][x] == e {
                    match new_map_info.get(&(x, y)) {
                        None => new_map_info.insert((x, y), rating),
                        Some(r) => new_map_info.insert((x, y), r + rating),
                    };
                }
            }
        }
        map_info = new_map_info;
    }

    map_info.values().sum()
}

fn main() {
    println!("Hello, this is Patrick!");

    let input = include_str!("../input.txt");

    let result_1 = solve_1(input);
    println!("The sum of scores of all trailheads is {}", result_1);

    let result_2 = solve_2(input);
    println!("The sum of ratings of all trailheads is {}", result_2);
}

#[cfg(test)]

mod tests {
    use super::*;

    #[test]
    fn test_1() {
        let test_input = include_str!("../test.txt");
        assert_eq!(solve_1(test_input), 36);
    }

    #[test]
    fn test_2() {
        let test_input = include_str!("../test.txt");
        assert_eq!(solve_2(test_input), 81);
    }
}