{-# 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, String)
type Rule = (Bag, Maybe [(Integer, Bag)])

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

-- |
-- >>> :main
-- 268
-- 7867
main :: IO ()
IO ()
main =
  do rules <- [format|2020 7 ((%s %s) bags contain (no other bags|(%lu (%s %s) bag(|s))&(, )).%n)*|]
     let tc = [(Bag, Maybe [(Integer, Bag)])] -> Map Bag (Map Bag Integer)
transClosBags [(Bag, Maybe [(Integer, Bag)])]
rules
         k = ([Char]
"shiny", [Char]
"gold")
     print (countBy (Map.member k) tc)
     print (sum (tc Map.! 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]