{-|
Module:         Day15
Description:    <https://adventofcode.com/2021/day/15 Day 15: Chiton>
-}
{-# LANGUAGE ViewPatterns #-}
module Day15 (day15a, day15b) where

import Control.Monad (foldM, guard)
import Control.Monad.ST (ST, runST)
import Data.Char (digitToInt, isDigit)
import Data.Heap (MinPrioHeap)
import qualified Data.Heap as Heap (insert, singleton, view)
import Data.List.NonEmpty (NonEmpty((:|)), nonEmpty)
import Data.Text (Text)
import qualified Data.Text as T (all, length, lines, unpack)
import Data.Vector.Unboxed (Vector)
import qualified Data.Vector.Unboxed as V ((!), fromList, length)
import Data.Vector.Unboxed.Mutable (MVector)
import qualified Data.Vector.Unboxed.Mutable as MV (read, replicate, write)

day15 :: Int -> Vector Int -> Maybe Int
day15 :: Int -> Vector Int -> Maybe Int
day15 Int
width Vector Int
risks = (forall s. ST s (Maybe Int)) -> Maybe Int
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Maybe Int)) -> Maybe Int)
-> (forall s. ST s (Maybe Int)) -> Maybe Int
forall a b. (a -> b) -> a -> b
$ do
    MVector s Int
bests <- Int -> Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
MV.replicate (Vector Int -> Int
forall a. Unbox a => Vector a -> Int
V.length Vector Int
risks) Int
forall a. Bounded a => a
maxBound
    MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
MV.write MVector s Int
MVector (PrimState (ST s)) Int
bests Int
0 Int
0
    MVector s Int -> MinPrioHeap Int Int -> ST s (Maybe Int)
forall s. MVector s Int -> MinPrioHeap Int Int -> ST s (Maybe Int)
go MVector s Int
bests (MinPrioHeap Int Int -> ST s (Maybe Int))
-> MinPrioHeap Int Int -> ST s (Maybe Int)
forall a b. (a -> b) -> a -> b
$ (Int, Int) -> MinPrioHeap Int Int
forall pol item. HeapItem pol item => item -> Heap pol item
Heap.singleton (Int
0, Int
0)
  where
    f :: MVector s Int -> Int -> MinPrioHeap Int Int -> (Int, Int) -> ST s (MinPrioHeap Int Int)
    f :: MVector s Int
-> Int
-> MinPrioHeap Int Int
-> (Int, Int)
-> ST s (MinPrioHeap Int Int)
f MVector s Int
bests Int
c MinPrioHeap Int Int
heap (Int
y, Int
x)
      | Int
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
width Bool -> Bool -> Bool
|| Int
y Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
y Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
width Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Vector Int -> Int
forall a. Unbox a => Vector a -> Int
V.length Vector Int
risks = HeapT (Prio FstMinPolicy (Int, Int)) Int
-> ST s (HeapT (Prio FstMinPolicy (Int, Int)) Int)
forall (f :: * -> *) a. Applicative f => a -> f a
pure HeapT (Prio FstMinPolicy (Int, Int)) Int
MinPrioHeap Int Int
heap
      | Bool
otherwise = do
        Int
best <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
MV.read MVector s Int
MVector (PrimState (ST s)) Int
bests Int
i
        if Int
risk Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
best then (Int, Int) -> MinPrioHeap Int Int -> MinPrioHeap Int Int
forall pol item.
HeapItem pol item =>
item -> Heap pol item -> Heap pol item
Heap.insert (Int
risk, Int
i) MinPrioHeap Int Int
heap HeapT (Prio FstMinPolicy (Int, Int)) Int
-> ST s () -> ST s (HeapT (Prio FstMinPolicy (Int, Int)) Int)
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
MV.write MVector s Int
MVector (PrimState (ST s)) Int
bests Int
i Int
risk else HeapT (Prio FstMinPolicy (Int, Int)) Int
-> ST s (HeapT (Prio FstMinPolicy (Int, Int)) Int)
forall (f :: * -> *) a. Applicative f => a -> f a
pure HeapT (Prio FstMinPolicy (Int, Int)) Int
MinPrioHeap Int Int
heap
      where
        i :: Int
i = Int
y Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
width Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
x
        risk :: Int
risk = Int
c Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Vector Int
risks Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
V.! Int
i
    go :: MVector s Int -> MinPrioHeap Int Int -> ST s (Maybe Int)
    go :: MVector s Int -> MinPrioHeap Int Int -> ST s (Maybe Int)
go MVector s Int
bests (MinPrioHeap Int Int -> Maybe ((Int, Int), MinPrioHeap Int Int)
forall pol item.
HeapItem pol item =>
Heap pol item -> Maybe (item, Heap pol item)
Heap.view -> Just ((Int
_, Int
i), MinPrioHeap Int Int
heap))
      | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Vector Int -> Int
forall a. Unbox a => Vector a -> Int
V.length Vector Int
risks Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 = Int -> Maybe Int
forall a. a -> Maybe a
Just (Int -> Maybe Int) -> ST s Int -> ST s (Maybe Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
MV.read MVector s Int
MVector (PrimState (ST s)) Int
bests Int
i
      | Bool
otherwise = do
        Int
c <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
MV.read MVector s Int
MVector (PrimState (ST s)) Int
bests Int
i
        HeapT (Prio FstMinPolicy (Int, Int)) Int
heap' <- (HeapT (Prio FstMinPolicy (Int, Int)) Int
 -> (Int, Int) -> ST s (HeapT (Prio FstMinPolicy (Int, Int)) Int))
-> HeapT (Prio FstMinPolicy (Int, Int)) Int
-> [(Int, Int)]
-> ST s (HeapT (Prio FstMinPolicy (Int, Int)) Int)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (MVector s Int
-> Int
-> MinPrioHeap Int Int
-> (Int, Int)
-> ST s (MinPrioHeap Int Int)
forall s.
MVector s Int
-> Int
-> MinPrioHeap Int Int
-> (Int, Int)
-> ST s (MinPrioHeap Int Int)
f MVector s Int
bests Int
c) HeapT (Prio FstMinPolicy (Int, Int)) Int
MinPrioHeap Int Int
heap ([(Int, Int)] -> ST s (HeapT (Prio FstMinPolicy (Int, Int)) Int))
-> [(Int, Int)] -> ST s (HeapT (Prio FstMinPolicy (Int, Int)) Int)
forall a b. (a -> b) -> a -> b
$
            let (Int
y, Int
x) = Int
i Int -> Int -> (Int, Int)
forall a. Integral a => a -> a -> (a, a)
`divMod` Int
width in [(Int
y, Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1), (Int
y Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1, Int
x), (Int
y Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1, Int
x), (Int
y, Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)]
        MVector s Int -> MinPrioHeap Int Int -> ST s (Maybe Int)
forall s. MVector s Int -> MinPrioHeap Int Int -> ST s (Maybe Int)
go MVector s Int
bests HeapT (Prio FstMinPolicy (Int, Int)) Int
MinPrioHeap Int Int
heap'
    go MVector s Int
_ MinPrioHeap Int Int
_ = Maybe Int -> ST s (Maybe Int)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Int
forall a. Maybe a
Nothing

day15a :: Text -> Maybe Int
day15a :: Text -> Maybe Int
day15a Text
input = do
    input' :: NonEmpty Text
input'@(Text
line :| [Text]
lines') <- [Text] -> Maybe (NonEmpty Text)
forall a. [a] -> Maybe (NonEmpty a)
nonEmpty ([Text] -> Maybe (NonEmpty Text))
-> [Text] -> Maybe (NonEmpty Text)
forall a b. (a -> b) -> a -> b
$ Text -> [Text]
T.lines Text
input
    let width :: Int
width = Text -> Int
T.length Text
line
    Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ Int
width Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 Bool -> Bool -> Bool
&& (Text -> Bool) -> [Text] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all ((Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
width) (Int -> Bool) -> (Text -> Int) -> Text -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Int
T.length) [Text]
lines' Bool -> Bool -> Bool
&& (Text -> Bool) -> NonEmpty Text -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all ((Char -> Bool) -> Text -> Bool
T.all Char -> Bool
isDigit) NonEmpty Text
input'
    Int -> Vector Int -> Maybe Int
day15 Int
width (Vector Int -> Maybe Int) -> Vector Int -> Maybe Int
forall a b. (a -> b) -> a -> b
$ [Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
V.fromList ([Int] -> Vector Int) -> [Int] -> Vector Int
forall a b. (a -> b) -> a -> b
$ (Text -> [Int]) -> NonEmpty Text -> [Int]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ((Char -> Int) -> [Char] -> [Int]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Char -> Int
digitToInt ([Char] -> [Int]) -> (Text -> [Char]) -> Text -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> [Char]
T.unpack) NonEmpty Text
input'

day15b :: Text -> Maybe Int
day15b :: Text -> Maybe Int
day15b Text
input = do
    input' :: NonEmpty Text
input'@(Text
line :| [Text]
lines') <- [Text] -> Maybe (NonEmpty Text)
forall a. [a] -> Maybe (NonEmpty a)
nonEmpty ([Text] -> Maybe (NonEmpty Text))
-> [Text] -> Maybe (NonEmpty Text)
forall a b. (a -> b) -> a -> b
$ Text -> [Text]
T.lines Text
input
    let width :: Int
width = Text -> Int
T.length Text
line
    Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ Int
width Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 Bool -> Bool -> Bool
&& (Text -> Bool) -> [Text] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all ((Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
width) (Int -> Bool) -> (Text -> Int) -> Text -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Int
T.length) [Text]
lines' Bool -> Bool -> Bool
&& (Text -> Bool) -> NonEmpty Text -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all ((Char -> Bool) -> Text -> Bool
T.all Char -> Bool
isDigit) NonEmpty Text
input'
    Int -> Vector Int -> Maybe Int
day15 (Int
5 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
width) (Vector Int -> Maybe Int) -> Vector Int -> Maybe Int
forall a b. (a -> b) -> a -> b
$ [Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
V.fromList ([Int] -> Vector Int) -> [Int] -> Vector Int
forall a b. (a -> b) -> a -> b
$ do
        Int
i <- [Int
0..Int
4]
        Text
line' <- Text
lineText -> [Text] -> [Text]
forall a. a -> [a] -> [a]
:[Text]
lines'
        Int
j <- [Int
0..Int
4]
        [(Char -> Int
digitToInt Char
c Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
j) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
9 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 | Char
c <- Text -> [Char]
T.unpack Text
line']