Advent of Code 2022 - Day 19

Solution.hs

{-# LANGUAGE BangPatterns #-}
-- Description : Advent of Code - 2022 / Tag 19
--
-- see [Advent of Code 2022 - day 19](https://adventofcode.com/2022/day/19)
{-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-}
{-# OPTIONS_GHC -Wno-name-shadowing #-}

module Y2022.Day19.Solution where

import CommonParsers (Parser, numberP, runParser)
import Data.Tuple (swap)
import qualified Text.Megaparsec as P
import qualified Text.Megaparsec.Char as PC

yearNr :: Int
yearNr = 2022

dayNr :: Int
dayNr = 19

run :: IO ()
run = do
  putStrLn $ "YEAR " <> show yearNr <> "/ DAY " <> show dayNr

  input <- loadInput

  let result1 = part1 input
  putStrLn $ "\t Part 1: " ++ show result1

  let result2 = part2 input
  putStrLn $ "\t Part 2: " ++ show result2

  putStrLn "---\n"

----------------------------------------------------------------------
-- solutions

part1 :: Input -> Int
part1 = qualityLevel 24

part2 :: Input -> Int
part2 = solve2

----------------------------------------------------------------------
-- data model

type Input = BluePrints

data Material
  = Ore
  | Clay
  | Obsidian
  | Geode
  deriving (Show, Eq, Ord)

type Cost = Int

data RobotBluePrint = Robot
  { robotMines :: Material
  , robotCosts :: [(Material, Cost)]
  }
  deriving (Show)

type BluePrint = [RobotBluePrint]

type BluePrints = [BluePrint]

data Robots = Robots
  { oreRobots :: Int
  , clayRobots :: Int
  , obsidianRobots :: Int
  , geodeRobots :: Int
  }
  deriving (Show)

data Materials = Materials
  { ores :: Int
  , clays :: Int
  , obsidians :: Int
  , geodes :: Int
  }
  deriving (Show)

----------------------------------------------------------------------
-- Algorithm

type Minutes = Int
type Count = Int

qualityLevel :: Minutes -> Input -> Count
qualityLevel totalTime =
  sum . zipWith (\i bp -> i * solveBluePrint bp totalTime) [1 ..]

solve2 :: Input -> Count
solve2 =
  product . map (`solveBluePrint` 32) . take 3

solveBluePrint :: BluePrint -> Minutes -> Count
solveBluePrint bluePrint =
  solve oRo cRo obRo obRc gRo gRO
 where
  oRo = snd $ head $ robotCosts oreRob
  cRo = snd $ head $ robotCosts clayRob
  [obRo, obRc] = map snd $ robotCosts obsidianRob
  [gRo, gRO] = map snd $ robotCosts geodeRob
  (oreRob, clayRob, obsidianRob, geodeRob) =
    case bluePrint of
      [oreRob, clayRob, obsidianRob, geodeRob] -> (oreRob, clayRob, obsidianRob, geodeRob)
      _ -> error "wrong config"

solve :: Int -> Int -> Int -> Int -> Int -> Int -> Minutes -> Count
solve oreRobotOreCost clayRobotOreCost obsidianRobotOreCost obsidianRobotClayCost geodeRobotOreCost geodeRobotObsidianCost =
  go 0 (1, 0, 0, 0) (0, 0, 0, 0) (False, False, False)
 where
  go !bestSoFar (_, _, _, _) (_, _, _, !geodes) _ 0 = max bestSoFar geodes
  go !bestSoFar (!oreRobots, !clayRobots, !obsidianRobots, !geodeRobots) (!ore, !clay, !obsidian, !geodes) (couldProduceOreRobot, couldProduceClayRobot, couldProduceObsidianRobot) timeLeft
    | (timeLeft * (timeLeft - 1) `div` 2) + geodes + geodeRobots * timeLeft <= bestSoFar = bestSoFar
    | obsidian >= geodeRobotObsidianCost && ore >= geodeRobotOreCost =
        go
          bestSoFar
          (oreRobots, clayRobots, obsidianRobots, geodeRobots + 1)
          (ore + oreRobots - geodeRobotOreCost, clay + clayRobots, obsidian + obsidianRobots - geodeRobotObsidianCost, geodes + geodeRobots)
          (False, False, False)
          (timeLeft - 1)
    | otherwise =
        accum bestSoFar (withNewOreRobot ++ withNewClayRobot ++ withNewObsidianRobot ++ justProduce)
   where
    accum best [] = best
    accum best (opt : more) =
      let best' = opt best
       in accum best' more
    justProduce =
      [ \best ->
          go
            best
            (oreRobots, clayRobots, obsidianRobots, geodeRobots)
            (ore + oreRobots, clay + clayRobots, obsidian + obsidianRobots, geodes + geodeRobots)
            (canProduceOreRobot, canProduceClayRobot, canProduceObidianRobot)
            (timeLeft - 1)
      ]
    canProduceObidianRobot = clay >= obsidianRobotClayCost && ore >= obsidianRobotOreCost
    withNewObsidianRobot
      | not couldProduceObsidianRobot && obsidianRobots < maxObsidianCost && canProduceObidianRobot =
          [ \best ->
              go
                best
                (oreRobots, clayRobots, obsidianRobots + 1, geodeRobots)
                (ore + oreRobots - obsidianRobotOreCost, clay + clayRobots - obsidianRobotClayCost, obsidian + obsidianRobots, geodes + geodeRobots)
                (False, False, False)
                (timeLeft - 1)
          ]
      | otherwise = []
    canProduceClayRobot = ore >= clayRobotOreCost
    withNewClayRobot
      | not couldProduceClayRobot && clayRobots < maxClayCost && canProduceClayRobot =
          [ \best ->
              go
                best
                (oreRobots, clayRobots + 1, obsidianRobots, geodeRobots)
                (ore + oreRobots - clayRobotOreCost, clay + clayRobots, obsidian + obsidianRobots, geodes + geodeRobots)
                (False, False, False)
                (timeLeft - 1)
          ]
      | otherwise = []
    canProduceOreRobot = ore >= oreRobotOreCost
    withNewOreRobot
      | not couldProduceOreRobot && oreRobots < maxOreCost && canProduceOreRobot =
          [ \best ->
              go
                best
                (oreRobots + 1, clayRobots, obsidianRobots, geodeRobots)
                (ore + oreRobots - oreRobotOreCost, clay + clayRobots, obsidian + obsidianRobots, geodes + geodeRobots)
                (False, False, False)
                (timeLeft - 1)
          ]
      | otherwise = []
  maxOreCost = maximum [oreRobotOreCost, clayRobotOreCost, obsidianRobotOreCost, geodeRobotOreCost]
  maxClayCost = obsidianRobotClayCost
  maxObsidianCost = geodeRobotObsidianCost

----------------------------------------------------------------------
-- load and parse input

loadInput :: IO Input
loadInput = loadFile "input.txt"

loadExample :: IO Input
loadExample = loadFile "example.txt"

loadFile :: FilePath -> IO Input
loadFile file = do
  txt <- readFile ("./src/Y" <> show yearNr <> "/Day" <> show dayNr <> "/" <> file)
  pure $ parseText txt

parseText :: String -> Input
parseText = map (runParser bluePrintP) . lines

bluePrintP :: Parser BluePrint
bluePrintP = do
  _ <- PC.string "Blueprint " *> (numberP :: Parser Int) <* PC.string ": "
  robotBluePrintP `P.sepBy1` PC.string " "

robotBluePrintP :: Parser RobotBluePrint
robotBluePrintP = do
  _ <- PC.string "Each "
  mat <- materialP
  _ <- PC.string " robot costs "
  costs <- costsP
  _ <- PC.string "."
  pure $ Robot mat costs

costsP :: Parser [(Material, Cost)]
costsP = (swap <$> costP) `P.sepBy1` PC.string " and "

costP :: Parser (Cost, Material)
costP = (,) <$> (numberP <* PC.string " ") <*> materialP

materialP :: Parser Material
materialP = P.choice [oreP, clayP, obisidanP, geodeP]
 where
  oreP = Ore <$ PC.string "ore"
  clayP = Clay <$ PC.string "clay"
  obisidanP = Obsidian <$ PC.string "obsidian"
  geodeP = Geode <$ PC.string "geode"