{-|
Module:         Day18
Description:    <https://adventofcode.com/2021/day/18 Day 18: Snailfish>
-}
{-# LANGUAGE TypeFamilies, ViewPatterns #-}
module Day18 (day18a, day18b) where

import Data.List (inits, tails)
import Data.List.NonEmpty (nonEmpty)
import Data.Text (Text)
import Data.Void (Void)
import Text.Megaparsec (MonadParsec, ParseErrorBundle, (<|>), between, eof, parse, sepEndBy, single, try)
import qualified Text.Megaparsec as P (Token)
import Text.Megaparsec.Char (newline)
import Text.Megaparsec.Char.Lexer (decimal)

data Token a = Open | Close | Value a

parser :: (MonadParsec e s m, P.Token s ~ Char, Num a) => m [Token a]
parser :: m [Token a]
parser = m Char -> m Char -> m [Token a] -> m [Token a]
forall (m :: * -> *) open close a.
Applicative m =>
m open -> m close -> m a -> m a
between (Token s -> m (Token s)
forall e s (m :: * -> *).
MonadParsec e s m =>
Token s -> m (Token s)
single Char
Token s
'[') (Token s -> m (Token s)
forall e s (m :: * -> *).
MonadParsec e s m =>
Token s -> m (Token s)
single Char
Token s
']') (m [Token a] -> m [Token a]) -> m [Token a] -> m [Token a]
forall a b. (a -> b) -> a -> b
$ do
    [Token a]
lhs <- m [Token a]
parser' m [Token a] -> m Char -> m [Token a]
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f a
<* Token s -> m (Token s)
forall e s (m :: * -> *).
MonadParsec e s m =>
Token s -> m (Token s)
single Char
Token s
','
    [Token a]
rhs <- m [Token a]
parser'
    [Token a] -> m [Token a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Token a] -> m [Token a]) -> [Token a] -> m [Token a]
forall a b. (a -> b) -> a -> b
$ Token a
forall a. Token a
Open Token a -> [Token a] -> [Token a]
forall a. a -> [a] -> [a]
: [Token a]
lhs [Token a] -> [Token a] -> [Token a]
forall a. [a] -> [a] -> [a]
++ [Token a]
rhs [Token a] -> [Token a] -> [Token a]
forall a. [a] -> [a] -> [a]
++ [Token a
forall a. Token a
Close]
  where parser' :: m [Token a]
parser' = (Token a -> [Token a] -> [Token a]
forall a. a -> [a] -> [a]
:[]) (Token a -> [Token a]) -> (a -> Token a) -> a -> [Token a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Token a
forall a. a -> Token a
Value (a -> [Token a]) -> m a -> m [Token a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m a -> m a
forall e s (m :: * -> *) a. MonadParsec e s m => m a -> m a
try m a
forall e s (m :: * -> *) a.
(MonadParsec e s m, Token s ~ Char, Num a) =>
m a
decimal m [Token a] -> m [Token a] -> m [Token a]
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> m [Token a]
forall e s (m :: * -> *) a.
(MonadParsec e s m, Token s ~ Char, Num a) =>
m [Token a]
parser

add :: (Integral a, Ord a) => [Token a] -> [Token a] -> [Token a]
add :: [Token a] -> [Token a] -> [Token a]
add [Token a]
lhs [Token a]
rhs = [Token a] -> [Token a]
forall a. Integral a => [Token a] -> [Token a]
explode ([Token a] -> [Token a]) -> [Token a] -> [Token a]
forall a b. (a -> b) -> a -> b
$ Token a
forall a. Token a
Open Token a -> [Token a] -> [Token a]
forall a. a -> [a] -> [a]
: [Token a]
lhs [Token a] -> [Token a] -> [Token a]
forall a. [a] -> [a] -> [a]
++ [Token a]
rhs [Token a] -> [Token a] -> [Token a]
forall a. [a] -> [a] -> [a]
++ [Token a
forall a. Token a
Close] where
    explode :: [Token a] -> [Token a]
explode [Token a]
tokens = Int -> [Token a] -> [Token a] -> [Token a]
explode' (Int
0 :: Int) [] [Token a]
tokens where
        explode' :: Int -> [Token a] -> [Token a] -> [Token a]
explode' Int
n [Token a]
pre (Token a
Open : Value a
x : Value a
y : Token a
Close : [Token a]
post) | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
4 = [Token a] -> [Token a]
explode ([Token a] -> [Token a]) -> [Token a] -> [Token a]
forall a b. (a -> b) -> a -> b
$
            [Token a] -> [Token a]
forall a. [a] -> [a]
reverse ((a -> a) -> [Token a] -> [Token a]
forall a. (a -> a) -> [Token a] -> [Token a]
modifyFirstValue (a -> a -> a
forall a. Num a => a -> a -> a
+ a
x) [Token a]
pre) [Token a] -> [Token a] -> [Token a]
forall a. [a] -> [a] -> [a]
++ a -> Token a
forall a. a -> Token a
Value a
0 Token a -> [Token a] -> [Token a]
forall a. a -> [a] -> [a]
: (a -> a) -> [Token a] -> [Token a]
forall a. (a -> a) -> [Token a] -> [Token a]
modifyFirstValue (a -> a -> a
forall a. Num a => a -> a -> a
+ a
y) [Token a]
post
        explode' Int
n [Token a]
pre (Token a
cur : [Token a]
post) = Int -> [Token a] -> [Token a] -> [Token a]
explode' Int
n' (Token a
cur Token a -> [Token a] -> [Token a]
forall a. a -> [a] -> [a]
: [Token a]
pre) [Token a]
post where
            n' :: Int
n' | Token a
Open <- Token a
cur = Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 | Token a
Close <- Token a
cur = Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 | Bool
otherwise = Int
n
        explode' Int
_ [Token a]
_ [] = [Token a] -> [Token a]
split [Token a]
tokens
    modifyFirstValue :: (a -> a) -> [Token a] -> [Token a]
modifyFirstValue a -> a
f (Value a
x : [Token a]
rest) = a -> Token a
forall a. a -> Token a
Value (a -> a
f a
x) Token a -> [Token a] -> [Token a]
forall a. a -> [a] -> [a]
: [Token a]
rest
    modifyFirstValue a -> a
f (Token a
x : [Token a]
rest) = Token a
x Token a -> [Token a] -> [Token a]
forall a. a -> [a] -> [a]
: (a -> a) -> [Token a] -> [Token a]
modifyFirstValue a -> a
f [Token a]
rest
    modifyFirstValue a -> a
_ [] = []
    split :: [Token a] -> [Token a]
split [Token a]
tokens = [Token a] -> [Token a] -> [Token a]
split' [] [Token a]
tokens where
        split' :: [Token a] -> [Token a] -> [Token a]
split' [Token a]
pre (Value a
x : [Token a]
post) | a
x a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> a
9 = [Token a] -> [Token a]
explode ([Token a] -> [Token a]) -> [Token a] -> [Token a]
forall a b. (a -> b) -> a -> b
$
            [Token a] -> [Token a]
forall a. [a] -> [a]
reverse [Token a]
pre [Token a] -> [Token a] -> [Token a]
forall a. [a] -> [a] -> [a]
++ Token a
forall a. Token a
Open Token a -> [Token a] -> [Token a]
forall a. a -> [a] -> [a]
: a -> Token a
forall a. a -> Token a
Value (a
x a -> a -> a
forall a. Integral a => a -> a -> a
`div` a
2) Token a -> [Token a] -> [Token a]
forall a. a -> [a] -> [a]
: a -> Token a
forall a. a -> Token a
Value ((a
x a -> a -> a
forall a. Num a => a -> a -> a
+ a
1) a -> a -> a
forall a. Integral a => a -> a -> a
`div` a
2) Token a -> [Token a] -> [Token a]
forall a. a -> [a] -> [a]
: Token a
forall a. Token a
Close Token a -> [Token a] -> [Token a]
forall a. a -> [a] -> [a]
: [Token a]
post
        split' [Token a]
pre (Token a
cur : [Token a]
post) = [Token a] -> [Token a] -> [Token a]
split' (Token a
cur Token a -> [Token a] -> [Token a]
forall a. a -> [a] -> [a]
: [Token a]
pre) [Token a]
post
        split' [Token a]
_ [] = [Token a]
tokens

magnitude :: (Num a) => [Token a] -> Maybe a
magnitude :: [Token a] -> Maybe a
magnitude [Token a]
input
  | (Maybe a
result, []) <- [Token a] -> (Maybe a, [Token a])
forall a. Num a => [Token a] -> (Maybe a, [Token a])
magnitude' [Token a]
input = Maybe a
result
  | Bool
otherwise = Maybe a
forall a. Maybe a
Nothing
  where
    magnitude' :: [Token a] -> (Maybe a, [Token a])
magnitude' (Token a
Open : ([Token a] -> (Maybe a, [Token a])
magnitude' -> (Just a
lhs, [Token a] -> (Maybe a, [Token a])
magnitude' -> (Just a
rhs, Token a
Close : [Token a]
rest)))) =
        (a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> a -> Maybe a
forall a b. (a -> b) -> a -> b
$ a
3 a -> a -> a
forall a. Num a => a -> a -> a
* a
lhs a -> a -> a
forall a. Num a => a -> a -> a
+ a
2 a -> a -> a
forall a. Num a => a -> a -> a
* a
rhs, [Token a]
rest)
    magnitude' (Value a
value : [Token a]
rest) = (a -> Maybe a
forall a. a -> Maybe a
Just a
value, [Token a]
rest)
    magnitude' [Token a]
rest = (Maybe a
forall a. Maybe a
Nothing, [Token a]
rest)

day18a :: Text -> Either (ParseErrorBundle Text Void) (Maybe Int)
day18a :: Text -> Either (ParseErrorBundle Text Void) (Maybe Int)
day18a Text
input = do
    [[Token Int]]
list <- Parsec Void Text [[Token Int]]
-> String
-> Text
-> Either (ParseErrorBundle Text Void) [[Token Int]]
forall e s a.
Parsec e s a -> String -> s -> Either (ParseErrorBundle s e) a
parse (ParsecT Void Text Identity [Token Int]
forall e s (m :: * -> *) a.
(MonadParsec e s m, Token s ~ Char, Num a) =>
m [Token a]
parser ParsecT Void Text Identity [Token Int]
-> ParsecT Void Text Identity Char
-> Parsec Void Text [[Token Int]]
forall (m :: * -> *) a sep. MonadPlus m => m a -> m sep -> m [a]
`sepEndBy` ParsecT Void Text Identity Char
forall e s (m :: * -> *).
(MonadParsec e s m, Token s ~ Char) =>
m (Token s)
newline Parsec Void Text [[Token Int]]
-> ParsecT Void Text Identity () -> Parsec Void Text [[Token Int]]
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f a
<* ParsecT Void Text Identity ()
forall e s (m :: * -> *). MonadParsec e s m => m ()
eof) String
"" Text
input
    Maybe Int -> Either (ParseErrorBundle Text Void) (Maybe Int)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe Int -> Either (ParseErrorBundle Text Void) (Maybe Int))
-> Maybe Int -> Either (ParseErrorBundle Text Void) (Maybe Int)
forall a b. (a -> b) -> a -> b
$ [[Token Int]] -> Maybe (NonEmpty [Token Int])
forall a. [a] -> Maybe (NonEmpty a)
nonEmpty [[Token Int]]
list Maybe (NonEmpty [Token Int])
-> (NonEmpty [Token Int] -> Maybe Int) -> Maybe Int
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= [Token Int] -> Maybe Int
forall a. Num a => [Token a] -> Maybe a
magnitude ([Token Int] -> Maybe Int)
-> (NonEmpty [Token Int] -> [Token Int])
-> NonEmpty [Token Int]
-> Maybe Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([Token Int] -> [Token Int] -> [Token Int])
-> NonEmpty [Token Int] -> [Token Int]
forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldl1 [Token Int] -> [Token Int] -> [Token Int]
forall a.
(Integral a, Ord a) =>
[Token a] -> [Token a] -> [Token a]
add

day18b :: Text -> Either (ParseErrorBundle Text Void) (Maybe Int)
day18b :: Text -> Either (ParseErrorBundle Text Void) (Maybe Int)
day18b Text
input = do
    [[Token Int]]
list <- Parsec Void Text [[Token Int]]
-> String
-> Text
-> Either (ParseErrorBundle Text Void) [[Token Int]]
forall e s a.
Parsec e s a -> String -> s -> Either (ParseErrorBundle s e) a
parse (ParsecT Void Text Identity [Token Int]
forall e s (m :: * -> *) a.
(MonadParsec e s m, Token s ~ Char, Num a) =>
m [Token a]
parser ParsecT Void Text Identity [Token Int]
-> ParsecT Void Text Identity Char
-> Parsec Void Text [[Token Int]]
forall (m :: * -> *) a sep. MonadPlus m => m a -> m sep -> m [a]
`sepEndBy` ParsecT Void Text Identity Char
forall e s (m :: * -> *).
(MonadParsec e s m, Token s ~ Char) =>
m (Token s)
newline Parsec Void Text [[Token Int]]
-> ParsecT Void Text Identity () -> Parsec Void Text [[Token Int]]
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f a
<* ParsecT Void Text Identity ()
forall e s (m :: * -> *). MonadParsec e s m => m ()
eof) String
"" Text
input
    let sums :: [[Token Int]]
sums = do
            ([[Token Int]]
pre, [Token Int]
x : [[Token Int]]
post) <- [[[Token Int]]]
-> [[[Token Int]]] -> [([[Token Int]], [[Token Int]])]
forall a b. [a] -> [b] -> [(a, b)]
zip ([[Token Int]] -> [[[Token Int]]]
forall a. [a] -> [[a]]
inits [[Token Int]]
list) ([[Token Int]] -> [[[Token Int]]]
forall a. [a] -> [[a]]
tails [[Token Int]]
list)
            [Token Int]
y <- [[Token Int]]
pre [[Token Int]] -> [[Token Int]] -> [[Token Int]]
forall a. [a] -> [a] -> [a]
++ [[Token Int]]
post
            [[Token Int] -> [Token Int] -> [Token Int]
forall a.
(Integral a, Ord a) =>
[Token a] -> [Token a] -> [Token a]
add [Token Int]
x [Token Int]
y, [Token Int] -> [Token Int] -> [Token Int]
forall a.
(Integral a, Ord a) =>
[Token a] -> [Token a] -> [Token a]
add [Token Int]
y [Token Int]
x]
    Maybe Int -> Either (ParseErrorBundle Text Void) (Maybe Int)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe Int -> Either (ParseErrorBundle Text Void) (Maybe Int))
-> Maybe Int -> Either (ParseErrorBundle Text Void) (Maybe Int)
forall a b. (a -> b) -> a -> b
$ (NonEmpty Int -> Int) -> Maybe (NonEmpty Int) -> Maybe Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap NonEmpty Int -> Int
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum (Maybe (NonEmpty Int) -> Maybe Int)
-> Maybe (NonEmpty Int) -> Maybe Int
forall a b. (a -> b) -> a -> b
$ ([Token Int] -> Maybe Int) -> [[Token Int]] -> Maybe [Int]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM [Token Int] -> Maybe Int
forall a. Num a => [Token a] -> Maybe a
magnitude [[Token Int]]
sums Maybe [Int]
-> ([Int] -> Maybe (NonEmpty Int)) -> Maybe (NonEmpty Int)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= [Int] -> Maybe (NonEmpty Int)
forall a. [a] -> Maybe (NonEmpty a)
nonEmpty