module Main where

import Data.Functor
import Data.List.Split
import Data.IntMap (IntMap)
import qualified Data.IntMap as IM

parse :: String -> IntMap Int
parse = foldr (flip (IM.insertWith (+)) 1 . read) mempty . words

rules :: [(Int -> Bool, Int -> IntMap Int)]
rules =
  [
   ((==0), const (IM.singleton 1 1))
  ,(even . length . show,
    \a -> let b = show a; s = length b in
      foldr (flip (IM.insertWith (+)) 1 . read) mempty (chunksOf (s `div` 2) b))
  ,(const otherwise, (`IM.singleton` 1) . (*2024))
  ]

applyRule :: Int -> IntMap Int
applyRule i = head . filter (/= mempty) $ rules >>= \(f,g) -> if f i then pure $ g i else pure mempty

step :: IntMap Int -> IntMap Int
step = IM.foldrWithKey (\s c n -> IM.unionWith (+) n ((c*) <$> applyRule s)) mempty

applyAmount :: Int -> a -> (a -> a) -> a
applyAmount 0 a _ = a
applyAmount n a f = f $ applyAmount (n-1) a f

solve :: Int -> IntMap Int -> Int
solve i = sum . IM.elems . flip (applyAmount i) step

main :: IO ()
main = readFile "inputs/11" <&> parse >>= \i ->
       print (solve 25 i) >>
       print (solve 75 i)