// Critical point to note: EVERY free space is on the shortest path.

use itertools::Itertools;

use crate::search::dijkstra;
use crate::solver::DaySolver;
use crate::grid::{Coordinate, Grid};

type MyGrid = Grid<MazeElement>;

enum MazeElement {
    Wall,
    Free,
    Start,
    End,
}

impl MazeElement {
    fn is_free(&self) -> bool {
        !matches!(self, MazeElement::Wall)
    }
}

fn parse_input(input: String) -> (MyGrid, Coordinate) {
    let grid: MyGrid = input.lines()
        .map(|line| line.chars().map(|c| match c {
            '#' => MazeElement::Wall,
            '.' => MazeElement::Free,
            'S' => MazeElement::Start,
            'E' => MazeElement::End,
            _ => panic!("Invalid maze element: {}", c)
        })
             .collect())
        .collect::<Vec<Vec<MazeElement>>>()
        .into();

    let start_pos = grid.iter().find(|(_, me)| matches!(me, MazeElement::Start)).unwrap().0;

    (grid, start_pos)
}

fn solve<R>(grid: &MyGrid, start_pos: Coordinate, min_save_time: usize, in_range: R) -> usize
where R: Fn(Coordinate, Coordinate) -> bool {
    let (cost_to, _) = dijkstra(
        start_pos,
        |&c| grid.orthogonal_neighbors_iter(c)
            .filter(|&n| grid[n].is_free())
            .map(|n| (n, 1))
            .collect()
    );

    cost_to.keys()
        .cartesian_product(cost_to.keys())
        .map(|(&c1, &c2)| (c1, c2))
        // We don't compare with ourself
        .filter(|&(c1, c2)| c1 != c2)
        // The two points need to be within range
        .filter(|&(c1, c2)| in_range(c1, c2))
        // We need to save at least `min_save_time`
        .filter(|(c1, c2)| {
            // Because we are working with `usize`, the subtraction
            // could underflow. So we need to use `checked_sub`
            // instead, and check that a) no underflow happened, and
            // b) that the time saved is at least the minimum.
            cost_to.get(c2).copied()
                .and_then(|n| n.checked_sub(*cost_to.get(c1).unwrap()))
                .and_then(|n| n.checked_sub(c1.distance_to(c2)))
                .map(|n| n >= min_save_time)
                .unwrap_or(false)
        })
        .count()
}

pub struct Day20Solver;

impl DaySolver for Day20Solver {
    fn part1(&self, input: String) -> String {
        let (grid, start_pos) = parse_input(input);
        solve(
            &grid,
            start_pos,
            100,
            |c1, c2| c1.distance_to(&c2) == 2,
        ).to_string()
    }

    fn part2(&self, input: String) -> String {
        let (grid, start_pos) = parse_input(input);
        solve(
            &grid,
            start_pos,
            100,
            |c1, c2| c1.distance_to(&c2) <= 20,
        ).to_string()
    }
}

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

    #[test]
    fn test_part1() {
        let input = include_str!("../../inputs/test/20");
        let (grid, start_pos) = parse_input(input.to_string());
        let actual = solve(&grid, start_pos, 1, |c1, c2| c1.distance_to(&c2) == 2);
        assert_eq!(44, actual);
    }

    #[test]
    fn test_part2() {
        let input = include_str!("../../inputs/test/20");
        let (grid, start_pos) = parse_input(input.to_string());
        let actual = solve(&grid, start_pos, 50, |c1, c2| c1.distance_to(&c2) <= 20);
        assert_eq!(285, actual);
    }
}