module Day11b

import Prelude
import Node
import Aoc
import SortedMap

-- Alternate version that doesn't use string

infixl 7 _%_
pfunc _%_ : Int → Int → Int := `(x,y) => x % y`

-- 32 bit ints are too small
pfunc div53 : Int → Int → Int := `(x,y) => Math.floor(x / y)`

stone : Int → Either Int (Int × Int)
stone num = if num == 0 then Left 1 else go num num 1
  where
    go : Int → Int → Int → Either Int (Int × Int)
    go a b mod =
      if b == 0 then Right (a, num % mod)
      else if b < 10 then Left (2024 * num)
      else go (div53 a 10) (div53 b 100) (mod * 10)

foldMap : ∀ a b. {{Ord a}} {{Eq a}} → (b → b → b) → SortedMap a b → List (a × b) → SortedMap a b
foldMap f m Nil = m
foldMap f m ((a,b) :: xs) = case lookupMap a m of
  Nothing => foldMap f (updateMap a b m) xs
  Just (_, b') => foldMap f (updateMap a (f b' b) m) xs

step : List (Int × Int) → List (Int × Int)
step xs = go Nil xs
  where
    go : List (Int × Int) → List (Int × Int) → List (Int × Int)
    go acc Nil = acc
    go acc ((x,c) :: xs) = case stone x of
      Left a => go ((a,c) :: acc) xs
      Right (a,b) => go ((a,c) :: (b,c) :: acc) xs

iter : Int → SortedMap Int Int → Int
iter count parts = let x = go count parts in foldl _+_ 0 $ map snd $ toList x
  where
    go : Int → SortedMap Int Int → SortedMap Int Int
    go 0 stuff = stuff
    go x stuff = go (x - 1) $ foldMap _+_ EmptyMap $ step $ toList stuff

run : String -> IO Unit
run fn = do
  putStrLn fn
  text <- readFile fn
  let stuff = foldMap _+_ EmptyMap $ map (\ x => (stringToInt x, 1)) $ split (trim text) " "
  let p1 = iter 25 stuff
  putStrLn $ "part1 " ++ show p1
  let p2 = iter 75 stuff
  putStrLn $ "part2 " ++ show p2

main : IO Unit
main = do
  run "aoc2024/day11/eg.txt"
  run "aoc2024/day11/input.txt"