{-# LANGUAGE TupleSections #-}
module Main where

import Control.Monad (forM_)
import Control.Monad.ST (ST)
import qualified Data.Array as A
import qualified Data.Array.ST as MA
import Data.Function (on)
import Data.List hiding (insert)
import Data.Ord

import Input


-- Priority queue that only performs okay if there are only a small number of
-- distinct keys in the queue at any one time.
newtype PQueue k v = PQueue [(k, [v])]
  deriving (Show)
qFromList :: Ord k => [(k, v)] -> PQueue k v
qFromList [] = PQueue []
qFromList l = PQueue (map ((,) <$> fst . head <*> map snd)
                      . groupBy ((==) `on` fst)
                      $ sortBy (comparing fst) l)
qAdds :: Ord k => [(k, v)] -> PQueue k v -> PQueue k v
qAdds = \pairs (PQueue l) -> PQueue (inserts (let PQueue pq = qFromList pairs in pq) l)
  where inserts :: Ord k => [(k, [v])] -> [(k, [v])] -> [(k, [v])]
        inserts [] l = l
        inserts news [] = news
        inserts news@((k, vs) : restnews) l@((k', vs') : rest) = case compare k k' of
          LT -> (k, vs) : inserts restnews l
          EQ -> (k, vs ++ vs') : inserts restnews rest
          GT -> (k', vs') : inserts news rest
qPop :: PQueue k v -> Maybe ((k, v), PQueue k v)
qPop (PQueue []) = Nothing
qPop (PQueue ((k, [v]) : l)) = Just ((k, v), PQueue l)
qPop (PQueue ((k, v:vs) : l)) = Just ((k, v), PQueue ((k, vs) : l))
qPop (PQueue ((_, []) : l)) = qPop (PQueue l)

dijkstra :: A.Array (Int, Int) Int -> A.Array (Int, Int) Int
dijkstra board = MA.runSTArray $ do
    arr <- MA.newArray (A.bounds board) maxBound
    MA.writeArray arr (0, 0) 0
    loop arr (qFromList [(0, (0, 0))])
    return arr
  where
    loop :: MA.STArray s (Int, Int) Int -> PQueue Int (Int, Int) -> ST s ()
    loop arr qu
      | Just ((curdist, (x, y)), qu') <- qPop qu = do
          neighbours <- mapM (\pos -> (pos,) <$> MA.readArray arr pos)
                        . filter (A.inRange (A.bounds board))
                        $ [(x+1,y), (x,y+1), (x-1,y), (x,y-1)]
          let updates = [(curdist + board A.! pos', pos')
                        | (pos', dist) <- neighbours
                        , dist > curdist + board A.! pos']
          forM_ updates $ \(newdist, pos') ->
              MA.writeArray arr pos' newdist
          loop arr (qAdds updates qu')
      | otherwise = return ()

expand :: Int -> (Int -> a -> a) -> A.Array (Int, Int) a -> A.Array (Int, Int) a
expand times f arr =
    let ((0, 0), (ymax, xmax)) = A.bounds arr
        (h, w) = (ymax + 1, xmax + 1)
    in A.listArray ((0, 0), (times * h - 1, times * w - 1))
         [f (x `div` w + y `div` w) (arr A.! (y `mod` h, x `mod` w))
         | y <- [0 .. times * h - 1]
         , x <- [0 .. times * w - 1]]

main :: IO ()
main = do
    board <- map (map (read . pure)) <$> getInput 15
    let (h, w) = (length board, length (head board))
    let boardarr = A.listArray ((0, 0), (h-1, w-1)) (concat board)
    print (dijkstra boardarr A.! (h-1, w-1))
    let expanded = expand 5 (\n x -> (x + n - 1) `mod` 9 + 1) boardarr
    print (dijkstra expanded A.! (5*h-1, 5*w-1))