module Day09 where
import Data.Char ( digitToInt )
import Data.Maybe (catMaybes)
import Data.Containers (IsMap(filterMap))
import qualified Data.HashSet as HashSet
import Data.HashSet (HashSet)
import Data.List (sort)

infixr 5 !?
(!?) :: [a] -> Int -> Maybe a
(!?) [] _ = Nothing
(!?) (x:xs) i
    | i < 0 = Nothing
    | i == 0 = Just x
    | otherwise = xs !? i-1

getGrid :: [[Int]] -> Int -> Int -> Maybe Int
getGrid grid x y = do
    row <- grid !? y
    row !? x

getMinimumNeighbor :: [[Int]] -> Int -> Int -> Int
getMinimumNeighbor grid x y = minimum $ catMaybes [g (x-1) y, g (x+1) y, g x (y-1), g x (y+1)]
    where g x y = getGrid grid x y

part1:: [[Int]] -> Int
part1 input = sum $ map sum $ zipWith (\y row -> zipWith (\x digit -> apply x y digit) [0 .. ] row) [0..] input
    where apply x y digit = if digit < getMinimumNeighbor input x y then digit + 1 else 0

getLows :: [[Int]] -> [(Int,Int)]
getLows input = concatMap catMaybes (zipWith (\y row -> zipWith (\x digit -> apply x y digit) [0 .. ] row) [0..] input)
    where apply x y digit = if digit < getMinimumNeighbor input x y then Just (x,y) else Nothing

stepBasin :: [[Int]] -> HashSet (Int,Int) -> HashSet (Int,Int)
stepBasin input set = HashSet.foldr HashSet.union HashSet.empty $ HashSet.map (\(x,y) -> HashSet.fromList $ getElems x y) set
    where
        getElems x y = (x,y) : map (\(x,y,n) -> (x,y)) (getHigherNeighs (Just $ input !! y !! x) x y)
        getHigherNeighs digit x y = filter (\(x,y,n) -> n > digit && n < Just 9) $ getNeighs x y
        getNeighs x y = [g (x-1) y, g (x+1) y, g x (y-1), g x (y+1)]
        g x y = (x,y,getGrid input x y)

getBasin :: [[Int]] -> (Int,Int) -> Int
getBasin input (x,y) = length $ HashSet.toList $ calcBasin $ HashSet.fromList [(x,y)]
    where
        calcBasin set = let nextSet = stepBasin input set in if length nextSet == length set then set else calcBasin nextSet

part2 :: [[Int]] -> Int
part2 input = product $ take 3 $ reverse $ sort $ map (getBasin input) $ getLows input

main :: IO ()
main = do
    input <- map (map digitToInt) . lines <$> getContents
    putStrLn $ "part 1: " ++ show (part1 input)
    putStrLn $ "part 2: " ++ show (part2 input)