{-|
Module      : Main
Description : Day 9 solution
Copyright   : (c) Eric Mertens, 2021
License     : ISC
Maintainer  : emertens@gmail.com

<https://adventofcode.com/2021/day/9>

Find the basins on the height map.

-}
module Main (main) where

import Advent (counts, getInputArray, arrIx)
import Advent.Coord (Coord(..), cardinal)
import Data.Array.Unboxed (Array, UArray, (!), array, assocs, bounds, elems)
import Data.Char (digitToInt)
import Data.Foldable (toList)
import Data.List (sortBy)
import Data.Maybe (catMaybes)

type Heights  = Array Coord (Maybe Int)
type BasinIds = Array Coord (Maybe Coord)

-- >>> :main
-- 588
-- 964712
main :: IO ()
main :: IO ()
main =
 do Heights
heights <- UArray Coord Char -> Heights
heightArray (UArray Coord Char -> Heights)
-> IO (UArray Coord Char) -> IO Heights
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> IO (UArray Coord Char)
getInputArray Int
9
    let basinIds :: BasinIds
basinIds = Heights -> BasinIds
heightsToBasinIds Heights
heights

    Int -> IO ()
forall a. Show a => a -> IO ()
print ([Int] -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Int
1Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
h | (Coord
c, Just Int
h) <- Heights -> [(Coord, Maybe Int)]
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> [(i, e)]
assocs Heights
heights, Coord -> Maybe Coord
forall a. a -> Maybe a
Just Coord
c Maybe Coord -> Maybe Coord -> Bool
forall a. Eq a => a -> a -> Bool
== BasinIds
basinIdsBasinIds -> Coord -> Maybe Coord
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
!Coord
c])
    Int -> IO ()
forall a. Show a => a -> IO ()
print ([Int] -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product (Int -> [Int] -> [Int]
forall a. Ord a => Int -> [a] -> [a]
top Int
3 (BasinIds -> [Int]
basinSizes BasinIds
basinIds)))

-- | Returns the @n@ largest elements of a list
top :: Ord a => Int -> [a] -> [a]
top :: forall a. Ord a => Int -> [a] -> [a]
top Int
n = Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
take Int
n ([a] -> [a]) -> ([a] -> [a]) -> [a] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> a -> Ordering) -> [a] -> [a]
forall a. (a -> a -> Ordering) -> [a] -> [a]
sortBy ((a -> a -> Ordering) -> a -> a -> Ordering
forall a b c. (a -> b -> c) -> b -> a -> c
flip a -> a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare)

-- | Convert that ASCII grid into a grid of heights
heightArray :: UArray Coord Char -> Heights
heightArray :: UArray Coord Char -> Heights
heightArray UArray Coord Char
a = (Coord, Coord) -> [(Coord, Maybe Int)] -> Heights
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
(i, i) -> [(i, e)] -> a i e
array (UArray Coord Char -> (Coord, Coord)
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> (i, i)
bounds UArray Coord Char
a) [(Coord
c, Char -> Maybe Int
cvt Char
x) | (Coord
c, Char
x) <- UArray Coord Char -> [(Coord, Char)]
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> [(i, e)]
assocs UArray Coord Char
a]
  where
    cvt :: Char -> Maybe Int
cvt Char
'9' = Maybe Int
forall a. Maybe a
Nothing -- 9s are defined to be out of bounds in the problem
    cvt Char
x   = Int -> Maybe Int
forall a. a -> Maybe a
Just (Char -> Int
digitToInt Char
x)

heightsToBasinIds :: Heights -> BasinIds
heightsToBasinIds :: Heights -> BasinIds
heightsToBasinIds Heights
heights = BasinIds
basinIds
  where
    basinIds :: BasinIds
basinIds = (Coord, Coord) -> [(Coord, Maybe Coord)] -> BasinIds
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
(i, i) -> [(i, e)] -> a i e
array (Heights -> (Coord, Coord)
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> (i, i)
bounds Heights
heights) [(Coord
c, Coord -> Int -> Maybe Coord
basinId Coord
c (Int -> Maybe Coord) -> Maybe Int -> Maybe Coord
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Maybe Int
mbh) | (Coord
c, Maybe Int
mbh) <- Heights -> [(Coord, Maybe Int)]
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> [(i, e)]
assocs Heights
heights]
    basinId :: Coord -> Int -> Maybe Coord
basinId Coord
c Int
h =
     do [Coord]
xs <- [Maybe Coord] -> Maybe [Coord]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence [BasinIds
basinIdsBasinIds -> Coord -> Maybe Coord
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
!Coord
x | Coord
x <- Coord -> [Coord]
cardinal Coord
c, Just (Just Int
xh) <- [Heights -> Coord -> Maybe (Maybe Int)
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> Maybe e
arrIx Heights
heights Coord
x], Int
xhInt -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<Int
h]
        case [Coord]
xs of
          []                  -> Coord -> Maybe Coord
forall a. a -> Maybe a
Just Coord
c -- this is the lowest point in the basin
          Coord
y:[Coord]
ys | (Coord -> Bool) -> [Coord] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Coord
yCoord -> Coord -> Bool
forall a. Eq a => a -> a -> Bool
==) [Coord]
ys -> Coord -> Maybe Coord
forall a. a -> Maybe a
Just Coord
y -- all flows go to a unique basin
               | Bool
otherwise    -> Maybe Coord
forall a. Maybe a
Nothing -- no unique basin

basinSizes :: BasinIds -> [Int]
basinSizes :: BasinIds -> [Int]
basinSizes = Map Coord Int -> [Int]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList (Map Coord Int -> [Int])
-> (BasinIds -> Map Coord Int) -> BasinIds -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Coord] -> Map Coord Int
forall (f :: * -> *) a. (Foldable f, Ord a) => f a -> Map a Int
counts ([Coord] -> Map Coord Int)
-> (BasinIds -> [Coord]) -> BasinIds -> Map Coord Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Maybe Coord] -> [Coord]
forall a. [Maybe a] -> [a]
catMaybes ([Maybe Coord] -> [Coord])
-> (BasinIds -> [Maybe Coord]) -> BasinIds -> [Coord]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasinIds -> [Maybe Coord]
forall (a :: * -> * -> *) e i. (IArray a e, Ix i) => a i e -> [e]
elems