{-|
Module:         Day9
Description:    <https://adventofcode.com/2021/day/9 Day 9: Smoke Basin>
-}
{-# LANGUAGE TypeApplications #-}
module Day9 (day9a, day9b) where

import Data.Char (digitToInt, isDigit)
import Data.Graph.Inductive (Gr, buildGr, components)
import Data.List (mapAccumL, mapAccumR, sortOn, zipWith5)
import Data.Maybe (catMaybes)
import Data.Ord (Down(Down))
import Data.Text (Text)
import qualified Data.Text as T (lines, unpack)

day9a :: Text -> Int
day9a :: Text -> Int
day9a Text
input = [Int] -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Int]
risks where
    heights :: [String]
heights = Text -> String
T.unpack (Text -> String) -> [Text] -> [String]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Text -> [Text]
T.lines Text
input
    risks :: [Int]
risks = [[Int]] -> [Int]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[Int]] -> [Int]) -> [[Int]] -> [Int]
forall a b. (a -> b) -> a -> b
$ (String -> String -> String -> [Int])
-> [String] -> [String] -> [String] -> [[Int]]
forall a b c d. (a -> b -> c -> d) -> [a] -> [b] -> [c] -> [d]
zipWith3 String -> String -> String -> [Int]
basins [String]
heights
        (Char -> String
forall a. a -> [a]
repeat Char
forall a. Bounded a => a
maxBound String -> [String] -> [String]
forall a. a -> [a] -> [a]
: [String]
heights) (Int -> [String] -> [String]
forall a. Int -> [a] -> [a]
drop Int
1 [String]
heights [String] -> [String] -> [String]
forall a. [a] -> [a] -> [a]
++ [Char -> String
forall a. a -> [a]
repeat Char
forall a. Bounded a => a
maxBound])
    basins :: String -> String -> String -> [Int]
basins String
row String
above String
below = [Maybe Int] -> [Int]
forall a. [Maybe a] -> [a]
catMaybes ([Maybe Int] -> [Int]) -> [Maybe Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ (Char -> Char -> Char -> Char -> Char -> Maybe Int)
-> String -> String -> String -> String -> String -> [Maybe Int]
forall a b c d e f.
(a -> b -> c -> d -> e -> f)
-> [a] -> [b] -> [c] -> [d] -> [e] -> [f]
zipWith5 Char -> Char -> Char -> Char -> Char -> Maybe Int
basin String
row String
above String
below
        (Char
forall a. Bounded a => a
maxBound Char -> String -> String
forall a. a -> [a] -> [a]
: String
row) (Int -> String -> String
forall a. Int -> [a] -> [a]
drop Int
1 String
row String -> String -> String
forall a. [a] -> [a] -> [a]
++ Char -> String
forall a. a -> [a]
repeat Char
forall a. Bounded a => a
maxBound)
    basin :: Char -> Char -> Char -> Char -> Char -> Maybe Int
basin Char
x Char
above Char
below Char
left Char
right
      | Char -> Bool
isDigit Char
x, Char
x Char -> Char -> Bool
forall a. Ord a => a -> a -> Bool
< Char
above, Char
x Char -> Char -> Bool
forall a. Ord a => a -> a -> Bool
< Char
below, Char
x Char -> Char -> Bool
forall a. Ord a => a -> a -> Bool
< Char
left, Char
x Char -> Char -> Bool
forall a. Ord a => a -> a -> Bool
< Char
right = Int -> Maybe Int
forall a. a -> Maybe a
Just (Int -> Maybe Int) -> Int -> Maybe Int
forall a b. (a -> b) -> a -> b
$ Char -> Int
digitToInt Char
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
      | Bool
otherwise = Maybe Int
forall a. Maybe a
Nothing

day9b :: Text -> Int
day9b :: Text -> Int
day9b Text
input = [Int] -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product ([Int] -> Int) -> [Int] -> Int
forall a b. (a -> b) -> a -> b
$ Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
take Int
3 ([Int] -> [Int]) -> [Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ (Int -> Down Int) -> [Int] -> [Int]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn Int -> Down Int
forall a. a -> Down a
Down ([Int] -> [Int]) -> [Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ ([Int] -> Int) -> [[Int]] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map [Int] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([[Int]] -> [Int]) -> [[Int]] -> [Int]
forall a b. (a -> b) -> a -> b
$ Gr () () -> [[Int]]
forall (gr :: * -> * -> *) a b. Graph gr => gr a b -> [[Int]]
components Gr () ()
basins where
    basins :: Gr () ()
basins = forall a b. DynGraph Gr => [Context a b] -> Gr a b
forall (gr :: * -> * -> *) a b.
DynGraph gr =>
[Context a b] -> gr a b
buildGr @Gr ([Context () ()] -> Gr () ()) -> [Context () ()] -> Gr () ()
forall a b. (a -> b) -> a -> b
$ [[Context () ()]] -> [Context () ()]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[Context () ()]] -> [Context () ()])
-> [[Context () ()]] -> [Context () ()]
forall a b. (a -> b) -> a -> b
$ ((Int, [Maybe Int]), [[Context () ()]]) -> [[Context () ()]]
forall a b. (a, b) -> b
snd (((Int, [Maybe Int]), [[Context () ()]]) -> [[Context () ()]])
-> ((Int, [Maybe Int]), [[Context () ()]]) -> [[Context () ()]]
forall a b. (a -> b) -> a -> b
$ ((Int, [Maybe Int])
 -> String -> ((Int, [Maybe Int]), [Context () ()]))
-> (Int, [Maybe Int])
-> [String]
-> ((Int, [Maybe Int]), [[Context () ()]])
forall (t :: * -> *) a b c.
Traversable t =>
(a -> b -> (a, c)) -> a -> t b -> (a, t c)
mapAccumR (Int, [Maybe Int])
-> String -> ((Int, [Maybe Int]), [Context () ()])
forall a a.
Num a =>
(a, [Maybe a])
-> String -> ((a, [Maybe a]), [([((), a)], a, (), [a])])
mkRow (Int
0, []) ([String] -> ((Int, [Maybe Int]), [[Context () ()]]))
-> [String] -> ((Int, [Maybe Int]), [[Context () ()]])
forall a b. (a -> b) -> a -> b
$ Text -> String
T.unpack (Text -> String) -> [Text] -> [String]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Text -> [Text]
T.lines Text
input
    mkRow :: (a, [Maybe a])
-> String -> ((a, [Maybe a]), [([((), a)], a, (), [a])])
mkRow (a
n, [Maybe a]
prev) String
line = ((a, [Maybe a])
next, [([((), a)], a, (), [a])]
forall a. [([((), a)], a, (), [a])]
contexts) where
        next :: (a, [Maybe a])
next@(a
_, [Maybe a]
ns) = (a -> Char -> (a, Maybe a)) -> a -> String -> (a, [Maybe a])
forall (t :: * -> *) a b c.
Traversable t =>
(a -> b -> (a, c)) -> a -> t b -> (a, t c)
mapAccumL a -> Char -> (a, Maybe a)
forall a. Num a => a -> Char -> (a, Maybe a)
f a
n String
line
        f :: a -> Char -> (a, Maybe a)
f a
n' Char
c | Char -> Bool
isDigit Char
c, Char -> Int
digitToInt Char
c Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
9 = (a
n' a -> a -> a
forall a. Num a => a -> a -> a
+ a
1, a -> Maybe a
forall a. a -> Maybe a
Just a
n')
        f a
n' Char
_ = (a
n', Maybe a
forall a. Maybe a
Nothing)
        contexts :: [([((), a)], a, (), [a])]
contexts = [Maybe ([((), a)], a, (), [a])] -> [([((), a)], a, (), [a])]
forall a. [Maybe a] -> [a]
catMaybes ([Maybe ([((), a)], a, (), [a])] -> [([((), a)], a, (), [a])])
-> [Maybe ([((), a)], a, (), [a])] -> [([((), a)], a, (), [a])]
forall a b. (a -> b) -> a -> b
$ (Maybe a -> Maybe a -> Maybe a -> Maybe ([((), a)], a, (), [a]))
-> [Maybe a]
-> [Maybe a]
-> [Maybe a]
-> [Maybe ([((), a)], a, (), [a])]
forall a b c d. (a -> b -> c -> d) -> [a] -> [b] -> [c] -> [d]
zipWith3 Maybe a -> Maybe a -> Maybe a -> Maybe ([((), a)], a, (), [a])
forall b b a.
Maybe b -> Maybe b -> Maybe b -> Maybe ([((), b)], b, (), [a])
g [Maybe a]
ns ([Maybe a]
prev [Maybe a] -> [Maybe a] -> [Maybe a]
forall a. [a] -> [a] -> [a]
++ Maybe a -> [Maybe a]
forall a. a -> [a]
repeat Maybe a
forall a. Maybe a
Nothing) (Int -> [Maybe a] -> [Maybe a]
forall a. Int -> [a] -> [a]
drop Int
1 [Maybe a]
ns [Maybe a] -> [Maybe a] -> [Maybe a]
forall a. [a] -> [a] -> [a]
++ Maybe a -> [Maybe a]
forall a. a -> [a]
repeat Maybe a
forall a. Maybe a
Nothing)
        g :: Maybe b -> Maybe b -> Maybe b -> Maybe ([((), b)], b, (), [a])
g (Just b
n') Maybe b
above' Maybe b
right =
            ([((), b)], b, (), [a]) -> Maybe ([((), b)], b, (), [a])
forall a. a -> Maybe a
Just ([((), b
m) | b
m <- [Maybe b] -> [b]
forall a. [Maybe a] -> [a]
catMaybes [Maybe b
above', Maybe b
right]], b
n', (), [])
        g Maybe b
_ Maybe b
_ Maybe b
_ = Maybe ([((), b)], b, (), [a])
forall a. Maybe a
Nothing