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

static DAY: u8 = 10;

fn main() {
    let input = advent::read_lines(DAY);
    println!("{DAY}a: {}", Map::new(&input).trailhead_score());
    println!("{DAY}b: {}", Map::new(&input).distinct_trails());
}

#[derive(Eq, PartialEq, Hash, Copy, Clone)]
struct Position {
    x: isize,
    y: isize,
}

struct Map {
    map: HashMap<Position, u32>,
}

impl Map {
    fn new(input: &[String]) -> Map {
        let mut map = HashMap::new();
        for (y, line) in input.iter().enumerate() {
            for (x, c) in line.chars().enumerate() {
                let pos = Position { x: x as isize, y: y as isize };
                let height = c.to_digit(10).unwrap();
                map.insert(pos, height);
            }
        }
        Map { map }
    }

    fn trailheads(&self) -> Vec<Position> {
        self.map.iter()
                .filter(|&(_, height)| *height == 0)
                .map(|(pos, _)| *pos)
                .collect::<Vec<_>>()
    }

    fn reachable_peaks(&self, pos: &Position) -> HashSet<Position> {
        let mut peaks = HashSet::new();
        let height = *self.map.get(pos).unwrap();
        if height == 9 {
            peaks.insert(*pos);
            return peaks;
        }
        for (x, y) in [(-1, 0), (1, 0), (0, -1), (0, 1)] {
            let neighbor = Position { x: pos.x + x, y: pos.y + y };
            if let Some(height_neighbor) = self.map.get(&neighbor) {
                if *height_neighbor == height + 1 {
                    for peak in self.reachable_peaks(&neighbor) {
                        peaks.insert(peak);
                    }
                }
            }
        }
        peaks
    }

    fn count_trails(&self, pos: &Position) -> usize {
        let height = *self.map.get(pos).unwrap();
        if height == 9 {
            return 1;
        }
        let mut trails = 0;
        for (x, y) in [(-1, 0), (1, 0), (0, -1), (0, 1)] {
            let neighbor = Position { x: pos.x + x, y: pos.y + y };
            if let Some(height_neighbor) = self.map.get(&neighbor) {
                if *height_neighbor == height + 1 {
                    trails += self.count_trails(&neighbor);
                }
            }
        }
        trails
    }

    fn trailhead_score(&self) -> usize {
        self.trailheads().iter()
                         .map(|trailhead| self.reachable_peaks(trailhead).len())
                         .sum()
    }

    fn distinct_trails(&self) -> usize {
        self.trailheads().iter()
                         .map(|trailhead| self.count_trails(trailhead))
                         .sum()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test() {
        let input = [
          "89010123",
          "78121874",
          "87430965",
          "96549874",
          "45678903",
          "32019012",
          "01329801",
          "10456732",
        ].iter().map(|&x| String::from(x)).collect::<Vec<_>>();
        assert_eq!(Map::new(&input).trailhead_score(), 36);
        assert_eq!(Map::new(&input).distinct_trails(), 81);
    }
}