{-# Language ViewPatterns #-}
{-|
Module      : Main
Description : Day 15 solution
Copyright   : (c) Eric Mertens, 2021
License     : ISC
Maintainer  : emertens@gmail.com

<https://adventofcode.com/2021/day/15>

Finding the shortest route through a cave, and then
finding the shortest route through a slightly larger cave.

This solution uses Dijkstra's Algorithm to perform a
shortest path search through the cave. (A* with a zero
heuristic degenerates to this)

For part 2 this solution transforms the lookup coordinates
rather than to build a larger cave array. The reason for this
is to reduce memory pressure especially when running the
search on much larger maps.

-}
module Main (main) where

import Advent (arrIx, getInputArray)
import Advent.Coord (Coord(..), cardinal, origin)
import Advent.Search (AStep(..), astar)
import Data.Array.Unboxed ((!), amap, IArray(bounds), UArray)
import Data.Char (digitToInt)
import Data.Maybe (fromMaybe, maybeToList)
import Data.Word (Word8)

-- | >>> :main
-- 698
-- 3022
main :: IO ()
IO ()
main =
 do UArray Coord Word8
inp <- (Char -> Word8) -> UArray Coord Char -> UArray Coord Word8
forall (a :: * -> * -> *) e' e i.
(IArray a e', IArray a e, Ix i) =>
(e' -> e) -> a i e' -> a i e
amap (Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Word8) -> (Char -> Int) -> Char -> Word8
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Char -> Int
digitToInt) (UArray Coord Char -> UArray Coord Word8)
-> IO (UArray Coord Char) -> IO (UArray Coord Word8)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Int -> IO (UArray Coord Char)
getInputArray Int
2021 Int
15
    let end :: Coord
end = (Coord, Coord) -> Coord
forall a b. (a, b) -> b
snd (UArray Coord Word8 -> (Coord, Coord)
forall i. Ix i => UArray i Word8 -> (i, i)
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> (i, i)
bounds UArray Coord Word8
inp)
    Int -> IO ()
forall a. Show a => a -> IO ()
print (Coord -> (Coord -> Maybe Word8) -> Int
solve Coord
end (UArray Coord Word8 -> Coord -> Maybe Word8
forall (a :: * -> * -> *) e i (f :: * -> *).
(IArray a e, Ix i, Alternative f) =>
a i e -> i -> f e
arrIx UArray Coord Word8
inp))
    let (Coord
end2, Coord -> Maybe Word8
inp2) = UArray Coord Word8 -> (Coord, Coord -> Maybe Word8)
extendCave UArray Coord Word8
inp
    Int -> IO ()
forall a. Show a => a -> IO ()
print (Coord -> (Coord -> Maybe Word8) -> Int
solve Coord
end2 Coord -> Maybe Word8
inp2)

-- | Compute the risk score traveling through a cave.
solve :: Coord -> (Coord -> Maybe Word8) -> Int
solve :: Coord -> (Coord -> Maybe Word8) -> Int
solve Coord
end Coord -> Maybe Word8
costAt = Int -> Maybe Int -> Int
forall a. a -> Maybe a -> a
fromMaybe ([Char] -> Int
forall a. HasCallStack => [Char] -> a
error [Char]
"no path") (Coord -> [(Coord, Int)] -> Maybe Int
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup Coord
end [(Coord, Int)]
costs)
  where
    costs :: [(Coord, Int)]
costs = (Coord -> [AStep Coord]) -> Coord -> [(Coord, Int)]
forall a. Ord a => (a -> [AStep a]) -> a -> [(a, Int)]
astar Coord -> [AStep Coord]
step Coord
origin
    step :: Coord -> [AStep Coord]
step Coord
here = [ Coord -> Int -> Int -> AStep Coord
forall a. a -> Int -> Int -> AStep a
AStep Coord
next (Word8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
cost) Int
0
                | Coord
next <- Coord -> [Coord]
cardinal Coord
here
                , Word8
cost <- Maybe Word8 -> [Word8]
forall a. Maybe a -> [a]
maybeToList (Coord -> Maybe Word8
costAt Coord
next)]

-- | Build a larger cave by tiling the input cave in a 5x5
-- grid. Added caves have their risk values updated according
-- to their new locations.
--
-- Rather than building a huge array this uses the original array
-- and computes a function to extracts the /tile/ coordinates
-- of the given coordinate and uses that to compute the risk on the fly.
extendCave :: UArray Coord Word8 -> (Coord, Coord -> Maybe Word8)
extendCave :: UArray Coord Word8 -> (Coord, Coord -> Maybe Word8)
extendCave UArray Coord Word8
m = Coord
end Coord
-> (Coord, Coord -> Maybe Word8) -> (Coord, Coord -> Maybe Word8)
forall a b. a -> b -> b
`seq` (Coord
end, Coord -> Maybe Word8
costAt)
  where
    (C Int
0 Int
0, (Coord
1Coord -> Coord -> Coord
forall a. Num a => a -> a -> a
+) -> C Int
wy Int
wx) = UArray Coord Word8 -> (Coord, Coord)
forall i. Ix i => UArray i Word8 -> (i, i)
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> (i, i)
bounds UArray Coord Word8
m
    end :: Coord
end = Int -> Int -> Coord
C (Int
5Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
wy) (Int
5Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
wx) Coord -> Coord -> Coord
forall a. Num a => a -> a -> a
- Coord
1
    view :: Int -> Int -> (Int, Int)
view = (Int -> Int -> (Int, Int)) -> Int -> Int -> (Int, Int)
forall a b c. (a -> b -> c) -> b -> a -> c
flip Int -> Int -> (Int, Int)
forall a. Integral a => a -> a -> (a, a)
divMod

    costAt :: Coord -> Maybe Word8
costAt (C (Int -> Int -> (Int, Int)
view Int
wy -> (Int
ty, Int
y)) (Int -> Int -> (Int, Int)
view Int
wx -> (Int
tx, Int
x)))
      | Int
0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
ty, Int
ty Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
5, Int
0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
tx, Int
tx Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
5 =
        Word8 -> Maybe Word8
forall a. a -> Maybe a
Just (Word8 -> Maybe Word8) -> Word8 -> Maybe Word8
forall a b. (a -> b) -> a -> b
$! Word8 -> Word8
fixRisk (UArray Coord Word8
m UArray Coord Word8 -> Coord -> Word8
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
! Int -> Int -> Coord
C Int
y Int
x Word8 -> Word8 -> Word8
forall a. Num a => a -> a -> a
+ Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
tx Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
ty))
      | Bool
otherwise = Maybe Word8
forall a. Maybe a
Nothing

-- | Risks are defined to roll over from 9 back to 1
--
-- >>> fixRisk <$> [1,5,9,10,12]
-- [1,5,9,1,3]
fixRisk :: Word8 -> Word8
fixRisk :: Word8 -> Word8
fixRisk Word8
x = (Word8
x Word8 -> Word8 -> Word8
forall a. Num a => a -> a -> a
- Word8
1) Word8 -> Word8 -> Word8
forall a. Integral a => a -> a -> a
`mod` Word8
9 Word8 -> Word8 -> Word8
forall a. Num a => a -> a -> a
+ Word8
1