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

<https://adventofcode.com/2022/day/19>

This solution uses a few optimizations to achieve lightning fast performance:

* Prune out any state that has the same number of bots at a given time with fewer resources
* Prune out any state that in the best case produces fewer geodes than the any state doing nothing
* Generate a new state for each purchase, not individual timesteps.
* Any bot bought is bought at the earliest possible time.

>>> :{
:main +
    "Blueprint 1: Each ore robot costs 4 ore. Each clay robot costs 2 ore. Each obsidian robot costs 3 ore and 14 clay. Each geode robot costs 2 ore and 7 obsidian.\n\
    \Blueprint 2: Each ore robot costs 2 ore. Each clay robot costs 3 ore. Each obsidian robot costs 3 ore and 8 clay. Each geode robot costs 3 ore and 12 obsidian.\n"
:}
33
3472

-}
module Main where

import Control.Parallel (par)
import Control.Parallel.Strategies (parMap, rseq)
import Data.List (foldl')
import Data.Map qualified as Map

import Advent (format)

type Blueprint = (Int, Int, Int, Int, Int, Int, Int)

-- |
-- >>> :main
-- 1306
-- 37604
main :: IO ()
IO ()
main = do
    [(Int, Int, Int, Int, Int, Int, Int)]
input <- [format|2022 19
        (Blueprint %u: Each ore robot costs %u ore.
         Each clay robot costs %u ore.
         Each obsidian robot costs %u ore and %u clay.
         Each geode robot costs %u ore and %u obsidian.%n)*|]
    let xs :: [Int]
xs = Strategy Int
-> ((Int, Int, Int, Int, Int, Int, Int) -> Int)
-> [(Int, Int, Int, Int, Int, Int, Int)]
-> [Int]
forall b a. Strategy b -> (a -> b) -> [a] -> [b]
parMap Strategy Int
forall a. Strategy a
rseq (\b :: (Int, Int, Int, Int, Int, Int, Int)
b@(Int
i,Int
_,Int
_,Int
_,Int
_,Int
_,Int
_) -> Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int -> (Int, Int, Int, Int, Int, Int, Int) -> Int
solve Int
24 (Int, Int, Int, Int, Int, Int, Int)
b) [(Int, Int, Int, Int, Int, Int, Int)]
input
    let ys :: [Int]
ys = Strategy Int
-> ((Int, Int, Int, Int, Int, Int, Int) -> Int)
-> [(Int, Int, Int, Int, Int, Int, Int)]
-> [Int]
forall b a. Strategy b -> (a -> b) -> [a] -> [b]
parMap Strategy Int
forall a. Strategy a
rseq (Int -> (Int, Int, Int, Int, Int, Int, Int) -> Int
solve Int
32) (Int
-> [(Int, Int, Int, Int, Int, Int, Int)]
-> [(Int, Int, Int, Int, Int, Int, Int)]
forall a. Int -> [a] -> [a]
take Int
3 [(Int, Int, Int, Int, Int, Int, Int)]
input)
    Int -> IO ()
forall a. Show a => a -> IO ()
print ([Int]
ys [Int] -> Int -> Int
forall a b. a -> b -> b
`par` [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Int]
xs)
    Int -> IO ()
forall a. Show a => a -> IO ()
print ([Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [Int]
ys)

solve :: Int -> Blueprint -> Int
solve :: Int -> (Int, Int, Int, Int, Int, Int, Int) -> Int
solve Int
t0 (Int, Int, Int, Int, Int, Int, Int)
blue = Map Int [State] -> Int
go (Int -> [State] -> Map Int [State]
forall k a. k -> a -> Map k a
Map.singleton Int
t0 [State
def])
  where
    go :: Map Int [State] -> Int
go Map Int [State]
q =
      case Map Int [State] -> Maybe ((Int, [State]), Map Int [State])
forall k a. Map k a -> Maybe ((k, a), Map k a)
Map.maxViewWithKey Map Int [State]
q of
        Maybe ((Int, [State]), Map Int [State])
Nothing -> Int
0
        Just ((Int
0,[State]
sts), Map Int [State]
_) -> [Int] -> Int
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum ((State -> Int) -> [State] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Res -> Int
geo (Res -> Int) -> (State -> Res) -> State -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. State -> Res
amts) [State]
sts)
        Just ((Int
t,[State]
sts), Map Int [State]
q') ->
          Map Int [State] -> Int
go (Map Int [State] -> Int) -> Map Int [State] -> Int
forall a b. (a -> b) -> a -> b
$ ([State] -> [State] -> [State])
-> Map Int [State] -> Map Int [State] -> Map Int [State]
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
Map.unionWith [State] -> [State] -> [State]
forall a. [a] -> [a] -> [a]
(++) Map Int [State]
q' (Map Int [State] -> Map Int [State])
-> Map Int [State] -> Map Int [State]
forall a b. (a -> b) -> a -> b
$
          ([State] -> [State] -> [State])
-> [(Int, [State])] -> Map Int [State]
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
Map.fromListWith [State] -> [State] -> [State]
forall a. [a] -> [a] -> [a]
(++)
          [ (Int
t',[State
v'])
          | (Res
k,[Res]
vs) <-
                Map Res [Res] -> [(Res, [Res])]
forall k a. Map k a -> [(k, a)]
Map.assocs (Map Res [Res] -> [(Res, [Res])])
-> Map Res [Res] -> [(Res, [Res])]
forall a b. (a -> b) -> a -> b
$
                ([Res] -> [Res] -> [Res]) -> [(Res, [Res])] -> Map Res [Res]
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
Map.fromListWith [Res] -> [Res] -> [Res]
forall a. [a] -> [a] -> [a]
(++)
                [ (State -> Res
bots State
st, [State -> Res
amts State
st])
                | let u :: Int
u = [Int] -> Int
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum ((State -> Int) -> [State] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> State -> Int
underapprox Int
t) [State]
sts)
                , State
st <- [State]
sts
                , Int -> State -> Int
overapprox Int
t State
st Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
u]
          , Res
v       <- [Res] -> [Res]
keepBest [Res]
vs
          , (Int
t',State
v') <- (Int, Int, Int, Int, Int, Int, Int)
-> Int -> Res -> Res -> [(Int, State)]
step (Int, Int, Int, Int, Int, Int, Int)
blue Int
t Res
k Res
v]

-- | amount of geodes we'd end with if we bought a geode bot every single timestep
overapprox :: Int -> State -> Int
overapprox :: Int -> State -> Int
overapprox Int
t State
st = Int -> State -> Int
underapprox Int
t State
st Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
t Int -> Int -> Int
forall a. Num a => a -> a -> a
* (Int
t Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2

-- | amount of geodes we'd end with if we didn't buy any more geode bots
underapprox :: Int -> State -> Int
underapprox :: Int -> State -> Int
underapprox Int
t State
st = Int
t Int -> Int -> Int
forall a. Num a => a -> a -> a
* State
st.bots.geo Int -> Int -> Int
forall a. Num a => a -> a -> a
+ State
st.amts.geo

-- | Remove all resource sets from the list that are dominated by another
-- entry in the list.
keepBest :: [Res] -> [Res]
keepBest :: [Res] -> [Res]
keepBest = ([Res] -> Res -> [Res]) -> [Res] -> [Res] -> [Res]
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' [Res] -> Res -> [Res]
f []
  where
    f :: [Res] -> Res -> [Res]
f [Res]
acc Res
x
      | (Res -> Bool) -> [Res] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Res -> Res -> Bool
`cover` Res
x) [Res]
acc = [Res]
acc
      | Bool
otherwise = Res
x Res -> [Res] -> [Res]
forall a. a -> [a] -> [a]
: (Res -> Bool) -> [Res] -> [Res]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool) -> (Res -> Bool) -> Res -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Res -> Res -> Bool
cover Res
x) [Res]
acc

-- | Relation for the first element dominating the second.
cover :: Res -> Res -> Bool
cover :: Res -> Res -> Bool
cover Res
a Res
b =
    Res
a.ore Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Res
b.ore Bool -> Bool -> Bool
&&
    Res
a.cla Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Res
b.cla Bool -> Bool -> Bool
&&
    Res
a.obs Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Res
b.obs Bool -> Bool -> Bool
&&
    Res
a.geo Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Res
b.geo

data State = State {
    State -> Res
bots :: !Res,
    State -> Res
amts :: !Res
} deriving (Int -> State -> ShowS
[State] -> ShowS
State -> String
(Int -> State -> ShowS)
-> (State -> String) -> ([State] -> ShowS) -> Show State
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> State -> ShowS
showsPrec :: Int -> State -> ShowS
$cshow :: State -> String
show :: State -> String
$cshowList :: [State] -> ShowS
showList :: [State] -> ShowS
Show, State -> State -> Bool
(State -> State -> Bool) -> (State -> State -> Bool) -> Eq State
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: State -> State -> Bool
== :: State -> State -> Bool
$c/= :: State -> State -> Bool
/= :: State -> State -> Bool
Eq, Eq State
Eq State =>
(State -> State -> Ordering)
-> (State -> State -> Bool)
-> (State -> State -> Bool)
-> (State -> State -> Bool)
-> (State -> State -> Bool)
-> (State -> State -> State)
-> (State -> State -> State)
-> Ord State
State -> State -> Bool
State -> State -> Ordering
State -> State -> State
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: State -> State -> Ordering
compare :: State -> State -> Ordering
$c< :: State -> State -> Bool
< :: State -> State -> Bool
$c<= :: State -> State -> Bool
<= :: State -> State -> Bool
$c> :: State -> State -> Bool
> :: State -> State -> Bool
$c>= :: State -> State -> Bool
>= :: State -> State -> Bool
$cmax :: State -> State -> State
max :: State -> State -> State
$cmin :: State -> State -> State
min :: State -> State -> State
Ord)

data Res = Res {
    Res -> Int
ore, Res -> Int
cla, Res -> Int
obs, Res -> Int
geo :: !Int
} deriving (Int -> Res -> ShowS
[Res] -> ShowS
Res -> String
(Int -> Res -> ShowS)
-> (Res -> String) -> ([Res] -> ShowS) -> Show Res
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Res -> ShowS
showsPrec :: Int -> Res -> ShowS
$cshow :: Res -> String
show :: Res -> String
$cshowList :: [Res] -> ShowS
showList :: [Res] -> ShowS
Show, Res -> Res -> Bool
(Res -> Res -> Bool) -> (Res -> Res -> Bool) -> Eq Res
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Res -> Res -> Bool
== :: Res -> Res -> Bool
$c/= :: Res -> Res -> Bool
/= :: Res -> Res -> Bool
Eq, Eq Res
Eq Res =>
(Res -> Res -> Ordering)
-> (Res -> Res -> Bool)
-> (Res -> Res -> Bool)
-> (Res -> Res -> Bool)
-> (Res -> Res -> Bool)
-> (Res -> Res -> Res)
-> (Res -> Res -> Res)
-> Ord Res
Res -> Res -> Bool
Res -> Res -> Ordering
Res -> Res -> Res
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: Res -> Res -> Ordering
compare :: Res -> Res -> Ordering
$c< :: Res -> Res -> Bool
< :: Res -> Res -> Bool
$c<= :: Res -> Res -> Bool
<= :: Res -> Res -> Bool
$c> :: Res -> Res -> Bool
> :: Res -> Res -> Bool
$c>= :: Res -> Res -> Bool
>= :: Res -> Res -> Bool
$cmax :: Res -> Res -> Res
max :: Res -> Res -> Res
$cmin :: Res -> Res -> Res
min :: Res -> Res -> Res
Ord)

def :: State
def :: State
def = State {
    bots :: Res
bots = Res { ore :: Int
ore=Int
1, cla :: Int
cla=Int
0, obs :: Int
obs=Int
0, geo :: Int
geo=Int
0 },
    amts :: Res
amts = Res { ore :: Int
ore=Int
0, cla :: Int
cla=Int
0, obs :: Int
obs=Int
0, geo :: Int
geo=Int
0 }
}

divUp :: Int -> Int -> Int
divUp :: Int -> Int -> Int
divUp Int
x Int
y = (Int
xInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
yInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
y

step :: Blueprint -> Int -> Res -> Res -> [(Int,State)]
step :: (Int, Int, Int, Int, Int, Int, Int)
-> Int -> Res -> Res -> [(Int, State)]
step (Int
_, Int
oreCostOre, Int
claCostOre, Int
obsCostOre, Int
obsCostCla, Int
geoCostOre, Int
geoCostObs) Int
t !Res
bs !Res
as
  | [(Int, State)] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [(Int, State)]
buys = [(Int
0, State{ bots :: Res
bots = Res
bs, amts :: Res
amts = Res
as { geo = as.geo + t * bs.geo }}) | Int
t Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0]
  | Bool
otherwise = [(Int, State)]
buys
    where
        oreCostMax :: Int
oreCostMax = Int
oreCostOre Int -> Int -> Int
forall a. Ord a => a -> a -> a
`max` Int
claCostOre Int -> Int -> Int
forall a. Ord a => a -> a -> a
`max` Int
obsCostOre Int -> Int -> Int
forall a. Ord a => a -> a -> a
`max` Int
geoCostOre

        cap :: Res -> Res -> State
cap Res
b Res
a = State { bots :: Res
bots = Res
b , amts :: Res
amts = Res
a{
            ore = if b.ore == oreCostMax then min b.ore a.ore else a.ore,
            cla = if b.cla == obsCostCla then min b.cla a.cla else a.cla,
            obs = if b.obs == geoCostObs then min b.obs a.obs else a.obs}}

        buys :: [(Int, State)]
buys =
            [(Int
t', Res -> Res -> State
cap
                  Res
bs{ geo = bs.geo + 1 }
                  Res
as{ ore = as.ore + bs.ore * dt - geoCostOre
                    , cla = as.cla + bs.cla * dt
                    , obs = as.obs + bs.obs * dt - geoCostObs
                    , geo = as.geo + bs.geo * dt })
                | Res
bs.obs Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0
                , let dt :: Int
dt = Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+
                        Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max
                        ((Int
geoCostOre Int -> Int -> Int
forall a. Num a => a -> a -> a
- Res
as.ore) Int -> Int -> Int
`divUp` Res
bs.ore)
                        ((Int
geoCostObs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Res
as.obs) Int -> Int -> Int
`divUp` Res
bs.obs))
                , let t' :: Int
t' = Int
t Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
dt
                , Int
t' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0] [(Int, State)] -> [(Int, State)] -> [(Int, State)]
forall a. [a] -> [a] -> [a]
++
            [(Int
t', Res -> Res -> State
cap
                  Res
bs{ obs = bs.obs + 1 }
                  Res
as{ ore = as.ore + bs.ore * dt - obsCostOre
                    , cla = as.cla + bs.cla * dt - obsCostCla
                    , obs = as.obs + bs.obs * dt
                    , geo = as.geo + bs.geo * dt})
                | Res
bs.obs Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
geoCostObs
                , Res
bs.cla Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0
                , let dt :: Int
dt = Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max
                        ((Int
obsCostOre Int -> Int -> Int
forall a. Num a => a -> a -> a
- Res
as.ore) Int -> Int -> Int
`divUp` Res
bs.ore)
                        ((Int
obsCostCla Int -> Int -> Int
forall a. Num a => a -> a -> a
- Res
as.cla) Int -> Int -> Int
`divUp` Res
bs.cla))
                , let t' :: Int
t' = Int
t Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
dt
                , Int
t' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0] [(Int, State)] -> [(Int, State)] -> [(Int, State)]
forall a. [a] -> [a] -> [a]
++
            [(Int
t', Res -> Res -> State
cap
                  Res
bs{ cla = bs.cla + 1 }
                  Res
as{ ore = as.ore + bs.ore * dt - claCostOre
                    , cla = as.cla + bs.cla * dt
                    , obs = as.obs + bs.obs * dt
                    , geo = as.geo + bs.geo * dt})
                | Res
bs.cla Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
obsCostCla
                , let dt :: Int
dt = Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0
                        ((Int
claCostOre Int -> Int -> Int
forall a. Num a => a -> a -> a
- Res
as.ore) Int -> Int -> Int
`divUp` Res
bs.ore)
                , let t' :: Int
t' = Int
t Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
dt
                , Int
t' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0] [(Int, State)] -> [(Int, State)] -> [(Int, State)]
forall a. [a] -> [a] -> [a]
++
            [(Int
t', Res -> Res -> State
cap
                  Res
bs{ ore = bs.ore + 1 }
                  Res
as{ ore = as.ore + bs.ore * dt - oreCostOre
                    , cla = as.cla + bs.cla * dt
                    , obs = as.obs + bs.obs * dt
                    , geo = as.geo + bs.geo * dt})
                | Res
bs.ore Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
oreCostMax
                , let dt :: Int
dt = Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0
                        ((Int
oreCostOre Int -> Int -> Int
forall a. Num a => a -> a -> a
- Res
as.ore) Int -> Int -> Int
`divUp` Res
bs.ore)
                , let t' :: Int
t' = Int
t Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
dt
                , Int
t' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0]