import Data.List
import Data.List.Split
import Data.Maybe

main :: IO ()
main = do
    input <- readFile "9/input.txt"
    let x = parse $ lines input
    -- Problem 1
    print $ f x
    -- Problem 2
    print $ g x
    where
        f = risk
        g = solve

parse :: [String] -> [[Int]]
parse = map (map read . chunksOf 1)

risk :: [[Int]] -> Int
risk a = sum $ map succ $ mapMaybe index (areas a)
    where
        index = ($ a) . uncurry idx

solve :: [[Int]] -> Int
solve x = product $ take 3 $ reverse $ sort $ map (basin x) $ areas x

-- | Find risk areas (low points).
areas :: [[Int]] -> [(Int,Int)]
areas a = [ (x,y) | y <- [ 0 .. (pred $ length a) ],
                    x <- [ 0 .. (length (a !! y)) ],
                    Just n <- [idx x y a],
                    all ((> n) . snd) (adj x y a) ]

-- | Get the size of a basin centered on the given point.
basin :: [[Int]] -> (Int, Int) -> Int
basin m (x,y) = length a
    where
        a = case idx x y m of
            Just v  -> let t = ((x,y),v) in t : search t m
            Nothing -> []

-- | Search for neighbors of the given node where the value of the neighbor
-- | is not 9 and is greater than the value of the given node. This search
-- | is then repeated for each result it finds until no nodes satisfying the
-- | constraint can be found in the immediate neighborhood of the cluster.
search :: ((Int, Int), Int) -> [[Int]] -> [((Int, Int), Int)]
search t @ ((x,y),v) b = nub (next ++ (next >>= (`search` b)))
    where
        neighbors = adj x y b
        next = nub $ filter (\(_,n) -> n /= 9 && n > v) neighbors

-- | Find elements that are directly adjacent to a point in a matrix.
adj :: Int -> Int -> [[a]] -> [((Int, Int), a)]
adj x y m = [ ((a, b), v) | a <- [pred x .. succ x],
                            b <- [pred y .. succ y],
                            x /= a || y /= b,
                            y == b || x == a,
                            Just v <- [idx a b m] ]

-- | Index into a matrix safely. If out of bounds on either axis, 
-- | returns `Nothing`.
idx :: Int -> Int -> [[a]] -> Maybe a
idx x y m = get y m >>= get x

-- | Index into a list safely. Returns `Nothing` if out of bounds.
get :: Int -> [a] -> Maybe a
get i l | i < 0         = Nothing
        | i >= length l = Nothing
        | otherwise     = Just (l !! i)