import Data.List

main :: IO ()
main = do
    input <- readFile "6/input.txt"
    let x = parse input
    -- Problem 1
    print $ f x
    -- Problem 2
    print $ g x
    where
        f = solve 80
        g = solve 256

parse :: String -> [Int]
parse = map read . words . map f
    where f ',' = ' '
          f x   = x

-- Solve the problem by creating an association list mapping a
-- value for the internal timer to a number of fishies as the seed,
-- and then simulating `n` iterations of growth.
solve :: Int -- ^ n: The number of generations to simulate
  -> [Int]   -- ^ The initial population
  -> Int     -- ^ The total population at the end of n generations
solve n = sum . map snd . simulate n . map assoc . group . sort
    where assoc x = (head x, length x)

-- Simulate `n` iterations of growth and return the state
-- after that many iterations.
simulate :: Int -> [(Int, Int)] -> [(Int, Int)]         
simulate 0 st = st
simulate n st = simulate (n - 1) (update st)

-- Update the simulation state by one iteration.
update :: [(Int, Int)] -> [(Int, Int)]
update = merge . (>>= f)
    where f (0, n) = [(6, n), (8, n)]
          f (t, n) = [(t - 1, n)]

-- Merge the values of duplicate key entries and produce a new list
-- where each key is mapped to the sum of all values associated with
-- that key.
merge :: [(Int, Int)] -> [(Int, Int)]
merge [] = []
merge ((t,n) : ls) = (t, n + sum (map snd x)) : merge r
    where (x,r) = partition ((== t) . fst) ls