1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147

use crate::input::Input;
use std::cmp::max;
use std::collections::hash_map::Entry;
use std::collections::{HashMap, HashSet, VecDeque};

pub fn visit_rooms<F>(input_string: &str, mut callback: F) -> Result<(), String>
where
    F: FnMut(i32),
{
    if input_string.len() == 1 {
        return Err("Invalid one character input".to_string());
    }
    let input_string = &input_string[1..input_string.len() - 1];

    let apply_direction = |direction, position: &mut (i32, i32)| match direction {
        'N' => {
            position.1 -= 1;
        }
        'E' => {
            position.0 += 1;
        }
        'S' => {
            position.1 += 1;
        }
        'W' => {
            position.0 -= 1;
        }
        _ => {}
    };

    let mut room_doors = HashMap::new();

    let mut current_positions = HashSet::new();
    let mut positions_at_start_of_branch = Vec::new();
    let mut new_possibilities = Vec::new();

    current_positions.insert((0, 0));
    for char in input_string.chars() {
        match char {
            'N' | 'E' | 'S' | 'W' => {
                let old_positions = current_positions;
                current_positions = HashSet::new();
                for possibility in old_positions {
                    let mut possibility = possibility;
                    room_doors
                        .entry(possibility)
                        .or_insert_with(Vec::new)
                        .push(char);
                    apply_direction(char, &mut possibility);
                    let reverse_direction = match char {
                        'N' => 'S',
                        'E' => 'W',
                        'S' => 'N',
                        'W' => 'E',
                        _ => '?',
                    };
                    room_doors
                        .entry(possibility)
                        .or_insert_with(Vec::new)
                        .push(reverse_direction);
                    current_positions.insert(possibility);
                }
            }
            '(' => {
                positions_at_start_of_branch.push(current_positions.clone());
                new_possibilities.push(Vec::new());
            }
            '|' => {
                new_possibilities
                    .last_mut()
                    .ok_or("No possibility to push to for |")?
                    .push(current_positions);
                current_positions = positions_at_start_of_branch
                    .last_mut()
                    .ok_or("No position at start of branch for |")?
                    .clone();
            }
            ')' => {
                new_possibilities
                    .last_mut()
                    .ok_or("No possibility to push to for )")?
                    .push(current_positions);
                current_positions = HashSet::new();
                let nn: Vec<HashSet<(i32, i32)>> =
                    new_possibilities.pop().ok_or("No new possibility for )")?;
                for n in nn {
                    for e in n {
                        current_positions.insert(e);
                    }
                }

                positions_at_start_of_branch.pop();
            }
            _ => {
                return Err(format!("Invalid map tile: {}", char));
            }
        }
    }

    let mut visited = HashSet::new();
    let mut to_visit = VecDeque::new();
    to_visit.push_back((0_i32, 0, 0));

    while let Some(visiting) = to_visit.pop_front() {
        callback(visiting.0);

        if let Entry::Occupied(doors) = room_doors.entry((visiting.1, visiting.2)) {
            for char in doors.get() {
                let mut adjacent_room = (visiting.1, visiting.2);
                apply_direction(*char, &mut adjacent_room);
                if visited.insert(adjacent_room) {
                    to_visit.push_back((visiting.0 + 1, adjacent_room.0, adjacent_room.1));
                }
            }
        };
    }
    Ok(())
}

pub fn solve(input: &mut Input) -> Result<i32, String> {
    let mut result = 0;
    visit_rooms(input.text, |cost| {
        if input.is_part_one() {
            result = max(result, cost);
        } else if cost >= 1000 {
            result += 1;
        }
    })?;
    Ok(result)
}

#[test]
fn tests() {
    use crate::input::{test_part_one, test_part_two};

    test_part_one!("^WNE$" => 3);
    test_part_one!("^ENWWW(NEEE|SSE(EE|N))$" => 10);
    test_part_one!("^ENNWSWW(NEWS|)SSSEEN(WNSE|)EE(SWEN|)NNN$" => 18);
    test_part_one!("^(SSS|EEESSSWWW)ENNES$" => 8);
    test_part_one!("^(E|SSEENNW)S$" => 4);
    test_part_one!("^(E|SEN)$" => 2);
    test_part_one!("^NNNNN(EEEEE|NNN)NNNNN$" => 15);

    let input = include_str!("day20_input.txt");
    test_part_one!(input => 3151);
    test_part_two!(input => 8784);
}