
  0 | module Days.Day11
  1 | 
  2 | import Data.String
  3 | import Data.List1
  4 | import Data.Nat
  5 | import Data.SortedMap
  6 | import Control.Monad.State
  7 | import Control.Monad.Identity
  8 | import System
  9 | 
 10 | import Util
 11 | 
 12 | ListHead : Type -> Type
 13 | ListHead a = Lazy (List a) -> List a
 14 | 
 15 | infixl 8 <::>
 16 | (<::>) : Lazy (ListHead a) -> Lazy (ListHead a) -> ListHead a
 17 | xs <::> ys = (\x => xs (ys x))
 18 | 
 19 | Input : Type
 20 | Input = List Integer
 21 | 
 22 | parseInput : String -> Maybe Input
 23 | parseInput str =
 24 |   traverse parsePositive . map trim . forget . split (== ' ') . trim $ str
 25 | 
 26 | digitCount : Integer -> Integer
 27 | digitCount i = digitCount' i 1
 28 |   where
 29 |     digitCount' : (i, acc : Integer) -> Integer
 30 |     digitCount' i acc =
 31 |       if i < 10
 32 |         then acc
 33 |         else digitCount' (i `div` 10) (acc + 1)
 34 | 
 35 | splitDigits : Integer -> Maybe (Integer, Integer)
 36 | splitDigits i =
 37 |   if (digitCount i `mod` 2) == 0
 38 |     then
 39 |       let new_len = integerToNat $ (digitCount i) `div` 2
 40 |           diver = natToInteger $ 10 `power` new_len
 41 |       in Just (i `div` diver, i `mod` diver)
 42 |     else Nothing
 43 | 
 44 | -- Blink an individual stone, counting how many stones result
 45 | blinkStone : (stone : Integer) -> (times : Nat)
 46 |            -> State (SortedMap (Integer, Nat) Nat) Nat
 47 | blinkStone stone 0 = pure 1
 48 | blinkStone stone (S times) = do
 49 |   cache <- get
 50 |   case lookup (stone, (S times)) cache of
 51 |     Just x => pure x
 52 |     Nothing =>
 53 |       if stone == 0
 54 |         then do
 55 |           rest <- blinkStone 1 times
 56 |           modify $ insert (stone, (S times)) rest
 57 |           pure rest
 58 |         else case splitDigits stone of
 59 |           Just (x, y) => do
 60 |             rest_1 <- blinkStone x times
 61 |             rest_2 <- blinkStone y times
 62 |             modify $ insert (stone, (S times)) (rest_1 + rest_2)
 63 |             pure (rest_1 + rest_2)
 64 |           Nothing => do
 65 |             rest <- blinkStone (stone * 2024) times
 66 |             modify $ insert (stone, (S times)) rest
 67 |             pure rest
 68 | 
 69 | -- Count the resulting stones from doing n blinks
 70 | countStones : (stones : List Integer) -> (times : Nat)
 71 |             -> State (SortedMap (Integer, Nat) Nat) Nat
 72 | countStones stones times = do
 73 |   xs <- traverse (\x => blinkStone x times) stones
 74 |   pure . sum $ xs
 75 | 
 76 | export
 77 | part1 : String -> IO (String, Input)
 78 | part1 str = do
 79 |   Just input <- pure $ parseInput str
 80 |     | _ => exitFailure
 81 |   let (_, count) = runState empty (countStones input 25)
 82 |   pure (show count, input)
 83 | 
 84 | export
 85 | part2 : Input -> String -> IO String
 86 | part2 input _ = do
 87 |   let (_, count) = runState empty (countStones input 75)
 88 |   pure . show $ count
 89 |