{-|
Module:         Day12
Description:    <https://adventofcode.com/2021/day/12 Day 12: Passage Pathing>
-}
{-# LANGUAGE FlexibleContexts, MultiWayIf, OverloadedStrings, TypeApplications #-}
module Day12 (day12a, day12b) where

import Control.Monad (forM, guard, when)
import Control.Monad.Writer (execWriter, tell)
import Data.Bits (finiteBitSize, setBit, testBit)
import Data.Char (isUpper)
import Data.Containers.ListUtils (nubOrd)
import Data.Graph.Inductive (Graph, Gr, lsuc, mkGraph, nodes)
import qualified Data.Map as Map ((!), (!?), fromList, toList)
import Data.Monoid (Sum(..))
import Data.Text (Text)
import qualified Data.Text as T (all, breakOn, lines, stripPrefix)
import Data.Tuple (swap)

parse :: (Graph gr) => Text -> Maybe (Int, Int, gr Text Bool)
parse :: Text -> Maybe (Int, Int, gr Text Bool)
parse Text
input = do
    [(Text, Text)]
conn <- [Text] -> (Text -> Maybe (Text, Text)) -> Maybe [(Text, Text)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM (Text -> [Text]
T.lines Text
input) ((Text -> Maybe (Text, Text)) -> Maybe [(Text, Text)])
-> (Text -> Maybe (Text, Text)) -> Maybe [(Text, Text)]
forall a b. (a -> b) -> a -> b
$ \Text
line ->
        let (Text
a, Text
b) = Text -> Text -> (Text, Text)
T.breakOn Text
"-" Text
line in (,) Text
a (Text -> (Text, Text)) -> Maybe Text -> Maybe (Text, Text)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Text -> Text -> Maybe Text
T.stripPrefix Text
"-" Text
b
    let names :: Map Text Int
names = [(Text, Int)] -> Map Text Int
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([(Text, Int)] -> Map Text Int) -> [(Text, Int)] -> Map Text Int
forall a b. (a -> b) -> a -> b
$ ([Text] -> [Int] -> [(Text, Int)])
-> [Int] -> [Text] -> [(Text, Int)]
forall a b c. (a -> b -> c) -> b -> a -> c
flip [Text] -> [Int] -> [(Text, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0..] ([Text] -> [(Text, Int)]) -> [Text] -> [(Text, Int)]
forall a b. (a -> b) -> a -> b
$ [Text] -> [Text]
forall a. Ord a => [a] -> [a]
nubOrd ([Text] -> [Text]) -> [Text] -> [Text]
forall a b. (a -> b) -> a -> b
$ [[Text]] -> [Text]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Text
a, Text
b] | (Text
a, Text
b) <- [(Text, Text)]
conn]
    Int
start <- Map Text Int
names Map Text Int -> Text -> Maybe Int
forall k a. Ord k => Map k a -> k -> Maybe a
Map.!? Text
"start"
    Int
end <- Map Text Int
names Map Text Int -> Text -> Maybe Int
forall k a. Ord k => Map k a -> k -> Maybe a
Map.!? Text
"end"
    (Int, Int, gr Text Bool) -> Maybe (Int, Int, gr Text Bool)
forall (f :: * -> *) a. Applicative f => a -> f a
pure
      ( Int
start
      , Int
end
      , [LNode Text] -> [LEdge Bool] -> gr Text Bool
forall (gr :: * -> * -> *) a b.
Graph gr =>
[LNode a] -> [LEdge b] -> gr a b
mkGraph ((Text, Int) -> LNode Text
forall a b. (a, b) -> (b, a)
swap ((Text, Int) -> LNode Text) -> [(Text, Int)] -> [LNode Text]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map Text Int -> [(Text, Int)]
forall k a. Map k a -> [(k, a)]
Map.toList Map Text Int
names) ([LEdge Bool] -> gr Text Bool) -> [LEdge Bool] -> gr Text Bool
forall a b. (a -> b) -> a -> b
$ do
            (Text
a, Text
b) <- [(Text, Text)]
conn
            let a' :: Int
a' = Map Text Int
names Map Text Int -> Text -> Int
forall k a. Ord k => Map k a -> k -> a
Map.! Text
a
                b' :: Int
b' = Map Text Int
names Map Text Int -> Text -> Int
forall k a. Ord k => Map k a -> k -> a
Map.! Text
b
            [(Int
a', Int
b', (Char -> Bool) -> Text -> Bool
T.all Char -> Bool
isUpper Text
b), (Int
b', Int
a', (Char -> Bool) -> Text -> Bool
T.all Char -> Bool
isUpper Text
a)]
      )

walk :: (Monad m, Traversable t) => (a -> m (t a)) -> a -> m ()
walk :: (a -> m (t a)) -> a -> m ()
walk a -> m (t a)
f = a -> m ()
walk' where walk' :: a -> m ()
walk' a
a = a -> m (t a)
f a
a m (t a) -> (t a -> m ()) -> m ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (a -> m ()) -> t a -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ a -> m ()
walk'

day12 :: Bool -> Text -> Maybe Int
day12 :: Bool -> Text -> Maybe Int
day12 Bool
bonus Text
input = do
    (Int
start, Int
end, Gr Text Bool
g) <- Text -> Maybe (Int, Int, Gr Text Bool)
forall (gr :: * -> * -> *).
Graph gr =>
Text -> Maybe (Int, Int, gr Text Bool)
parse @Gr Text
input
    Bool -> Maybe () -> Maybe ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when ((Int -> Bool) -> [Int] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
bonusBit) ([Int] -> Bool) -> [Int] -> Bool
forall a b. (a -> b) -> a -> b
$ Gr Text Bool -> [Int]
forall (gr :: * -> * -> *) a b. Graph gr => gr a b -> [Int]
nodes Gr Text Bool
g) (Maybe () -> Maybe ()) -> Maybe () -> Maybe ()
forall a b. (a -> b) -> a -> b
$ [Char] -> Maybe ()
forall a. HasCallStack => [Char] -> a
error [Char]
"input too large"
    let step :: (a, Int) -> f [(a, Int)]
step (a
state, Int
i)
          | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
end = [(a, Int)]
forall a. Monoid a => a
mempty [(a, Int)] -> f () -> f [(a, Int)]
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Sum a -> f ()
forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell (a -> Sum a
forall a. a -> Sum a
Sum a
1)
          | Bool
otherwise = [(a, Int)] -> f [(a, Int)]
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([(a, Int)] -> f [(a, Int)]) -> [(a, Int)] -> f [(a, Int)]
forall a b. (a -> b) -> a -> b
$ do
                (Int
j, Bool
big) <- Gr Text Bool -> Int -> [(Int, Bool)]
forall (gr :: * -> * -> *) a b.
Graph gr =>
gr a b -> Int -> [(Int, b)]
lsuc Gr Text Bool
g Int
i
                Bool -> [()]
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> [()]) -> Bool -> [()]
forall a b. (a -> b) -> a -> b
$ Int
j Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
start
                if
                  | Bool
big -> (a, Int) -> [(a, Int)]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
state, Int
j)
                  | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ a -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit a
state Int
j -> (a, Int) -> [(a, Int)]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> Int -> a
forall a. Bits a => a -> Int -> a
setBit a
state Int
j, Int
j)
                  | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ a -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit a
state Int
bonusBit -> (a, Int) -> [(a, Int)]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> Int -> a
forall a. Bits a => a -> Int -> a
setBit a
state Int
bonusBit, Int
j)
                  | Bool
otherwise -> [(a, Int)]
forall a. Monoid a => a
mempty
    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
$ Sum Int -> Int
forall a. Sum a -> a
getSum (Sum Int -> Int) -> Sum Int -> Int
forall a b. (a -> b) -> a -> b
$ Writer (Sum Int) () -> Sum Int
forall w a. Writer w a -> w
execWriter (Writer (Sum Int) () -> Sum Int) -> Writer (Sum Int) () -> Sum Int
forall a b. (a -> b) -> a -> b
$ ((Int, Int) -> WriterT (Sum Int) Identity [(Int, Int)])
-> (Int, Int) -> Writer (Sum Int) ()
forall (m :: * -> *) (t :: * -> *) a.
(Monad m, Traversable t) =>
(a -> m (t a)) -> a -> m ()
walk (Int, Int) -> WriterT (Sum Int) Identity [(Int, Int)]
forall (f :: * -> *) a a.
(MonadWriter (Sum a) f, Num a, Bits a) =>
(a, Int) -> f [(a, Int)]
step (if Bool
bonus then Int
0 else Int -> Int -> Int
forall a. Bits a => a -> Int -> a
setBit Int
0 Int
bonusBit :: Int, Int
start)
  where bonusBit :: Int
bonusBit = Int -> Int
forall b. FiniteBits b => b -> Int
finiteBitSize (Int
0 :: Int) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1

day12a :: Text -> Maybe Int
day12a :: Text -> Maybe Int
day12a = Bool -> Text -> Maybe Int
day12 Bool
False

day12b :: Text -> Maybe Int
day12b :: Text -> Maybe Int
day12b = Bool -> Text -> Maybe Int
day12 Bool
True