#include "common.h"

#define GZ 145

enum {NN, EE, SS, WW};

static const int dx[]={0,1,0,-1}, dy[]={-1,0,1,0};

static char g[GZ][GZ];		/* with 1 tile border */
static int cost[GZ][GZ][4];	/* per direction, starts at 1, 0=no info */

static int traversible(char c) { return c=='.' || c=='S' || c=='E'; }

static int
minat(int x, int y)
{
	int acc=0, d;

	for (d=0; d<4; d++)
		if (cost[y][x][d] && (!acc || cost[y][x][d] < acc))
			acc = cost[y][x][d];

	return acc;
}


static int
count_exits(int x, int y)
{
	int acc=0, i;

	assert(x>0); assert(x<GZ-1);
	assert(y>0); assert(y<GZ-1);

	for (i=0; i<4; i++)
		acc += traversible(g[y+dy[i]][x+dx[i]]);
	
	return acc;
}

/* remove all dead ends */
static void
prune_dead(void)
{
	int dirty=1, x,y;

	while (dirty) {
		dirty = 0;

		for (y=1; y<GZ-1; y++)
		for (x=1; x<GZ-1; x++)
			if (g[y][x]=='.' && count_exits(x,y) < 2)
				{ dirty = 1; g[y][x] = '#'; }
	}
}

/* remove all dead ends from cost[], leaves only optimal paths */
static void
prune_subopt(void)
{
	int dirty=1, x,y,d;

	while (dirty) {
		dirty = 0;

		for (y=1; y<GZ-1; y++)
		for (x=1; x<GZ-1; x++)
		for (d=0; d<4; d++) {
			if (!cost[y][x][d])
				continue;

			if (g[y][x]=='E') {
				if (cost[y][x][d] != minat(x,y))
					{ dirty = 1; cost[y][x][d] = 0; }
				continue;
			}

			if (cost[y][x][d]+1 > cost[y+dy[d]][x+dx[d]][d] &&
			    cost[y][x][d]+1000 > cost[y][x][(d+1)%4] &&
			    cost[y][x][d]+1000 > cost[y][x][(d+3)%4])
				{ dirty = 1; cost[y][x][d] = 0; }
		}
	}
}

static void
propagate_costs(void)
{
	int dirty=1, cost1, x,y,d;

	while (dirty) {
		dirty = 0;

		for (y=1; y<GZ-1; y++)
		for (x=1; x<GZ-1; x++)
		for (d=0; d<4; d++) {
			if (!traversible(g[y][x]))
				continue;

			/* from back */
			if ((cost1 = cost[y-dy[d]][x-dx[d]][d]) &&
			    (cost1+1 < cost[y][x][d] || !cost[y][x][d]))
				{ dirty = 1; cost[y][x][d] = cost1+1; }

			/* from right */
			if ((cost1 = cost[y][x][(d+1)%4]) &&
			    (cost1+1000 < cost[y][x][d] || !cost[y][x][d]))
				{ dirty = 1; cost[y][x][d] = cost1+1000; }

			/* from left */
			if ((cost1 = cost[y][x][(d+3)%4]) &&
			    (cost1+1000 < cost[y][x][d] || !cost[y][x][d]))
				{ dirty = 1; cost[y][x][d] = cost1+1000; }
		}
	}
}

int
main(int argc, char **argv)
{
	int p1=0,p2=0, sx=0,sy=0, ex=0,ey=0, x,y;
	char *p;

	if (argc > 1)
		DISCARD(freopen(argv[1], "r", stdin));

	for (y=1; fgets(g[y]+1, GZ-1, stdin); y++) {
		if ((p = strchr(g[y]+1, 'S'))) { sy=y; sx=p-g[y]; }
		if ((p = strchr(g[y]+1, 'E'))) { ey=y; ex=p-g[y]; }
		assert(y+1 < GZ-1);
	}

	cost[sy][sx][EE] = 1;

	prune_dead();
	propagate_costs();
	prune_subopt();

	p1 = minat(ex, ey) -1;	/* costs[] values start at 1! */

	for (y=1; y<GZ-1; y++)
	for (x=1; x<GZ-1; x++)
		p2 += minat(x,y) > 0;

	printf("16: %d %d\n", p1, p2);
	return 0;
}

import Control.Arrow
import Control.Monad
import Control.Monad.RWS
import Control.Monad.Trans.Maybe
import Data.Array.Unboxed
import Data.List
import Data.Map qualified as M
import Data.Maybe
import Data.Set qualified as S

data Dir = N | S | W | E deriving (Show, Eq, Ord)
type Maze = UArray Pos Char
type Pos = (Int, Int)
type Node = (Pos, Dir)
type CostNode = (Int, Node)
type Problem = RWS Maze [(Node, [Node])] (M.Map Node Int, S.Set (CostNode, Maybe Node))

parse = toMaze . lines

toMaze :: [String] -> Maze
toMaze b = listArray ((0, 0), (n - 1, m - 1)) $ concat b
  where
    n = length b
    m = length $ head b

next :: Int -> (Pos, Dir) -> Problem [CostNode]
next c (p, d) = do
    m <- ask

    let straigth = fmap ((1,) . (,d)) . filter ((/= '#') . (m !)) . return $ move d p
        turn = (1000,) . (p,) <$> rot d

    return $ first (+ c) <$> straigth ++ turn

move N = first (subtract 1)
move S = first (+ 1)
move W = second (subtract 1)
move E = second (+ 1)

rot d
    | d `elem` [N, S] = [E, W]
    | otherwise = [N, S]

dijkstra :: MaybeT Problem ()
dijkstra = do
    m <- ask
    visited <- gets fst
    Just (((cost, vertex@(p, _)), father), queue) <- gets (S.minView . snd)

    let (prevCost, visited') = M.insertLookupWithKey (\_ a _ -> a) vertex cost visited

    case prevCost of
        Nothing -> do
            queue' <- lift $ foldr S.insert queue <$> (fmap (,Just vertex) <$> next cost vertex)
            put (visited', queue')
            tell [(vertex, maybeToList father)]
        Just c -> do
            if c == cost
                then tell [(vertex, maybeToList father)]
                else guard $ m ! p /= 'E'
            put (visited, queue)
    dijkstra

solve b = do
    start <- getStart b
    end <- getEnd b
    let ((m, _), w) = execRWS (runMaybeT dijkstra) b (M.empty, S.singleton (start, Nothing))
        parents = M.fromListWith (++) w
        endDirs = (end,) <$> [N, S, E, W]
        min = minimum $ mapMaybe (`M.lookup` m) endDirs
        ends = filter ((== Just min) . (`M.lookup` m)) endDirs
        part2 =
            S.size . S.fromList . fmap fst . concat . takeWhile (not . null) $
                iterate (>>= flip (M.findWithDefault []) parents) ends
    return (min, part2)

getStart :: Maze -> Maybe CostNode
getStart = fmap ((0,) . (,E) . fst) . find ((== 'S') . snd) . assocs

getEnd :: Maze -> Maybe Pos
getEnd = fmap fst . find ((== 'E') . snd) . assocs

main = getContents >>= print . solve . parse

import Control.Arrow

import Data.Map (Map)
import Data.Set (Set)
import Data.Array.ST (STArray)
import Data.Array (Array)
import Control.Monad.ST (ST, runST)

import qualified Data.Char as Char
import qualified Data.List as List
import qualified Data.Map as Map
import qualified Data.Set as Set
import qualified Data.Array.ST as MutableArray
import qualified Data.Array as Array
import qualified Data.Maybe as Maybe

data Direction = East | West | South | North
        deriving (Show, Eq, Ord)
data MazeTile = Start | End | Wall | Unknown | Explored (Map Direction ExplorationScore)
        deriving Eq

--      instance Show MazeTile where
--              show Wall = "#"
--              show Start = "S"
--              show End = "E"
--              show Unknown = "."
--              show (Explored (East, _))  = ">"
--              show (Explored (South, _)) = "v"
--              show (Explored (West, _))  = "<"
--              show (Explored (North, _)) = "^"

type Position = (Int, Int)
type ExplorationScore = Int

translate '#' = Wall
translate '.' = Unknown
translate 'S' = Start
translate 'E' = End

parse :: String -> Array (Int, Int) MazeTile
parse s = Array.listArray ((1, 1), (height - 1, width)) . map translate . filter (/= '\n') $ s
        where
                width = length . takeWhile (/= '\n') $ s
                height = length . filter (== '\n') $ s

(a1, b1) .+. (a2, b2) = (a1+a2, b1+b2)
(a1, b1) .-. (a2, b2) = (a1-a2, b1-b2)

directions = [East, West, South, North]
directionVector East  = (0,  1)
directionVector West  = (0, -1)
directionVector North = (-1, 0)
directionVector South = ( 1, 0)

turnRight East  = South
turnRight South = West
turnRight West  = North
turnRight North = East

walkableNeighbors a p = do
        let neighbors = List.map ((.+. p) . directionVector) directions
        tiles <- mapM (MutableArray.readArray a) neighbors
        let neighborPosition = List.map fst . List.filter ((/= Wall). snd) . zip neighbors $ tiles
        return $ neighborPosition


findDeadEnds a = Array.assocs
        >>> List.filter (snd >>> (== Unknown))
        >>> List.map (fst)
        >>> List.filter (isDeadEnd a)
        $ a
isDeadEnd a p = List.map directionVector
        >>> List.map (.+. p)
        >>> List.map (a Array.!)
        >>> List.filter (/= Wall)
        >>> List.length
        >>> (== 1)
        $ directions

fillDeadEnds :: Array (Int, Int) MazeTile -> ST s (Array (Int, Int) MazeTile)
fillDeadEnds a = do
        ma <- MutableArray.thaw a
        let deadEnds = findDeadEnds a
        mapM_ (fillDeadEnd ma) deadEnds
        MutableArray.freeze ma

fillDeadEnd :: STArray s (Int, Int) MazeTile -> Position -> ST s ()
fillDeadEnd a p = do
        MutableArray.writeArray a p Wall
        p' <- walkableNeighbors a p >>= return . head
        t <- MutableArray.readArray a p'
        n <- walkableNeighbors a p' >>= return . List.length
        if n == 1 && t == Unknown then fillDeadEnd a p' else return ()

thawArray :: Array (Int, Int) MazeTile -> ST s (STArray s (Int, Int) MazeTile)
thawArray a = do
        a' <- MutableArray.thaw a
        return a'

solveMaze a = do
        a' <- fillDeadEnds a
        a'' <- thawArray a'
        let s = Array.assocs
                >>> List.filter ((== Start) . snd)
                >>> Maybe.listToMaybe
                >>> Maybe.maybe (error "Start not in map") fst
                $ a
        let e = Array.assocs
                >>> List.filter ((== End) . snd)
                >>> Maybe.listToMaybe
                >>> Maybe.maybe (error "End not in map") fst
                $ a
        MutableArray.writeArray a'' s $ Explored (Map.singleton East 0)
        MutableArray.writeArray a'' e $ Unknown
        solveMaze' (s, East) a''
        fa <- MutableArray.freeze a''
        t <- MutableArray.readArray a'' e
        case t of
                Wall  -> error "Unreachable code"
                Start -> error "Unreachable code"
                End   -> error "Unreachable code"
                Unknown -> error "End was not explored yet"
                Explored m -> return (List.minimum . List.map snd . Map.toList $ m, countTiles fa s e)

countTiles a s p = Set.size . countTiles' a s p $ South

countTiles' :: Array (Int, Int) MazeTile -> Position -> Position -> Direction -> Set Position
countTiles' a s p d
        | p == s    = Set.singleton p
        | otherwise = Set.unions 
                . List.map (Set.insert p) 
                . List.map (uncurry (countTiles' a s)) 
                $ (zip minCostNeighbors minCostDirections)
        where
                minCostNeighbors   = List.map ((p .-.) . directionVector) minCostDirections
                minCostDirections  = List.map fst . List.filter ((== minCost) . snd) . Map.toList $ visits
                visits = case a Array.! p of
                        Explored m -> Map.adjust (+ (-1000)) d m
                minCost = List.minimum . List.map snd . Map.toList $ visits

maybeExplore c p d a = do
        t <- MutableArray.readArray a p
        case t of
                Wall     -> return ()
                Start    -> error "Unreachable code"
                End      -> error "Unreachable code"
                Unknown  -> do
                        MutableArray.writeArray a p $ Explored (Map.singleton d c)
                        solveMaze' (p, d) a
                Explored m -> do
                        let c' = Maybe.maybe c id (m Map.!? d)
                        if c <= c' then do
                                let m' = Map.insert d c m
                                MutableArray.writeArray a p (Explored m')
                                solveMaze' (p, d) a
                        else
                                return ()

solveMaze' :: (Position, Direction) -> STArray s (Int, Int) MazeTile -> ST s ()
solveMaze' s@(p, d) a = do
        t <- MutableArray.readArray a p
        case t of
                Wall -> return ()
                Start -> error "Unreachable code"
                End -> error "Unreachable code"
                Unknown -> error "Starting on unexplored field"
                Explored m -> do
                        let c = m Map.! d
                        maybeExplore (c+1)    (p .+. directionVector d)  d a
                        let d' = turnRight d
                        maybeExplore (c+1001) (p .+. directionVector d') d' a
                        let d'' = turnRight d'
                        maybeExplore (c+1001) (p .+. directionVector d'') d'' a
                        let d''' = turnRight d''
                        maybeExplore (c+1001) (p .+. directionVector d''') d''' a

part1 a = runST (solveMaze a)

main = getContents
        >>= print
        . part1
        . parse

#[cfg(test)]
mod tests {
    use crate::day_16::tests::State::{CELL, END, SHORTPATH, START, WALL};
    use std::cmp::PartialEq;

    fn get_cell(board: &[Vec<MazeCell>], row: isize, col: isize) -> &MazeCell {
        &board[row as usize][col as usize]
    }

    fn set_cell(board: &mut [Vec<MazeCell>], value: &MazeStep) {
        let cell = &mut board[value.i as usize][value.j as usize];
        cell.dir = value.dir;
        cell.cost = value.cost;
        cell.state = value.state.clone();
    }

    fn find_cell(board: &mut [Vec<MazeCell>], state: State) -> (isize, isize) {
        for i in 0..board.len() {
            for j in 0..board[i].len() {
                if get_cell(board, i as isize, j as isize).state == state {
                    return (i as isize, j as isize);
                }
            }
        }
        unreachable!();
    }

    static DIRECTIONS: [(isize, isize); 4] = [(0, 1), (1, 0), (0, -1), (-1, 0)];

    #[derive(PartialEq, Debug, Clone)]
    enum State {
        CELL,
        WALL,
        START,
        END,
        SHORTPATH,
    }

    struct MazeCell {
        dir: i8,
        cost: isize,
        state: State,
    }

    struct MazeStep {
        i: isize,
        j: isize,
        dir: i8,
        cost: isize,
        state: State,
    }

    fn walk_maze(board: &mut [Vec<MazeCell>]) -> isize {
        let start = find_cell(board, START);
        let mut moves = vec![MazeStep {
            i: start.0,
            j: start.1,
            cost: 0,
            dir: 0,
            state: START,
        }];

        let mut best = isize::MAX;

        loop {
            if moves.is_empty() {
                break;
            }
            let cell = moves.pop().unwrap();
            let current_cost = get_cell(board, cell.i, cell.j);
            if current_cost.state == END {
                if cell.cost < best {
                    best = cell.cost;
                }
                continue;
            }
            if current_cost.state == WALL {
                continue;
            }
            if current_cost.cost < cell.cost {
                continue;
            }
            set_cell(board, &cell);

            for (i, dir) in DIRECTIONS.iter().enumerate() {
                let cost = match (i as i8) - cell.dir {
                    0 => cell.cost + 1,
                    -2 | 2 => continue,
                    _ => cell.cost + 1001,
                };
                moves.push(MazeStep {
                    i: cell.i + dir.0,
                    j: cell.j + dir.1,
                    dir: i as i8,
                    cost,
                    state: State::CELL,
                });
            }
        }

        best
    }

    fn unwalk_path(board: &mut [Vec<MazeCell>], total_cost: isize) -> usize {
        let end = find_cell(board, END);

        let mut cells = vec![MazeStep {
            i: end.0,
            j: end.1,
            dir: 0,
            cost: total_cost,
            state: State::END,
        }];

        set_cell(board, &cells[0]);

        while let Some(mut cell) = cells.pop() {
            for dir in DIRECTIONS {
                let next_cell = get_cell(board, cell.i + dir.0, cell.j + dir.1);
                if next_cell.cost == 0 {
                    continue;
                }
                if next_cell.state == WALL {
                    continue;
                }
                if next_cell.state == CELL
                    && (next_cell.cost == &cell.cost - 1001 || next_cell.cost == &cell.cost - 1)
                {
                    cells.push(MazeStep {
                        i: cell.i + dir.0,
                        j: cell.j + dir.1,
                        dir: 0,
                        cost: next_cell.cost,
                        state: CELL,
                    });
                } else {
                    let prev_cell = get_cell(board, cell.i - dir.0, cell.j - dir.1);
                    if prev_cell.state == SHORTPATH && prev_cell.cost - 2 == next_cell.cost {
                        cells.push(MazeStep {
                            i: cell.i + dir.0,
                            j: cell.j + dir.1,
                            dir: 0,
                            cost: next_cell.cost,
                            state: CELL,
                        });
                    }
                }
            }
            cell.state = SHORTPATH;
            set_cell(board, &cell);
        }
        let mut count = 0;
        for row in board {
            for cell in row {
                if cell.state == SHORTPATH {
                    count += 1;
                }
                if cell.state == END {
                    count += 1;
                }
                if cell.state == START {
                    count += 1;
                }
            }
        }
        count
    }

    #[test]
    fn day15_part2_test() {
        let input = std::fs::read_to_string("src/input/day_16.txt").unwrap();

        let mut board = input
            .split('\n')
            .map(|line| {
                line.chars()
                    .map(|c| match c {
                        '#' => MazeCell {
                            dir: 0,
                            cost: isize::MAX,
                            state: WALL,
                        },
                        'S' => MazeCell {
                            dir: 0,
                            cost: isize::MAX,
                            state: START,
                        },
                        'E' => MazeCell {
                            dir: 0,
                            cost: isize::MAX,
                            state: END,
                        },
                        '.' => MazeCell {
                            dir: 0,
                            cost: isize::MAX,
                            state: CELL,
                        },
                        _ => unreachable!(),
                    })
                    .collect::<Vec<MazeCell>>()
            })
            .collect::<Vec<Vec<MazeCell>>>();

        let cost = walk_maze(&mut board);

        let count = unwalk_path(&mut board, cost);

        println!("{count}");
    }
}

use std::cmp::{Ordering, Reverse};

use euclid::{default::*, vec2};
use priority_queue::PriorityQueue;
use rustc_hash::{FxHashMap, FxHashSet};

const DIRS: [Vector2D<i32>; 4] = [vec2(1, 0), vec2(0, 1), vec2(-1, 0), vec2(0, -1)];

type Node = (Point2D<i32>, u8);

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, l) in input.lines().enumerate() {
        let mut row = Vec::new();
        for (x, b) in l.bytes().enumerate() {
            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());
            }
            row.push(b == b'#');
        }
        field.push(row);
    }
    (field, start.unwrap(), end.unwrap())
}

fn adj(field: &[Vec<bool>], (v, dir): Node) -> Vec<(Node, u32)> {
    let mut adj = Vec::with_capacity(3);
    let next = v + DIRS[dir as usize];
    if !field[next.y as usize][next.x as usize] {
        adj.push(((next, dir), 1));
    }
    adj.push(((v, (dir + 1) % 4), 1000));
    adj.push(((v, (dir + 3) % 4), 1000));
    adj
}

fn shortest_path_length(field: &[Vec<bool>], start: Node, end: Point2D<i32>) -> u32 {
    let mut dist: FxHashMap<Node, u32> = FxHashMap::default();
    dist.insert(start, 0);
    let mut pq: PriorityQueue<Node, Reverse<u32>> = PriorityQueue::new();
    pq.push(start, Reverse(0));
    while let Some((v, _)) = pq.pop() {
        for (w, weight) in adj(field, v) {
            let dist_w = dist.get(&w).copied().unwrap_or(u32::MAX);
            let new_dist = dist[&v] + weight;
            if dist_w > new_dist {
                dist.insert(w, new_dist);
                pq.push_increase(w, Reverse(new_dist));
            }
        }
    }
    // Shortest distance to end, regardless of final direction
    (0..4).map(|dir| dist[&(end, dir)]).min().unwrap()
}

fn part1(input: String) {
    let (field, start, end) = parse(&input);
    let distance = shortest_path_length(&field, (start, 0), end);
    println!("{distance}");
}

fn shortest_path_tiles(field: &[Vec<bool>], start: Node, end: Point2D<i32>) -> u32 {
    let mut parents: FxHashMap<Node, Vec<Node>> = FxHashMap::default();
    let mut dist: FxHashMap<Node, u32> = FxHashMap::default();
    dist.insert(start, 0);
    let mut pq: PriorityQueue<Node, Reverse<u32>> = PriorityQueue::new();
    pq.push(start, Reverse(0));
    while let Some((v, _)) = pq.pop() {
        for (w, weight) in adj(field, v) {
            let dist_w = dist.get(&w).copied().unwrap_or(u32::MAX);
            let new_dist = dist[&v] + weight;
            match dist_w.cmp(&new_dist) {
                Ordering::Greater => {
                    parents.insert(w, vec![v]);
                    dist.insert(w, new_dist);
                    pq.push_increase(w, Reverse(new_dist));
                }
                // Remember both parents if distance is equal
                Ordering::Equal => parents.get_mut(&w).unwrap().push(v),
                Ordering::Less => {}
            }
        }
    }
    let mut path_tiles: FxHashSet<Point2D<i32>> = FxHashSet::default();
    path_tiles.insert(end);

    // Shortest distance to end, regardless of final direction
    let shortest_dist = (0..4).map(|dir| dist[&(end, dir)]).min().unwrap();
    for dir in 0..4 {
        if dist[&(end, dir)] == shortest_dist {
            collect_tiles(&parents, &mut path_tiles, (end, dir));
        }
    }
    path_tiles.len() as u32
}

fn collect_tiles(
    parents: &FxHashMap<Node, Vec<Node>>,
    tiles: &mut FxHashSet<Point2D<i32>>,
    cur: Node,
) {
    if let Some(pars) = parents.get(&cur) {
        for p in pars {
            tiles.insert(p.0);
            collect_tiles(parents, tiles, *p);
        }
    }
}

fn part2(input: String) {
    let (field, start, end) = parse(&input);
    let tiles = shortest_path_tiles(&field, (start, 0), end);
    println!("{tiles}");
}

util::aoc_main!();