{-# LANGUAGE OverloadedStrings #-} -- for building Text strings
module Main where
import Control.Applicative (Alternative((<|>)))
import Control.Monad (void, when)
import Data.Text.Lazy (Text)
import Data.Void (Void)
import Prelude
import qualified Data.Text.Lazy.Encoding as Text
import qualified Text.Megaparsec as P
import qualified Text.Megaparsec.Char as P
import qualified Text.Megaparsec.Char.Lexer as P
-- import qualified Text.Megaparsec.Debug as P
import qualified Data.ByteString.Lazy as BSL
import Paths_AoC2020
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map

-- | See https://adventofcode.com/2020/day/7 for the problem statements
data Day07Results = Day07Results
  { example1 :: Int
  , batch1 :: Int
  , example2 :: Int
  , example3 :: Int
  , batch2 :: Int
  } deriving (Show)

main :: IO ()
main = do
  putStr "Day07Inputs: " >> getDataFileName "" >>= putStrLn
  print =<< Day07Results
    <$> (countTransitiveParentBags "shiny gold" <$> parse "example1")
    <*> (countTransitiveParentBags "shiny gold" <$> parse "batch")
    <*> (countChildrenBags "shiny gold" <$> parse "example1")
    <*> (countChildrenBags "shiny gold" <$> parse "example2")
    <*> (countChildrenBags "shiny gold" <$> parse "batch")

type Parser output = P.Parsec {-error-}Void {-input-}Text output

type Bag = Text
type Rule = (Bag, Map Bag Int)
type Rules = Map Bag (Map Bag Int)

parserBag :: Parser Bag
parserBag =
  (fst <$>) $ P.match $ do
    void $ P.takeWhileP (Just "adj") (/=' ')
    void $ P.char ' '
    void $ P.takeWhileP (Just "col") (/=' ')

parserRule :: Parser Rule
parserRule = {- P.dbg "parserRule" $ -} do
  outBag <- parserBag
  void $ P.string " bags contain "
  inBags <- P.choice
    [ [] <$ P.string "no other bags"
    , (`P.sepBy` P.string ", ") $ do
        cnt <- P.decimal
        void $ P.char ' '
        bag <- parserBag
        void $ P.string " bag"
        when (cnt > 1) $
          void $ P.optional (P.char 's')
        return (bag, cnt)
    ]
  void $ P.char '.'
  return (outBag, Map.fromList inBags)

parserRules :: Parser Rules
parserRules =
  -- P.dbg "parserRules" $
  (Map.fromList <$>) $
  P.many $ parserRule <* (() <$ P.char '\n' <|> P.eof)

parse :: FilePath -> IO Rules
parse input = do
  content <- Text.decodeUtf8 <$> (BSL.readFile =<< getDataFileName input)
  case P.parse (parserRules <* P.eof) input content of
    Left err -> error (P.errorBundlePretty err)
    Right rules -> return rules

type ReverseRules = Map Bag (Map Bag ())

reverseRules :: Rules -> ReverseRules
reverseRules =
  Map.foldlWithKey (\acc outBag inBags ->
    Map.unionWith (<>) acc (Map.singleton outBag () <$ inBags)
  ) Map.empty

parentBags :: Bag -> ReverseRules -> Map Bag ()
parentBags = Map.findWithDefault Map.empty

transitiveParentBags :: ReverseRules -> Bag -> Map Bag ()
transitiveParentBags revRules bag =
  parentBags bag revRules <>
  Map.unions (transitiveParentBags revRules <$> Map.keys (parentBags bag revRules))

countTransitiveParentBags :: Bag -> Rules -> Int
countTransitiveParentBags bag rules = Map.size $ transitiveParentBags (reverseRules rules) bag

countChildrenBags :: Bag -> Rules -> Int
countChildrenBags bag rules =
  sum $ Map.mapWithKey (\childBag childBagCount ->
    childBagCount + childBagCount * countChildrenBags childBag rules
  ) childrenBags
  where childrenBags = Map.findWithDefault Map.empty bag rules