{-|
Module:         Day14
Description:    <https://adventofcode.com/2021/day/14 Day 14: Extended Polymerization>
-}
{-# LANGUAGE OverloadedStrings, TupleSections #-}
module Day14 (day14a, day14b) where

import Control.Arrow ((&&&))
import Control.Monad (guard)
import qualified Data.Map as Map ((!), elems, fromList, fromListWith, size, toList)
import Data.Semigroup (Max(..), Min(..))
import qualified Data.Set as Set (fromList, size)
import Data.Text (Text)
import qualified Data.Text as T (head, last, length, lines, null, stripPrefix, tail, uncons, unpack, zip)

day14 :: Text -> Maybe [Int]
day14 :: Text -> Maybe [Int]
day14 Text
input = do
    Text
initial : Text
e : [Text]
rest <- [Text] -> Maybe [Text]
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Text] -> Maybe [Text]) -> [Text] -> Maybe [Text]
forall a b. (a -> b) -> a -> b
$ Text -> [Text]
T.lines Text
input
    Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ Text -> Bool
T.null Text
e Bool -> Bool -> Bool
&& Text -> Int
T.length Text
initial Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
2
    Map (Char, Char) [(Char, Char)]
rules <- [((Char, Char), [(Char, Char)])] -> Map (Char, Char) [(Char, Char)]
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([((Char, Char), [(Char, Char)])]
 -> Map (Char, Char) [(Char, Char)])
-> Maybe [((Char, Char), [(Char, Char)])]
-> Maybe (Map (Char, Char) [(Char, Char)])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Text -> Maybe ((Char, Char), [(Char, Char)]))
-> [Text] -> Maybe [((Char, Char), [(Char, Char)])]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Text -> Maybe ((Char, Char), [(Char, Char)])
parseRule [Text]
rest
    let nChars :: Int
nChars = Set Char -> Int
forall a. Set a -> Int
Set.size (Set Char -> Int) -> Set Char -> Int
forall a b. (a -> b) -> a -> b
$ [Char] -> Set Char
forall a. Ord a => [a] -> Set a
Set.fromList ([Char] -> Set Char) -> [Char] -> Set Char
forall a b. (a -> b) -> a -> b
$ Text -> [Char]
T.unpack Text
initial [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ do
            ((Char
x, Char
y), [(Char, Char)]
z) <- Map (Char, Char) [(Char, Char)] -> [((Char, Char), [(Char, Char)])]
forall k a. Map k a -> [(k, a)]
Map.toList Map (Char, Char) [(Char, Char)]
rules
            Char
x Char -> [Char] -> [Char]
forall a. a -> [a] -> [a]
: Char
y Char -> [Char] -> [Char]
forall a. a -> [a] -> [a]
: ((Char, Char) -> [Char]) -> [(Char, Char)] -> [Char]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (\(Char
a, Char
b) -> [Char
a, Char
b]) [(Char, Char)]
z
    Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ Map (Char, Char) [(Char, Char)] -> Int
forall k a. Map k a -> Int
Map.size Map (Char, Char) [(Char, Char)]
rules Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
nChars Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
nChars
    let state0 :: Map (Char, Char) Int
state0 = (Int -> Int -> Int)
-> [((Char, Char), Int)] -> Map (Char, Char) Int
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
Map.fromListWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) ([((Char, Char), Int)] -> Map (Char, Char) Int)
-> [((Char, Char), Int)] -> Map (Char, Char) Int
forall a b. (a -> b) -> a -> b
$ (, Int
1) ((Char, Char) -> ((Char, Char), Int))
-> [(Char, Char)] -> [((Char, Char), Int)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Text -> Text -> [(Char, Char)]
T.zip Text
initial (Text -> Text
T.tail Text
initial)
        step :: Map (Char, Char) a -> Map (Char, Char) a
step Map (Char, Char) a
state = (a -> a -> a) -> [((Char, Char), a)] -> Map (Char, Char) a
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
Map.fromListWith a -> a -> a
forall a. Num a => a -> a -> a
(+)
          [ ((Char, Char)
dst, a
n)
          | ((Char, Char)
src, a
n) <- Map (Char, Char) a -> [((Char, Char), a)]
forall k a. Map k a -> [(k, a)]
Map.toList Map (Char, Char) a
state
          , (Char, Char)
dst <- Map (Char, Char) [(Char, Char)]
rules Map (Char, Char) [(Char, Char)] -> (Char, Char) -> [(Char, Char)]
forall k a. Ord k => Map k a -> k -> a
Map.! (Char, Char)
src
          ]
        extract :: Map (Char, Char) a -> a
extract Map (Char, Char) a
state = (a
hi a -> a -> a
forall a. Num a => a -> a -> a
- a
lo) a -> a -> a
forall a. Integral a => a -> a -> a
`div` a
2 where
            (Min a
lo, Max a
hi) = [(Min a, Max a)] -> (Min a, Max a)
forall a. Monoid a => [a] -> a
mconcat ([(Min a, Max a)] -> (Min a, Max a))
-> [(Min a, Max a)] -> (Min a, Max a)
forall a b. (a -> b) -> a -> b
$ (a -> (Min a, Max a)) -> [a] -> [(Min a, Max a)]
forall a b. (a -> b) -> [a] -> [b]
map (a -> Min a
forall a. a -> Min a
Min (a -> Min a) -> (a -> Max a) -> a -> (Min a, Max a)
forall (a :: * -> * -> *) b c c'.
Arrow a =>
a b c -> a b c' -> a b (c, c')
&&& a -> Max a
forall a. a -> Max a
Max) ([a] -> [(Min a, Max a)]) -> [a] -> [(Min a, Max a)]
forall a b. (a -> b) -> a -> b
$ Map Char a -> [a]
forall k a. Map k a -> [a]
Map.elems (Map Char a -> [a]) -> Map Char a -> [a]
forall a b. (a -> b) -> a -> b
$ (a -> a -> a) -> [(Char, a)] -> Map Char a
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
Map.fromListWith a -> a -> a
forall a. Num a => a -> a -> a
(+) ([(Char, a)] -> Map Char a) -> [(Char, a)] -> Map Char a
forall a b. (a -> b) -> a -> b
$
                (Text -> Char
T.head Text
initial, a
1) (Char, a) -> [(Char, a)] -> [(Char, a)]
forall a. a -> [a] -> [a]
: (Text -> Char
T.last Text
initial, a
1) (Char, a) -> [(Char, a)] -> [(Char, a)]
forall a. a -> [a] -> [a]
:
                    [(Char
c, a
n) | ((Char
a, Char
b), a
n) <- Map (Char, Char) a -> [((Char, Char), a)]
forall k a. Map k a -> [(k, a)]
Map.toList Map (Char, Char) a
state, Char
c <- [Char
a, Char
b]]
    [Int] -> Maybe [Int]
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Int] -> Maybe [Int]) -> [Int] -> Maybe [Int]
forall a b. (a -> b) -> a -> b
$ Map (Char, Char) Int -> Int
forall a. (Integral a, Bounded a) => Map (Char, Char) a -> a
extract (Map (Char, Char) Int -> Int) -> [Map (Char, Char) Int] -> [Int]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Map (Char, Char) Int -> Map (Char, Char) Int)
-> Map (Char, Char) Int -> [Map (Char, Char) Int]
forall a. (a -> a) -> a -> [a]
iterate Map (Char, Char) Int -> Map (Char, Char) Int
forall a. Num a => Map (Char, Char) a -> Map (Char, Char) a
step Map (Char, Char) Int
state0
  where
    parseRule :: Text -> Maybe ((Char, Char), [(Char, Char)])
parseRule Text
line = do
        (Char
x, Text
line') <- Text -> Maybe (Char, Text)
T.uncons Text
line
        (Char
y, Text
line'') <- Text -> Maybe (Char, Text)
T.uncons Text
line'
        (Char
z, Text
line''') <- Text -> Maybe (Char, Text)
T.uncons (Text -> Maybe (Char, Text)) -> Maybe Text -> Maybe (Char, Text)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Text -> Text -> Maybe Text
T.stripPrefix Text
" -> " Text
line''
        Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ Text -> Bool
T.null Text
line'''
        ((Char, Char), [(Char, Char)])
-> Maybe ((Char, Char), [(Char, Char)])
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Char
x, Char
y), [(Char
x, Char
z), (Char
z, Char
y)])

day14a :: Text -> Maybe Int
day14a :: Text -> Maybe Int
day14a = ([Int] -> Int) -> Maybe [Int] -> Maybe Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ([Int] -> Int -> Int
forall a. [a] -> Int -> a
!! Int
10) (Maybe [Int] -> Maybe Int)
-> (Text -> Maybe [Int]) -> Text -> Maybe Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Maybe [Int]
day14

day14b :: Text -> Maybe Int
day14b :: Text -> Maybe Int
day14b = ([Int] -> Int) -> Maybe [Int] -> Maybe Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ([Int] -> Int -> Int
forall a. [a] -> Int -> a
!! Int
40) (Maybe [Int] -> Maybe Int)
-> (Text -> Maybe [Int]) -> Text -> Maybe Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Maybe [Int]
day14