module Main where

import Data.Functor
import Data.Char
import Data.Matrix hiding (trace)
import Data.List

parse :: String -> Matrix Int
parse = fromLists . ((digitToInt <$>) <$>) . lines

getAllDigit :: Int -> Matrix Int -> [(Int,Int)]
getAllDigit c s = zip [1..] (toLists s) >>= \(y,l) -> zip [1..] l >>= \(x,c') ->
  if c == c' then pure (y,x) else mempty

getAbove :: (Int,Int) -> Matrix Int -> [(Int,Int)]
getAbove (y,x) s = let target = getElem y x s + 1 in
  filter ((Just target ==) . flip (uncurry safeGet) s) [(y-1,x),(y+1,x),(y,x-1),(y,x+1)]

travel :: (Int,Int) -> Matrix Int -> [(Int,Int)]
travel p s = let curr = uncurry getElem p s in
  if curr == 9 then [p] else getAbove p s >>= (`travel` s)

travel' :: (Int,Int) -> Matrix Int -> Int
travel' p s = let curr = uncurry getElem p s in
  if curr == 9 then 1 else sum $ getAbove p s <&> (`travel'` s)

solve1 :: Matrix Int -> Int
solve1 s = length $ nub . (`travel` s) =<< getAllDigit 0 s

solve2 :: Matrix Int -> Int
solve2 s = sum $ getAllDigit 0 s <&> (`travel'` s)

main :: IO ()
main = readFile "inputs/10" <&> parse >>= \i ->
       print (solve1 i) >>
       print (solve2 i)