Day19.hs

{-|
Module:         Day19
Description:    <https://adventofcode.com/2022/day/19 Day 19: Not Enough Minerals>
-}
{-# LANGUAGE OverloadedStrings, TypeFamilies #-}
module Day19 (day19a, day19b) where

import Control.Arrow (second)
import Control.Parallel.Strategies (parMap, rseq)
import Data.Char (isAlphaNum)
import qualified Data.Heap as Heap (FstMaxPolicy, insert, singleton, view)
import Data.List (foldl', scanl', transpose)
import Data.Map (Map)
import qualified Data.Map as Map ((!), (!?), elems, fromDistinctAscList, fromList, keys, insert, insertWith, toAscList, toList, unionWith, unionsWith)
import Data.Maybe (fromMaybe)
import Data.String (IsString)
import Data.Text (Text)
import Data.Void (Void)
import Debug.Trace (traceShowId)
import Text.Megaparsec (MonadParsec, ParseErrorBundle, Token, Tokens, between, eof, many, parse, sepBy1, sepEndBy, takeWhile1P)
import Text.Megaparsec.Char (char, space, string)
import qualified Text.Megaparsec.Char.Lexer as L (decimal)

parser :: (Integral a, Integral b, MonadParsec e s m, IsString (Tokens s), Token s ~ Char) => m [(a, Map (Tokens s) (Map (Tokens s) b))]
parser = many blueprint where
    blueprint = (,) <$> between (string "Blueprint ") (char ':' >> space) L.decimal <*>
        (Map.fromList <$> (robot <* char '.') `sepEndBy` space)
    robot = (,) <$> between (string "Each ") (string " robot costs ") name <*>
        (Map.fromList <$> ore `sepBy1` string " and ")
    ore = flip (,) <$> L.decimal <*> (char ' ' >> name)
    name = takeWhile1P Nothing isAlphaNum

geodes :: (IsString k, Ord k, Num v, Ord v) => Int -> Map k (Map k v) -> v
geodes n blueprint = go 0 (initialRobots, initialResources, n) where
    maxValues = Map.unionsWith max $ Map.elems blueprint
    initialRobots = Map.insert "ore" 1 $ const 0 <$> blueprint
    initialResources = const 0 <$> blueprint
    potential robots resources m = potentialResources !! m Map.! "geode" where
        potentialRobots = Map.fromDistinctAscList . zip (Map.keys blueprint) <$> transpose
          [ (robots Map.! robot +) <$> scanl' f 0 potentialResources
          | (robot, costs) <- Map.toAscList blueprint
          , let f new resources'
                  | and [resources' Map.! key >= new * cost | (key, cost) <- Map.toList costs]
                  = new + 1
                  | otherwise = new
          ]
        potentialResources = scanl' (Map.unionWith (+)) resources potentialRobots
    go k (_, _, 0) = k
    go k (robots, resources, m)
      | k > potential robots resources m = k
      | otherwise = foldl' go
          (max k $ resources Map.! "geode" + fromIntegral m * robots Map.! "geode")
          [ (robots', resources'', m - d - 1)
          | (robot, costs) <- Map.toList blueprint
          , maybe True (robots Map.! robot <) $ maxValues Map.!? robot
          , (d, resources') <- take 1 $ dropWhile (any (< 0) . snd) $
                zip [0..m - 1] $ iterate (Map.unionWith (+) robots) $
                Map.unionWith (+) resources $ negate <$> costs
          , let robots' = Map.insertWith (+) robot 1 robots
                resources'' = Map.unionWith (+) robots resources'
          ]

day19a :: Text -> Either (ParseErrorBundle Text Void) Int
day19a input = sum . fmap (uncurry (*)) . parMap rseq (traceShowId . second (geodes 24)) <$>
    parse (parser @Int @Int @Void <* eof) "day19.txt" input

day19b :: Text -> Either (ParseErrorBundle Text Void) Int
day19b input = product . parMap rseq (snd . traceShowId . second (geodes 32)) . take 3 <$>
    parse (parser @Int @Int @Void <* eof) "day19.txt" input