use std::collections::HashMap;

fn main() {
    let input = advent::read_file(3);
    println!("3a: {}", houses_visited(&input));
    println!("3b: {}", houses_visited_robo(&input));
}

fn new_pos(pos: (i32, i32), c: char) -> (i32, i32) {
    match c {
        'v' => (pos.0, pos.1 - 1),
        '^' => (pos.0, pos.1 + 1),
        '<' => (pos.0 - 1, pos.1),
        '>' => (pos.0 + 1, pos.1),
        _ => panic!("invalid character"),
    }
}

fn houses_visited(input: &str) -> usize {
    let mut map = HashMap::new();
    let mut pos = (0, 0);
    map.insert(pos, 1);

    for c in input.chars() {
        pos = new_pos(pos, c);
        let count = map.entry(pos).or_insert(0);
        *count += 1;
    }

    map.keys().len()
}

fn houses_visited_robo(input: &str) -> usize {
    let mut map = HashMap::new();
    let mut pos_santa = (0, 0);
    let mut pos_robo = (0, 0);

    map.insert(pos_santa, 2);

    for c in input.chars().step_by(2) {
        pos_santa = new_pos(pos_santa, c);

        let count = map.entry(pos_santa).or_insert(0);
        *count += 1;
    }
    for c in input.chars().skip(1).step_by(2) {
        pos_robo = new_pos(pos_robo, c);

        let count = map.entry(pos_robo).or_insert(0);
        *count += 1;
    }

    map.keys().len()
}

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

    #[test]
    fn test() {
        assert_eq!(houses_visited(">"), 2);
        assert_eq!(houses_visited("^>v<"), 4);
        assert_eq!(houses_visited("^v^v^v^v^v"), 2);
    }

    #[test]
    fn test_robo() {
        assert_eq!(houses_visited_robo("^v"), 3);
        assert_eq!(houses_visited_robo("^>v<"), 3);
        assert_eq!(houses_visited_robo("^v^v^v^v^v"), 11);
    }
}