{-# Language QuasiQuotes, ImportQualifiedPost #-}
{-|
Module      : Main
Description : Day 22 solution
Copyright   : (c) Eric Mertens, 2018
License     : ISC
Maintainer  : emertens@gmail.com

<https://adventofcode.com/2018/day/22>

This solution uses an A* graph search to find the shortest path through
the cave to the target.

I picked an arbitrary cave size to memoize that was big enough to avoid
indexing errors. I use a boxed array so that I can lazily compute the
geologic indexes of the various locations in the cave. This also allows
me to recursively define 'geologic' and 'erosion'.

My A* heuristic is manhattan distance to the target plus a penalty for
not holding the torch. (Accounting for the torch saves a small, but
positive amount of time.)

Instead of modelling the tool being held directly I simply keep track of
the risk number of the area I'm not allowed to enter.

-}
module Main (main) where

import Advent (format)
import Advent.Coord (Coord(C), above, left, cardinal, manhattan, origin)
import Advent.Memo (memo)
import Advent.Search (astar, AStep(AStep))
import Data.Array qualified as A
import Data.List (delete)

-- | Print the answers to day 22
main :: IO ()
IO ()
main =
  do (Int
depth, Int
tx, Int
ty) <- [format|2018 22 depth: %u%ntarget: %u,%u%n|]
     let target :: Coord
target = Int -> Int -> Coord
C Int
ty Int
tx
     let risk :: Coord -> Tool
risk = Int -> Coord -> Coord -> Tool
mkRisk Int
depth Coord
target
     Int -> IO ()
forall a. Show a => a -> IO ()
print ((Coord -> Tool) -> Coord -> Int
part1 Coord -> Tool
risk Coord
target)
     Int -> IO ()
forall a. Show a => a -> IO ()
print ((Coord -> Tool) -> Coord -> Int
part2 Coord -> Tool
risk Coord
target)

-- | Sum of risk values in rectangle defined by origin and target
part1 :: (Coord -> Tool) -> Coord -> Int
part1 :: (Coord -> Tool) -> Coord -> Int
part1 Coord -> Tool
risk Coord
target = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Tool -> Int
toolId (Coord -> Tool
risk Coord
c) | Coord
c <- (Coord, Coord) -> [Coord]
forall a. Ix a => (a, a) -> [a]
A.range (Coord
origin, Coord
target)]

-- | Minimum cost of traveling to the target from the origin
part2 :: (Coord -> Tool) -> Coord -> Int
part2 :: (Coord -> Tool) -> Coord -> Int
part2 Coord -> Tool
risk Coord
target = Int
n
  where
    Just Int
n = Node -> [(Node, Int)] -> Maybe Int
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup Node
goal ((Node -> [AStep Node]) -> Node -> [(Node, Int)]
forall a. Ord a => (a -> [AStep a]) -> a -> [(a, Int)]
astar ((Coord -> Tool) -> Coord -> Node -> [AStep Node]
steps Coord -> Tool
risk Coord
target) Node
start)
    start :: Node
start  = Coord -> Tool -> Node
Node Coord
origin Tool
torch
    goal :: Node
goal   = Coord -> Tool -> Node
Node Coord
target Tool
torch

-- tool representation -------------------------------------------------

-- | Tools track the risk index that they are incompatible with.
newtype Tool = Tool { Tool -> Int
toolId :: Int } deriving (Int -> Tool -> ShowS
[Tool] -> ShowS
Tool -> String
(Int -> Tool -> ShowS)
-> (Tool -> String) -> ([Tool] -> ShowS) -> Show Tool
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Tool -> ShowS
showsPrec :: Int -> Tool -> ShowS
$cshow :: Tool -> String
show :: Tool -> String
$cshowList :: [Tool] -> ShowS
showList :: [Tool] -> ShowS
Show, Tool -> Tool -> Bool
(Tool -> Tool -> Bool) -> (Tool -> Tool -> Bool) -> Eq Tool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Tool -> Tool -> Bool
== :: Tool -> Tool -> Bool
$c/= :: Tool -> Tool -> Bool
/= :: Tool -> Tool -> Bool
Eq, Eq Tool
Eq Tool =>
(Tool -> Tool -> Ordering)
-> (Tool -> Tool -> Bool)
-> (Tool -> Tool -> Bool)
-> (Tool -> Tool -> Bool)
-> (Tool -> Tool -> Bool)
-> (Tool -> Tool -> Tool)
-> (Tool -> Tool -> Tool)
-> Ord Tool
Tool -> Tool -> Bool
Tool -> Tool -> Ordering
Tool -> Tool -> Tool
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: Tool -> Tool -> Ordering
compare :: Tool -> Tool -> Ordering
$c< :: Tool -> Tool -> Bool
< :: Tool -> Tool -> Bool
$c<= :: Tool -> Tool -> Bool
<= :: Tool -> Tool -> Bool
$c> :: Tool -> Tool -> Bool
> :: Tool -> Tool -> Bool
$c>= :: Tool -> Tool -> Bool
>= :: Tool -> Tool -> Bool
$cmax :: Tool -> Tool -> Tool
max :: Tool -> Tool -> Tool
$cmin :: Tool -> Tool -> Tool
min :: Tool -> Tool -> Tool
Ord)

-- | The torch tool is used at the beginning and end of the trip.
torch :: Tool
torch :: Tool
torch = Int -> Tool
Tool Int
1 -- torch is excluded from wet (1) squares

-- | List of all three tools.
tools :: [Tool]
tools :: [Tool]
tools = [Int -> Tool
Tool Int
0, Int -> Tool
Tool Int
1, Int -> Tool
Tool Int
2]

-- movement rules ------------------------------------------------------

-- | Graph search node. There will be a lot of these and this
-- representation is much more compact than a tuple. This will represent
-- a 3D location in the graph search: current position and current tool.
data Node = Node {-# Unpack #-}!Coord {-# Unpack #-}!Tool deriving (Node -> Node -> Bool
(Node -> Node -> Bool) -> (Node -> Node -> Bool) -> Eq Node
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Node -> Node -> Bool
== :: Node -> Node -> Bool
$c/= :: Node -> Node -> Bool
/= :: Node -> Node -> Bool
Eq, Eq Node
Eq Node =>
(Node -> Node -> Ordering)
-> (Node -> Node -> Bool)
-> (Node -> Node -> Bool)
-> (Node -> Node -> Bool)
-> (Node -> Node -> Bool)
-> (Node -> Node -> Node)
-> (Node -> Node -> Node)
-> Ord Node
Node -> Node -> Bool
Node -> Node -> Ordering
Node -> Node -> Node
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: Node -> Node -> Ordering
compare :: Node -> Node -> Ordering
$c< :: Node -> Node -> Bool
< :: Node -> Node -> Bool
$c<= :: Node -> Node -> Bool
<= :: Node -> Node -> Bool
$c> :: Node -> Node -> Bool
> :: Node -> Node -> Bool
$c>= :: Node -> Node -> Bool
>= :: Node -> Node -> Bool
$cmax :: Node -> Node -> Node
max :: Node -> Node -> Node
$cmin :: Node -> Node -> Node
min :: Node -> Node -> Node
Ord)

-- | Compute the states reachable from the given state. Cost is the
-- incremental cost of choosing that state. Heuristic is lower-bound on
-- the distance remaining until the target. This lower-bound is an
-- admissible heuristic that enables A* to find the optimal path.
steps ::
  (Coord -> Tool) {- ^ location to banned tool         -} ->
  Coord           {- ^ target location                 -} ->
  Node            {- ^ location, tool                  -} ->
  [AStep Node]    {- ^ location, tool, cost, heuristic -}
steps :: (Coord -> Tool) -> Coord -> Node -> [AStep Node]
steps Coord -> Tool
risk Coord
target (Node Coord
here Tool
tool) =
  [ Node -> Int -> Int -> AStep Node
forall a. a -> Int -> Int -> AStep a
AStep (Coord -> Tool -> Node
Node Coord
dest Tool
tool') Int
cost Int
heuristic
     | (Node Coord
dest Tool
tool', Int
cost) <- [(Node, Int)]
changeTool [(Node, Int)] -> [(Node, Int)] -> [(Node, Int)]
forall a. [a] -> [a] -> [a]
++ [(Node, Int)]
move
     , Coord -> Tool
risk Coord
dest Tool -> Tool -> Bool
forall a. Eq a => a -> a -> Bool
/= Tool
tool'
     , let heuristic :: Int
heuristic = Coord -> Coord -> Int
manhattan Coord
dest Coord
target
                     Int -> Int -> Int
forall a. Num a => a -> a -> a
+ if Tool
tool' Tool -> Tool -> Bool
forall a. Eq a => a -> a -> Bool
== Tool
torch then Int
0 else Int
7
     ]
  where
    changeTool :: [(Node, Int)]
changeTool = [(Coord -> Tool -> Node
Node Coord
here Tool
tool', Int
7) | Tool
tool' <- Tool -> [Tool] -> [Tool]
forall a. Eq a => a -> [a] -> [a]
delete Tool
tool [Tool]
tools ]
    move :: [(Node, Int)]
move = [(Coord -> Tool -> Node
Node Coord
dst Tool
tool, Int
1) | dst :: Coord
dst@(C Int
y Int
x) <- Coord -> [Coord]
cardinal Coord
here, Int
y Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0, Int
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0]

-- cave characterization -----------------------------------------------

-- | Computes a function that can query the risk index at a given query
-- coordinate. The query function is backed by an array to efficiently
-- compute risks for a given depth and target value.
mkRisk ::
  Int   {- ^ layer depth                    -} ->
  Coord {- ^ target coordinate              -} ->
  Coord {- ^ query coordinate               -} ->
  Tool  {- ^ risk index at query coordinate -}
mkRisk :: Int -> Coord -> Coord -> Tool
mkRisk Int
depth Coord
target = \Coord
i -> Int -> Tool
Tool (Coord -> Int
erosion Coord
i Int -> Int -> Int
forall a. Integral a => a -> a -> a
`rem` Int
3)
  where
    geologic :: Coord -> Int
geologic c :: Coord
c@(C Int
y Int
x)
      | Int
y Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0      = Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
16807
      | Int
x Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0      = Int
y Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
48271
      | Coord
c Coord -> Coord -> Bool
forall a. Eq a => a -> a -> Bool
== Coord
target = Int
0
      | Bool
otherwise   = Coord -> Int
erosion (Coord -> Coord
above Coord
c) Int -> Int -> Int
forall a. Num a => a -> a -> a
* Coord -> Int
erosion (Coord -> Coord
left Coord
c)

    erosion :: Coord -> Int
erosion = (Coord -> Int) -> Coord -> Int
forall t a. HasTrie t => (t -> a) -> t -> a
memo (\Coord
i -> (Coord -> Int
geologic Coord
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
depth) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`rem` Int
20183)