{-# Language BlockArguments, ImportQualifiedPost, QuasiQuotes, GeneralisedNewtypeDeriving #-}
{-|
Module      : Main
Description : Day 21 solution
Copyright   : (c) Eric Mertens, 2021
License     : ISC
Maintainer  : emertens@gmail.com

<https://adventofcode.com/2021/day/21>

-}
module Main (main) where

import Advent.Format (format)
import Advent.Memo (memo4)
import Control.Applicative (Alternative((<|>)))
import Control.Monad (replicateM)
import Control.Monad.Trans.Writer.CPS (runWriterT, writerT, WriterT)
import Data.Coerce (coerce)
import Data.Map.Strict qualified as Map
import Data.Monoid (Product(Product))

-- | >>> :main
-- 428736
-- 57328067654557
main :: IO ()
main :: IO ()
main =
 do (Int
p1,Int
p2) <- [format|21 Player 1 starting position: %u%nPlayer 2 starting position: %u%n|]
    Int -> IO ()
forall a. Show a => a -> IO ()
print (Int -> Int -> Int -> Int -> Int -> Int
part1 Int
0 Int
p1 Int
p2 Int
0 Int
0)
    Int -> IO ()
forall a. Show a => a -> IO ()
print ([Int] -> Int
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum ((Bool, Int) -> Int
forall a b. (a, b) -> b
snd ((Bool, Int) -> Int) -> [(Bool, Int)] -> [Int]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Paths Bool -> [(Bool, Int)]
forall a. Paths a -> [(a, Int)]
runPaths (Int -> Int -> Int -> Int -> Paths Bool
part2 Int
p1 Int
p2 Int
0 Int
0)))

-- | Compute the @die rolls * losing score@ once one player
-- wins with 1000 points.
part1 ::
  Int {- ^ turn counter -} ->
  Int {- ^ player 1 location -} ->
  Int {- ^ player 2 location -} ->
  Int {- ^ player 1 score    -} ->
  Int {- ^ player 2 score    -} ->
  Int {- ^ player 2 score * 3 * turns -}
part1 :: Int -> Int -> Int -> Int -> Int -> Int
part1 Int
turns Int
p1 Int
p2 Int
p1s Int
p2s
  | Int
p1s' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
1000 = Int
3 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
turns' Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
p2s
  | Bool
otherwise    = Int -> Int -> Int -> Int -> Int -> Int
part1 Int
turns' Int
p2 Int
p1' Int
p2s Int
p1s'
  where
    turns' :: Int
turns' = Int
turns Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
    p1' :: Int
p1'    = Int -> Int -> Int
wrap (Int
p1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
6 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
turns) Int
10
    p1s' :: Int
p1s'   = Int
p1s Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
p1'

-- | Compute the possible ways a the players can win while playing with
-- a 3-sided dice given some starting conditions.
part2 ::
  Int   {- ^ player 1 location -} ->
  Int   {- ^ player 2 location -} ->
  Int   {- ^ player 1 score    -} ->
  Int   {- ^ player 2 score    -} ->
  Paths Bool {- ^ player 1 won -}
part2 :: Int -> Int -> Int -> Int -> Paths Bool
part2 = (Int -> Int -> Int -> Int -> Paths Bool)
-> Int -> Int -> Int -> Int -> Paths Bool
forall a b c d e.
(HasTrie a, HasTrie b, HasTrie c, HasTrie d) =>
(a -> b -> c -> d -> e) -> a -> b -> c -> d -> e
memo4 \Int
p1 Int
p2 Int
p1s Int
p2s ->
  Paths Bool -> Paths Bool
forall a. Ord a => Paths a -> Paths a
gather
   do Int
move <- Paths Int
threeRolls 
      let p1' :: Int
p1' = Int -> Int -> Int
wrap (Int
p1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
move) Int
10
      let p1s' :: Int
p1s' = Int
p1s Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
p1'
      if Int
p1s' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
21
        then Bool -> Paths Bool
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True
        else Bool -> Bool
not (Bool -> Bool) -> Paths Bool -> Paths Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Int -> Int -> Int -> Paths Bool
part2 Int
p2 Int
p1' Int
p2s Int
p1s'

-- | Sum of 3d3.
threeRolls :: Paths Int
threeRolls :: Paths Int
threeRolls = Paths Int -> Paths Int
forall a. Ord a => Paths a -> Paths a
gather ([Int] -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Int] -> Int) -> Paths [Int] -> Paths Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Paths Int -> Paths [Int]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
3 (Int -> Paths Int
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
1 Paths Int -> Paths Int -> Paths Int
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Int -> Paths Int
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
2 Paths Int -> Paths Int -> Paths Int
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Int -> Paths Int
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
3))

-- * Counting Nondeterminism Computations

-- | Nondeterministic computation that can consolidate
-- paths returning the same value.
newtype Paths a = Paths (WriterT (Product Int) [] a)
  deriving ((forall a b. (a -> b) -> Paths a -> Paths b)
-> (forall a b. a -> Paths b -> Paths a) -> Functor Paths
forall a b. a -> Paths b -> Paths a
forall a b. (a -> b) -> Paths a -> Paths b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: forall a b. a -> Paths b -> Paths a
$c<$ :: forall a b. a -> Paths b -> Paths a
fmap :: forall a b. (a -> b) -> Paths a -> Paths b
$cfmap :: forall a b. (a -> b) -> Paths a -> Paths b
Functor, Functor Paths
Functor Paths
-> (forall a. a -> Paths a)
-> (forall a b. Paths (a -> b) -> Paths a -> Paths b)
-> (forall a b c. (a -> b -> c) -> Paths a -> Paths b -> Paths c)
-> (forall a b. Paths a -> Paths b -> Paths b)
-> (forall a b. Paths a -> Paths b -> Paths a)
-> Applicative Paths
forall a. a -> Paths a
forall a b. Paths a -> Paths b -> Paths a
forall a b. Paths a -> Paths b -> Paths b
forall a b. Paths (a -> b) -> Paths a -> Paths b
forall a b c. (a -> b -> c) -> Paths a -> Paths b -> Paths c
forall (f :: * -> *).
Functor f
-> (forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
<* :: forall a b. Paths a -> Paths b -> Paths a
$c<* :: forall a b. Paths a -> Paths b -> Paths a
*> :: forall a b. Paths a -> Paths b -> Paths b
$c*> :: forall a b. Paths a -> Paths b -> Paths b
liftA2 :: forall a b c. (a -> b -> c) -> Paths a -> Paths b -> Paths c
$cliftA2 :: forall a b c. (a -> b -> c) -> Paths a -> Paths b -> Paths c
<*> :: forall a b. Paths (a -> b) -> Paths a -> Paths b
$c<*> :: forall a b. Paths (a -> b) -> Paths a -> Paths b
pure :: forall a. a -> Paths a
$cpure :: forall a. a -> Paths a
Applicative, Applicative Paths
Applicative Paths
-> (forall a b. Paths a -> (a -> Paths b) -> Paths b)
-> (forall a b. Paths a -> Paths b -> Paths b)
-> (forall a. a -> Paths a)
-> Monad Paths
forall a. a -> Paths a
forall a b. Paths a -> Paths b -> Paths b
forall a b. Paths a -> (a -> Paths b) -> Paths b
forall (m :: * -> *).
Applicative m
-> (forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
return :: forall a. a -> Paths a
$creturn :: forall a. a -> Paths a
>> :: forall a b. Paths a -> Paths b -> Paths b
$c>> :: forall a b. Paths a -> Paths b -> Paths b
>>= :: forall a b. Paths a -> (a -> Paths b) -> Paths b
$c>>= :: forall a b. Paths a -> (a -> Paths b) -> Paths b
Monad, Applicative Paths
Applicative Paths
-> (forall a. Paths a)
-> (forall a. Paths a -> Paths a -> Paths a)
-> (forall a. Paths a -> Paths [a])
-> (forall a. Paths a -> Paths [a])
-> Alternative Paths
forall a. Paths a
forall a. Paths a -> Paths [a]
forall a. Paths a -> Paths a -> Paths a
forall (f :: * -> *).
Applicative f
-> (forall a. f a)
-> (forall a. f a -> f a -> f a)
-> (forall a. f a -> f [a])
-> (forall a. f a -> f [a])
-> Alternative f
many :: forall a. Paths a -> Paths [a]
$cmany :: forall a. Paths a -> Paths [a]
some :: forall a. Paths a -> Paths [a]
$csome :: forall a. Paths a -> Paths [a]
<|> :: forall a. Paths a -> Paths a -> Paths a
$c<|> :: forall a. Paths a -> Paths a -> Paths a
empty :: forall a. Paths a
$cempty :: forall a. Paths a
Alternative)

-- | Return all values and counts from all the paths.
runPaths :: Paths a -> [(a, Int)]
runPaths :: forall a. Paths a -> [(a, Int)]
runPaths (Paths WriterT (Product Int) [] a
m) = [(a, Product Int)] -> [(a, Int)]
coerce (WriterT (Product Int) [] a -> [(a, Product Int)]
forall w (m :: * -> *) a. Monoid w => WriterT w m a -> m (a, w)
runWriterT WriterT (Product Int) [] a
m)

-- | Combine the counts of equal outputs to reduce braching factor.
gather :: Ord a => Paths a -> Paths a
gather :: forall a. Ord a => Paths a -> Paths a
gather (Paths WriterT (Product Int) [] a
xs) =
  WriterT (Product Int) [] a -> Paths a
forall a. WriterT (Product Int) [] a -> Paths a
Paths ([(a, Product Int)] -> WriterT (Product Int) [] a
forall (m :: * -> *) w a.
(Functor m, Monoid w) =>
m (a, w) -> WriterT w m a
writerT (Map a (Product Int) -> [(a, Product Int)]
forall k a. Map k a -> [(k, a)]
Map.toList ((Product Int -> Product Int -> Product Int)
-> [(a, Product Int)] -> Map a (Product Int)
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
Map.fromListWith Product Int -> Product Int -> Product Int
forall a. Num a => a -> a -> a
(+) (WriterT (Product Int) [] a -> [(a, Product Int)]
forall w (m :: * -> *) a. Monoid w => WriterT w m a -> m (a, w)
runWriterT WriterT (Product Int) [] a
xs))))

-- * Modular arithmetic

-- | Wrap number between @1@ and an inclusive upper bound
wrap :: Int {- ^ value -} -> Int {- ^ bound -} -> Int
wrap :: Int -> Int -> Int
wrap Int
x Int
n = (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1