import Control.Arrow
import Data.Array.Unboxed
import Data.Functor
import Data.List
import Data.Map qualified as M
import Data.Set qualified as S

type Pos = (Int, Int)
type Board = Array Pos Char
type Path = M.Map Pos Int

parse board = listArray ((1, 1), (length l, length $ head l)) (concat l)
  where
    l = lines board

moves :: Pos -> [Pos]
moves p = [first succ p, first pred p, second succ p, second pred p]

getOrigin :: Board -> Maybe Pos
getOrigin = fmap fst . find ((== 'S') . snd) . assocs

getPath :: Board -> Pos -> [Pos]
getPath board p
    | not $ inRange (bounds board) p = []
    | board ! p == 'E' = [p]
    | board ! p == '#' = []
    | otherwise = p : (moves p >>= getPath (board // [(p, '#')]))

taxiCab (xa, ya) (xb, yb) = abs (xa - xb) + abs (ya - yb)

solve dist board = do
    path <- M.fromList . flip zip [1 ..] <$> (getOrigin board <&> getPath board)
    let positions = M.keys path
        jumps = [ (path M.! a) - (path M.! b) - d | a <- positions, b <- positions, d <- [taxiCab a b], d <= dist]
    return $ length $ filter (>=100) jumps

main = getContents >>= print . (solve 2 &&& solve 20) . parse

// 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);
    }
}

#include "common.h"

#define GB 2	/* safety border on X/Y plane */
#define GZ (GB+141+2+GB)

static char G[GZ][GZ];		/* grid */
static int  C[GZ][GZ];		/* costs */
static int sx,sy, ex,ey;	/* start, end pos */

static void
flood(int x, int y)
{
	int lo = INT_MAX;

	if (x<1 || x>=GZ-1 ||
	    y<1 || y>=GZ-1 || G[y][x]!='.')
		return;

	if (C[y-1][x]) lo = MIN(lo, C[y-1][x]+1);
	if (C[y+1][x]) lo = MIN(lo, C[y+1][x]+1);
	if (C[y][x-1]) lo = MIN(lo, C[y][x-1]+1);
	if (C[y][x+1]) lo = MIN(lo, C[y][x+1]+1);

	if (lo != INT_MAX && (!C[y][x] || lo < C[y][x])) {
		C[y][x] = lo;

		flood(x, y-1); flood(x-1, y);
		flood(x, y+1); flood(x+1, y);
	}
}

static int
count_shortcuts(int lim, int min)
{
	int acc=0, x,y, dx,dy;

	for (y=GB; y<GZ-GB; y++)
	for (x=GB; x<GZ-GB; x++)
	for (dy = -lim; dy <= lim; dy++)
	for (dx = abs(dy)-lim; dx <= lim-abs(dy); dx++)
		acc += x+dx >= 0 && x+dx < GZ &&
		       y+dy >= 0 && y+dy < GZ && C[y][x] &&
		       C[y][x]+abs(dx)+abs(dy) <= C[y+dy][x+dx]-min;

	return acc;
}

int
main(int argc, char **argv)
{
	int x,y;

	if (argc > 1)
		DISCARD(freopen(argv[1], "r", stdin));
	for (y=2; fgets(G[y]+GB, GZ-GB*2, stdin); y++)
		assert(y < GZ-3);

	for (y=GB; y<GZ-GB; y++)
	for (x=GB; x<GZ-GB; x++)
		if (G[y][x] == 'S') { G[y][x]='.'; sx=x; sy=y; } else
		if (G[y][x] == 'E') { G[y][x]='.'; ex=x; ey=y; }

	C[sy][sx] = 1;

	flood(sx, sy-1); flood(sx-1, sy);
	flood(sx, sy+1); flood(sx+1, sy);

	printf("20: %d %d\n",
	    count_shortcuts(2, 100),
	    count_shortcuts(20, 100));

	return 0;
}

#include <boost/log/trivial.hpp>
#include <boost/unordered/unordered_flat_map.hpp>
#include <boost/unordered/unordered_flat_set.hpp>
#include <cstddef>
#include <fstream>
#include <limits>
#include <queue>
#include <stdexcept>

namespace {

using Loc = std::pair<int, int>;
using Dir = std::pair<int, int>;
template <class T>
using Score = std::pair<size_t, T>;
template <class T>
using MinHeap = std::priority_queue<Score<T>, std::vector<Score<T>>, std::greater<Score<T>>>;
using Map = boost::unordered_flat_set<Loc>;

auto operator+(const Loc &l, const Dir &d) {
  return Loc{l.first + d.first, l.second + d.second};
}

auto manhattan(const Loc &a, const Loc &b) {
  return std::abs(a.first - b.first) + std::abs(a.second - b.second);
}

auto dirs = std::vector<Dir>{
    {0,  -1},
    {0,  1 },
    {-1, 0 },
    {1,  0 }
};

struct Maze {
  Map map;
  Loc start;
  Loc end;
};

auto parse() {
  auto rval = Maze{};
  auto line = std::string{};
  auto ih = std::ifstream{"input/20"};
  auto row = 0;
  while (std::getline(ih, line)) {
    for (auto col = 0; col < int(line.size()); ++col) {
      auto t = line.at(col);
      switch (t) {
      case 'S':
        rval.start = Loc{col, row};
        rval.map.insert(Loc{col, row});
        break;
      case 'E':
        rval.end = Loc{col, row};
        rval.map.insert(Loc{col, row});
        break;
      case '.':
        rval.map.insert(Loc{col, row});
        break;
      case '#':
        break;
      default:
        throw std::runtime_error{"oops"};
      }
    }
    ++row;
  }
  return rval;
}

auto dijkstra(const Maze &m) {
  auto unvisited = MinHeap<Loc>{};
  auto visited = boost::unordered_flat_map<Loc, size_t>{};

  for (const auto &e : m.map)
    visited[e] = std::numeric_limits<size_t>::max();

  visited[m.start] = 0;
  unvisited.push({0, {m.start}});

  while (!unvisited.empty()) {
    auto next = unvisited.top();
    unvisited.pop();

    if (visited.at(next.second) < next.first)
      continue;

    for (const auto &dir : dirs) {
      auto prospective = Loc{next.second + dir};
      if (!visited.contains(prospective))
        continue;
      auto pscore = next.first + 1;
      if (visited.at(prospective) > pscore) {
        visited[prospective] = pscore;
        unvisited.push({pscore, prospective});
      }
    }
  }

  return visited;
}

using Walk = decltype(dijkstra(Maze{}));

constexpr auto GOOD_CHEAT = 100;

auto evaluate_cheats(const Walk &walk, int skip) {
  auto rval = size_t{};
  for (auto &start : walk) {
    for (auto &end : walk) {
      auto distance = manhattan(start.first, end.first);
      if (distance <= skip && end.second > start.second) {
        auto improvement = int(end.second) - int(start.second) - distance;
        if (improvement >= GOOD_CHEAT)
          ++rval;
      }
    }
  }
  return rval;
}

} // namespace

auto main() -> int {
  auto p = parse();
  auto walk = dijkstra(p);
  BOOST_LOG_TRIVIAL(info) << "01: " << evaluate_cheats(walk, 2);
  BOOST_LOG_TRIVIAL(info) << "02: " << evaluate_cheats(walk, 20);
}

use euclid::{default::*, vec2};

const DIRS: [Vector2D<i32>; 4] = [vec2(1, 0), vec2(0, 1), vec2(-1, 0), vec2(0, -1)];
const MIN_SAVE: u32 = 100;
const MAX_DIST: i32 = 20;

fn parse(input: &str) -> (Vec<Vec<bool>>, Point2D<i32>, Point2D<i32>) {
    let mut start = None;
    let mut end = None;
    let mut field = Vec::new();
    for (y, line) in input.lines().enumerate() {
        let mut row = Vec::new();
        for (x, b) in line.bytes().enumerate() {
            row.push(b == b'#');
            if b == b'S' {
                start = Some(Point2D::new(x, y).to_i32());
            } else if b == b'E' {
                end = Some(Point2D::new(x, y).to_i32());
            }
        }
        field.push(row);
    }
    (field, start.unwrap(), end.unwrap())
}

fn distances(
    field: &[Vec<bool>],
    start: Point2D<i32>,
    end: Point2D<i32>,
) -> (Vec<Vec<u32>>, Vec<Vec<u32>>) {
    let width = field[0].len();
    let height = field.len();
    let mut dist_to_start = vec![vec![u32::MAX; width]; height];
    let bounds = Rect::new(Point2D::origin(), Size2D::new(width, height)).to_i32();

    let mut cur = start;
    let mut dist = 0;
    dist_to_start[cur.y as usize][cur.x as usize] = dist;
    while cur != end {
        for dir in DIRS {
            let next = cur + dir;
            if bounds.contains(next)
                && !field[next.y as usize][next.x as usize]
                && dist_to_start[next.y as usize][next.x as usize] == u32::MAX
            {
                cur = next;
                break;
            }
        }
        dist += 1;
        dist_to_start[cur.y as usize][cur.x as usize] = dist;
    }
    let total_dist = dist_to_start[end.y as usize][end.x as usize];
    let dist_to_end = dist_to_start
        .iter()
        .map(|row| {
            row.iter()
                .map(|&d| {
                    if d == u32::MAX {
                        u32::MAX
                    } else {
                        total_dist - d
                    }
                })
                .collect()
        })
        .collect();
    (dist_to_start, dist_to_end)
}

fn cheats(
    field: &[Vec<bool>],
    dist_to_start: &[Vec<u32>],
    dist_to_end: &[Vec<u32>],
    total_dist: u32,
) -> u32 {
    let width = field[0].len();
    let height = field.len();
    let bounds = Rect::new(Point2D::origin(), Size2D::new(width, height)).to_i32();
    let mut count = 0;
    for (y, row) in field.iter().enumerate() {
        for (x, _w) in row.iter().enumerate().filter(|&(_i, w)| *w) {
            let pos = Point2D::new(x, y).to_i32();
            for (d0, &dir0) in DIRS.iter().enumerate().skip(1) {
                for &dir1 in DIRS.iter().take(d0) {
                    let p0 = pos + dir0;
                    let p1 = pos + dir1;
                    if bounds.contains(p0) && bounds.contains(p1) {
                        let p0 = p0.to_usize();
                        let p1 = p1.to_usize();
                        if !field[p0.y][p0.x] && !field[p1.y][p1.x] {
                            let dist = dist_to_start[p0.y][p0.x].min(dist_to_start[p1.y][p1.x])
                                + dist_to_end[p1.y][p1.x].min(dist_to_end[p0.y][p0.x])
                                + 2; // Add 2 for cutting across the wall
                            if total_dist - dist >= MIN_SAVE {
                                count += 1;
                            }
                        }
                    }
                }
            }
        }
    }
    count
}

fn part1(input: String) {
    let (field, start, end) = parse(&input);
    let (dist_to_start, dist_to_end) = distances(&field, start, end);
    let total_dist = dist_to_start[end.y as usize][end.x as usize];
    println!(
        "{}",
        cheats(&field, &dist_to_start, &dist_to_end, total_dist)
    );
}

// Half of all vectors with manhattan distance <= MAX_DIST.
// Only vectors with positive x or going straight down are considered to avoid using the same
// vector twice in both directions.
fn cheat_vectors() -> Vec<Vector2D<i32>> {
    let mut vectors = Vec::new();
    for y in -MAX_DIST..=MAX_DIST {
        let start = if y > 0 { 0 } else { 1 };
        for x in start..=(MAX_DIST - y.abs()) {
            assert!(x + y <= MAX_DIST);
            vectors.push(vec2(x, y));
        }
    }
    vectors
}

fn cheats20(
    field: &[Vec<bool>],
    dist_to_start: &[Vec<u32>],
    dist_to_end: &[Vec<u32>],
    total_dist: u32,
) -> u32 {
    let vectors = cheat_vectors();
    let width = field[0].len();
    let height = field.len();
    let bounds = Rect::new(Point2D::origin(), Size2D::new(width, height)).to_i32();
    let mut count = 0;
    for (y, row) in field.iter().enumerate() {
        for (x, _w) in row.iter().enumerate().filter(|&(_i, w)| !*w) {
            let p0 = Point2D::new(x, y);
            for &v in &vectors {
                let pi1 = p0.to_i32() + v;
                if bounds.contains(pi1) {
                    let p1 = pi1.to_usize();
                    if !field[p1.y][p1.x] {
                        let dist = dist_to_start[p0.y][p0.x].min(dist_to_start[p1.y][p1.x])
                            + dist_to_end[p1.y][p1.x].min(dist_to_end[p0.y][p0.x])
                            + v.x.unsigned_abs()  // Manhattan distance of vector
                            + v.y.unsigned_abs();
                        if total_dist - dist >= MIN_SAVE {
                            count += 1;
                        }
                    }
                }
            }
        }
    }
    count
}

fn part2(input: String) {
    let (field, start, end) = parse(&input);
    let (dist_to_start, dist_to_end) = distances(&field, start, end);
    let total_dist = dist_to_start[end.y as usize][end.x as usize];
    println!(
        "{}",
        cheats20(&field, &dist_to_start, &dist_to_end, total_dist)
    );
}

util::aoc_main!();