{-# Language QuasiQuotes #-}
{-# OPTIONS_GHC -funbox-strict-fields #-}
{-|
Module      : Main
Description : Day 17 solution
Copyright   : (c) Eric Mertens, 2023
License     : ISC
Maintainer  : emertens@gmail.com

<https://adventofcode.com/2023/day/17>

Shortest-path graph search where the graph states are the triple of
a location, direction.

Distance traveled doesn't need to be stored because all of the distances
that can be traveled from a starting location are added to the work
queue at the same time for each starting point.

>>> :{
:main +
"2413432311323
3215453535623
3255245654254
3446585845452
4546657867536
1438598798454
4457876987766
3637877979653
4654967986887
4564679986453
1224686865563
2546548887735
4322674655533
"
:}
102
94

>>> :{
:main +
"111111111111
999999999991
999999999991
999999999991
999999999991
"
:}
59
71

-}
module Main (main) where

import Advent (getInputArray, arrIx)
import Advent.Coord (east, south, turnLeft, turnRight, Coord)
import Advent.Search (astarN, AStep(..))
import Data.Array.Unboxed (amap, bounds, UArray)
import Data.Char (digitToInt)

-- | Parse input grid and print both answer parts.
--
-- >>> :main
-- 866
-- 1010
main :: IO ()
IO ()
main =
 do input <- (Char -> Int) -> UArray Coord Char -> UArray Coord Int
forall (a :: * -> * -> *) e' e i.
(IArray a e', IArray a e, Ix i) =>
(e' -> e) -> a i e' -> a i e
amap Char -> Int
digitToInt (UArray Coord Char -> UArray Coord Int)
-> IO (UArray Coord Char) -> IO (UArray Coord Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Int -> IO (UArray Coord Char)
getInputArray Int
2023 Int
17
    print (solve 1  3 input)
    print (solve 4 10 input)

solve :: Int -> Int -> UArray Coord Int -> Int
solve :: Int -> Int -> UArray Coord Int -> Int
solve Int
lo Int
hi UArray Coord Int
input =
  [Int] -> Int
forall a. HasCallStack => [a] -> a
head [Int
cost | (S Coord
loc Coord
_, Int
cost) <- (S -> [AStep S]) -> [S] -> [(S, Int)]
forall a. Ord a => (a -> [AStep a]) -> [a] -> [(a, Int)]
astarN (Int -> Int -> UArray Coord Int -> S -> [AStep S]
step Int
lo Int
hi UArray Coord Int
input) [S]
begin
             , Coord
loc Coord -> Coord -> Bool
forall a. Eq a => a -> a -> Bool
== Coord
end -- at the end
             ]
  where
    (Coord
start, Coord
end) = UArray Coord Int -> (Coord, Coord)
forall i. Ix i => UArray i Int -> (i, i)
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> (i, i)
bounds UArray Coord Int
input
    begin :: [S]
begin = [Coord -> Coord -> S
S Coord
start Coord
east, Coord -> Coord -> S
S Coord
start Coord
south]

data S = S !Coord !Coord -- ^ location, direction
  deriving (S -> S -> Bool
(S -> S -> Bool) -> (S -> S -> Bool) -> Eq S
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: S -> S -> Bool
== :: S -> S -> Bool
$c/= :: S -> S -> Bool
/= :: S -> S -> Bool
Eq, Eq S
Eq S =>
(S -> S -> Ordering)
-> (S -> S -> Bool)
-> (S -> S -> Bool)
-> (S -> S -> Bool)
-> (S -> S -> Bool)
-> (S -> S -> S)
-> (S -> S -> S)
-> Ord S
S -> S -> Bool
S -> S -> Ordering
S -> S -> S
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 :: S -> S -> Ordering
compare :: S -> S -> Ordering
$c< :: S -> S -> Bool
< :: S -> S -> Bool
$c<= :: S -> S -> Bool
<= :: S -> S -> Bool
$c> :: S -> S -> Bool
> :: S -> S -> Bool
$c>= :: S -> S -> Bool
>= :: S -> S -> Bool
$cmax :: S -> S -> S
max :: S -> S -> S
$cmin :: S -> S -> S
min :: S -> S -> S
Ord, Int -> S -> ShowS
[S] -> ShowS
S -> String
(Int -> S -> ShowS) -> (S -> String) -> ([S] -> ShowS) -> Show S
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> S -> ShowS
showsPrec :: Int -> S -> ShowS
$cshow :: S -> String
show :: S -> String
$cshowList :: [S] -> ShowS
showList :: [S] -> ShowS
Show)

step :: Int -> Int -> UArray Coord Int -> S -> [AStep S]
step :: Int -> Int -> UArray Coord Int -> S -> [AStep S]
step Int
lo Int
hi UArray Coord Int
input (S Coord
here Coord
dir) =
  [ AStep {
      astepNext :: S
astepNext      = Coord -> Coord -> S
S Coord
here' Coord
dir',
      astepCost :: Int
astepCost      = Int
cost,
      astepHeuristic :: Int
astepHeuristic = Int
0}
  | Coord
dir' <- [Coord -> Coord
turnLeft Coord
dir, Coord -> Coord
turnRight Coord
dir]
  , (Coord
here', Int
cost) <- Int -> [(Coord, Int)] -> [(Coord, Int)]
forall a. Int -> [a] -> [a]
take (Int
hi Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
lo Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int -> [(Coord, Int)] -> [(Coord, Int)]
forall a. Int -> [a] -> [a]
drop (Int
lo Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (UArray Coord Int -> Coord -> Coord -> [(Coord, Int)]
costs UArray Coord Int
input Coord
dir' Coord
here))
  ]

-- lazy lists of the locations and costs to get there in order of distance
-- starting at 1 away
costs :: UArray Coord Int -> Coord -> Coord -> [(Coord, Int)]
costs :: UArray Coord Int -> Coord -> Coord -> [(Coord, Int)]
costs UArray Coord Int
input Coord
v = Int -> Coord -> [(Coord, Int)]
go Int
0
  where
    go :: Int -> Coord -> [(Coord, Int)]
go Int
acc Coord
l =
     do let l' :: Coord
l' = Coord
l Coord -> Coord -> Coord
forall a. Num a => a -> a -> a
+ Coord
v
        c <- UArray Coord Int -> Coord -> [Int]
forall (a :: * -> * -> *) e i (f :: * -> *).
(IArray a e, Ix i, Alternative f) =>
a i e -> i -> f e
arrIx UArray Coord Int
input Coord
l'
        let acc' = Int
acc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
c
        (l', acc') : go acc' l'