use aoc_runner_derive::{aoc, aoc_generator};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Cell {
    Empty,
    Wall,
}

type Input = ((usize, usize), (usize, usize), Vec<Vec<Cell>>);

#[aoc_generator(day20)]
fn parse(input: &str) -> Input {
    let mut start = (0, 0);
    let mut end = (0, 0);
    let map = input
        .lines()
        .enumerate()
        .map(|(y, line)| {
            line.chars()
                .enumerate()
                .map(|(x, c)| match c {
                    '.' => Cell::Empty,
                    '#' => Cell::Wall,
                    'S' => {
                        start = (x, y);
                        Cell::Empty
                    }
                    'E' => {
                        end = (x, y);
                        Cell::Empty
                    }
                    _ => panic!("unexpected character: {}", c),
                })
                .collect()
        })
        .collect();
    (start, end, map)
}

#[aoc(day20, part1)]
fn part1(input: &Input) -> usize {
    let start = input.0;
    let map = &input.2;

    #[cfg(not(test))]
    let min_saving = 100;
    #[cfg(test)]
    let min_saving = 2;

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

    let mut dist = vec![vec![usize::MAX; w]; h];
    dist[start.1][start.0] = 0;

    let mut deque = std::collections::VecDeque::new();
    deque.push_back((0, start));

    while let Some((d, pos)) = deque.pop_front() {
        if dist[pos.1][pos.0] < d {
            continue;
        }

        dist[pos.1][pos.0] = d;

        [(1, 0), (-1, 0), (0, -1), (0, 1)]
            .iter()
            .for_each(|(dx, dy)| {
                let x = (pos.0 as isize + dx) as usize;
                let y = (pos.1 as isize + dy) as usize;

                if x < w && y < h && map[y][x] == Cell::Empty {
                    let new_d = d + 1;
                    if new_d < dist[y][x] {
                        deque.push_back((new_d, (x, y)));
                    }
                }
            })
    }

    map.iter()
        .enumerate()
        .map(|(y, line)| {
            line.iter()
                .enumerate()
                .map(|(x, cell)| {
                    let y = y as isize;
                    if *cell == Cell::Empty {
                        [
                            (0, 2),
                            (0, -2),
                            (2, 0),
                            (-2, 0),
                            (1, 1),
                            (1, -1),
                            (-1, 1),
                            (-1, -1),
                        ]
                        .iter()
                        .filter(|(dx, dy)| {
                            let nx = x as isize + dx;
                            let ny = y + dy;
                            nx >= 0
                                && nx < w as isize
                                && ny >= 0
                                && ny < h as isize
                                && map[ny as usize][nx as usize] == Cell::Empty
                                && dist[ny as usize][nx as usize]
                                    >= dist[y as usize][x] + min_saving + 2
                        })
                        .count()
                    } else {
                        0
                    }
                })
                .sum::<usize>()
        })
        .sum()
}

#[aoc(day20, part2)]
fn part2(input: &Input) -> usize {
    let start = input.0;
    let map = &input.2;

    #[cfg(not(test))]
    let min_saving = 100;
    #[cfg(test)]
    let min_saving = 60;

    let max_skip = 20;

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

    let mut dist = vec![vec![usize::MAX; w]; h];
    dist[start.1][start.0] = 0;

    let mut deque = std::collections::VecDeque::new();
    deque.push_back((0, start));

    while let Some((d, pos)) = deque.pop_front() {
        if dist[pos.1][pos.0] < d {
            continue;
        }

        dist[pos.1][pos.0] = d;

        [(1, 0), (-1, 0), (0, -1), (0, 1)]
            .iter()
            .for_each(|(dx, dy)| {
                let x = (pos.0 as isize + dx) as usize;
                let y = (pos.1 as isize + dy) as usize;

                if x < w && y < h && map[y][x] == Cell::Empty {
                    let new_d = d + 1;
                    if new_d < dist[y][x] {
                        deque.push_back((new_d, (x, y)));
                    }
                }
            })
    }

    map.iter()
        .enumerate()
        .map(|(y, line)| {
            line.iter()
                .enumerate()
                .map(|(x, cell)| {
                    let y = y as isize;
                    if *cell == Cell::Empty {
                        (-max_skip..=max_skip)
                            .flat_map(|dx: isize| {
                                let skip = max_skip - dx.abs();
                                (-skip..=skip).map(move |dy| (dx, dy))
                            })
                            .filter(|(dx, dy)| {
                                let nx = x as isize + dx;
                                let ny = y + dy;
                                nx >= 0
                                    && nx < w as isize
                                    && ny >= 0
                                    && ny < h as isize
                                    && map[ny as usize][nx as usize] == Cell::Empty
                                    && dist[ny as usize][nx as usize]
                                        >= dist[y as usize][x]
                                            + min_saving
                                            + dx.unsigned_abs()
                                            + dy.unsigned_abs()
                            })
                            .count()
                    } else {
                        0
                    }
                })
                .sum::<usize>()
        })
        .sum()
}

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

    #[test]
    fn part1_example() {
        assert_eq!(
            part1(&parse(
                "###############
#...#...#.....#
#.#.#.#.#.###.#
#S#...#.#.#...#
#######.#.#.###
#######.#.#...#
#######.#.###.#
###..E#...#...#
###.#######.###
#...###...#...#
#.#####.#.###.#
#.#...#.#.#...#
#.#.#.#.#.#.###
#...#...#...###
###############"
            )),
            44
        );
    }

    #[test]
    fn part2_example() {
        assert_eq!(
            part2(&parse(
                "###############
#...#...#.....#
#.#.#.#.#.###.#
#S#...#.#.#...#
#######.#.#.###
#######.#.#...#
#######.#.###.#
###..E#...#...#
###.#######.###
#...###...#...#
#.#####.#.###.#
#.#...#.#.#...#
#.#.#.#.#.#.###
#...#...#...###
###############"
            )),
            129
        );
    }
}