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

<https://adventofcode.com/2020/day/7>

The problem gives us a list of rules about the immediate contents
of each color of bag. We use this to compute the transitive
closure of bag contents in order to answer queries about a shiny
gold bag.

-}
module Main (main) where

import Advent (countBy)
import Advent.Format (format)
import Data.Map (Map)
import Data.Map qualified as Map
import Data.Maybe (fromMaybe)

type Bag = String
type Rule = (String, Maybe [(Integer, String)])

------------------------------------------------------------------------

-- |
-- >>> :main
-- 268
-- 7867
main :: IO ()
IO ()
main =
  do [(Bag, Maybe [(Integer, Bag)])]
rules <- [format|2020 7 ((%s %s)! bags contain (no other bags|(%lu (%s %s)! bag(|s))&(, )).%n)*|]
     let tc :: Map Bag (Map Bag Integer)
tc = [(Bag, Maybe [(Integer, Bag)])] -> Map Bag (Map Bag Integer)
transClosBags [(Bag, Maybe [(Integer, Bag)])]
rules
         k :: Bag
k = Bag
"shiny gold"
     Int -> IO ()
forall a. Show a => a -> IO ()
print ((Map Bag Integer -> Bool) -> Map Bag (Map Bag Integer) -> Int
forall (f :: * -> *) a. Foldable f => (a -> Bool) -> f a -> Int
countBy (Bag -> Map Bag Integer -> Bool
forall k a. Ord k => k -> Map k a -> Bool
Map.member Bag
k) Map Bag (Map Bag Integer)
tc)
     Integer -> IO ()
forall a. Show a => a -> IO ()
print (Map Bag Integer -> Integer
forall a. Num a => Map Bag a -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (Map Bag (Map Bag Integer)
tc Map Bag (Map Bag Integer) -> Bag -> Map Bag Integer
forall k a. Ord k => Map k a -> k -> a
Map.! Bag
k))

transClosBags :: [Rule] -> Map Bag (Map Bag Integer)
transClosBags :: [(Bag, Maybe [(Integer, Bag)])] -> Map Bag (Map Bag Integer)
transClosBags [(Bag, Maybe [(Integer, Bag)])]
rules = Map Bag (Map Bag Integer)
tc
  where
    tc :: Map Bag (Map Bag Integer)
tc = Maybe [(Integer, Bag)] -> Map Bag Integer
expand (Maybe [(Integer, Bag)] -> Map Bag Integer)
-> Map Bag (Maybe [(Integer, Bag)]) -> Map Bag (Map Bag Integer)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(Bag, Maybe [(Integer, Bag)])] -> Map Bag (Maybe [(Integer, Bag)])
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [(Bag, Maybe [(Integer, Bag)])]
rules

    expand :: Maybe [(Integer, Bag)] -> Map Bag Integer
expand Maybe [(Integer, Bag)]
contents =
      (Integer -> Integer -> Integer)
-> [Map Bag Integer] -> Map Bag Integer
forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
(a -> a -> a) -> f (Map k a) -> Map k a
Map.unionsWith Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
(+)
        [ (Integer
nInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
*) (Integer -> Integer) -> Map Bag Integer -> Map Bag Integer
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Integer -> Integer -> Integer)
-> Bag -> Integer -> Map Bag Integer -> Map Bag Integer
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
Map.insertWith Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
(+) Bag
bag Integer
1 (Map Bag (Map Bag Integer)
tc Map Bag (Map Bag Integer) -> Bag -> Map Bag Integer
forall k a. Ord k => Map k a -> k -> a
Map.! Bag
bag)
        | (Integer
n, Bag
bag) <- [(Integer, Bag)] -> Maybe [(Integer, Bag)] -> [(Integer, Bag)]
forall a. a -> Maybe a -> a
fromMaybe [] Maybe [(Integer, Bag)]
contents]