{-# Language ImportQualifiedPost, QuasiQuotes #-}
{-|
Module      : Main
Description : Day 6 solution
Copyright   : (c) Eric Mertens, 2021
License     : ISC
Maintainer  : emertens@gmail.com

<https://adventofcode.com/2021/day/6>

Multiplying fish! To make this problem tractable
track how many of each age of fish we have rather
than tracking all the individual ages of the fish.

-}
module Main (main) where

import Advent (counts, format, power)
import Data.Map (Map)
import Data.Map qualified as Map

-- | >>> :main
-- 376194
-- 1693022481538
main :: IO ()
IO ()
main =
 do Map Int Int
inp <- [Int] -> Map Int Int
forall (f :: * -> *) a. (Foldable f, Ord a) => f a -> Map a Int
counts ([Int] -> Map Int Int) -> IO [Int] -> IO (Map Int Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [format|2021 6 %u&,%n|]
    let bigFish :: Int
bigFish = [Int] -> Int
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum (Map Int Int -> [Int]
forall k a. Map k a -> [k]
Map.keys Map Int Int
inp)
    let oneStep :: Map Int (Map Int Int)
oneStep = Int -> Map Int (Map Int Int)
rule Int
bigFish
    let nSteps :: Integer -> Map Int (Map Int Int)
nSteps = (Map Int (Map Int Int)
 -> Map Int (Map Int Int) -> Map Int (Map Int Int))
-> Map Int (Map Int Int) -> Integer -> Map Int (Map Int Int)
forall a. HasCallStack => (a -> a -> a) -> a -> Integer -> a
power ((Map Int Int -> Map Int Int)
-> Map Int (Map Int Int) -> Map Int (Map Int Int)
forall a b. (a -> b) -> Map Int a -> Map Int b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Map Int Int -> Map Int Int)
 -> Map Int (Map Int Int) -> Map Int (Map Int Int))
-> (Map Int (Map Int Int) -> Map Int Int -> Map Int Int)
-> Map Int (Map Int Int)
-> Map Int (Map Int Int)
-> Map Int (Map Int Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Int (Map Int Int) -> Map Int Int -> Map Int Int
forall a b.
(Ord a, Ord b) =>
Map a (Map b Int) -> Map a Int -> Map b Int
applyRule) Map Int (Map Int Int)
oneStep
    Int -> IO ()
forall a. Show a => a -> IO ()
print (Map Int Int -> Int
forall a. Num a => Map Int a -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (Integer -> Map Int (Map Int Int)
nSteps  Integer
80 Map Int (Map Int Int) -> Map Int Int -> Map Int Int
forall a b.
(Ord a, Ord b) =>
Map a (Map b Int) -> Map a Int -> Map b Int
`applyRule` Map Int Int
inp))
    Int -> IO ()
forall a. Show a => a -> IO ()
print (Map Int Int -> Int
forall a. Num a => Map Int a -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (Integer -> Map Int (Map Int Int)
nSteps Integer
256 Map Int (Map Int Int) -> Map Int Int -> Map Int Int
forall a b.
(Ord a, Ord b) =>
Map a (Map b Int) -> Map a Int -> Map b Int
`applyRule` Map Int Int
inp))

rule :: Int -> Map Int (Map Int Int)
rule :: Int -> Map Int (Map Int Int)
rule Int
n = [Int] -> Map Int Int
forall (f :: * -> *) a. (Foldable f, Ord a) => f a -> Map a Int
counts ([Int] -> Map Int Int) -> Map Int [Int] -> Map Int (Map Int Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(Int, [Int])] -> Map Int [Int]
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ((Int
0, [Int
6,Int
8]) (Int, [Int]) -> [(Int, [Int])] -> [(Int, [Int])]
forall a. a -> [a] -> [a]
: [(Int
i, [Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1]) | Int
i <- [Int
1..Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
n Int
8]])

applyRule :: (Ord a, Ord b) => Map a (Map b Int) -> Map a Int -> Map b Int
applyRule :: forall a b.
(Ord a, Ord b) =>
Map a (Map b Int) -> Map a Int -> Map b Int
applyRule Map a (Map b Int)
r Map a Int
m = (Int -> Int -> Int) -> [Map b Int] -> Map b Int
forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
(a -> a -> a) -> f (Map k a) -> Map k a
Map.unionsWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) [(Int
v Int -> Int -> Int
forall a. Num a => a -> a -> a
*) (Int -> Int) -> Map b Int -> Map b Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Map a (Map b Int)
r Map a (Map b Int) -> a -> Map b Int
forall k a. Ord k => Map k a -> k -> a
Map.! a
k) | (a
k,Int
v) <- Map a Int -> [(a, Int)]
forall k a. Map k a -> [(k, a)]
Map.toList Map a Int
m]