{-# LANGUAGE FlexibleContexts, LambdaCase, TypeApplications, ViewPatterns #-}
module Day12 (day12) where
import Control.Monad.State (evalState, gets, modify)
import Data.List (dropWhileEnd)
import Data.List.Split (divvy)
import qualified Data.Map as Map (empty, insert, lookup)
import Data.Set (Set)
import qualified Data.Set as Set (fromList, member)
import Text.Megaparsec (MonadParsec, between, count, parseMaybe, sepEndBy, some)
import Text.Megaparsec.Char (newline, oneOf, string)
parser :: (MonadParsec e String m) => m ([Bool], Set [Bool])
parser = (,) <$> between (string "initial state: ") newline (some spot)
<*> (Set.fromList . map fst . filter snd <$> (newline *> sepEndBy line newline))
where line = (,) <$> count 5 spot <*> (string " => " *> spot)
spot = (== '#') <$> oneOf "#."
day12 :: Int -> String -> Maybe Int
day12 n (parseMaybe @() parser -> Just (initial, mappings)) | n >= 0, not $ null initial =
Just $ flip evalState Map.empty $ loop 0 0 initial
where loop i offset state | i == n = return $ sum $ map fst $ filter snd $ zip [offset..] state
loop i offset state = do
modify (Map.insert state (i, offset))
let (i1, (offset1, state')) = (i + 1, step offset state)
gets (Map.lookup state') >>= \case
Just (i0, offset0) -> loop (n - (n - i1) `mod` (i1 - i0))
(offset1 + (n - i1) `div` (i1 - i0) * (offset1 - offset0)) state'
Nothing -> loop i1 offset1 state'
step offset state = (offset + add - 2, state'')
where extended = False : False : False : False : state ++ [False, False, False, False]
state' = dropWhileEnd not $ map (`Set.member` mappings) $ divvy 5 1 extended
(length -> add, state'') = span not state'
day12 _ _ = Nothing