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

<https://adventofcode.com/2015/day/20>

Run a simulation of elves delivering presents which each elf taking a larger
step size than the previous.

-}
module Main where

import Advent.Format (format)
import Control.Monad.Loop (for, exec_)
import Control.Monad.Trans.Class ( MonadTrans(lift) )
import Data.Array.ST (readArray, writeArray, MArray(newArray), runSTUArray)
import Data.Array.Unboxed (UArray, assocs)

-- | >>> :main
-- 831600
-- 884520
main :: IO ()
IO ()
main =
 do target <- [format|2015 20 %u%n|]
    print (findHouse target (solve1 target))
    print (findHouse target (solve2 target))

-- | Return the house number with at least the given number of presents.
findHouse :: Int -> UArray Int Int -> Int
findHouse :: Int -> UArray Int Int -> Int
findHouse Int
target UArray Int Int
a = [Int] -> Int
forall a. HasCallStack => [a] -> a
head [Int
h | (Int
h,Int
t) <- UArray Int Int -> [(Int, Int)]
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> [(i, e)]
assocs UArray Int Int
a, Int
t Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
target]

solve1 :: Int -> UArray Int Int
solve1 :: Int -> UArray Int Int
solve1 Int
target = (forall s. ST s (STUArray s Int Int)) -> UArray Int Int
forall i e. (forall s. ST s (STUArray s i e)) -> UArray i e
runSTUArray
 do let top :: Int
top = Int
target Int -> Int -> Int
forall a. Integral a => a -> a -> a
`quot` Int
10
    a <- (Int, Int) -> Int -> ST s (STUArray s Int Int)
forall i. Ix i => (i, i) -> Int -> ST s (STUArray s i Int)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> e -> m (a i e)
newArray (Int
1,Int
top) Int
0
    a <$ exec_
     do elf   <- for   1 (<= top) (1+)
        house <- for elf (<= top) (elf+)
        lift do old <- readArray a house
                writeArray a house (old + elf*10)

solve2 :: Int -> UArray Int Int
solve2 :: Int -> UArray Int Int
solve2 Int
target = (forall s. ST s (STUArray s Int Int)) -> UArray Int Int
forall i e. (forall s. ST s (STUArray s i e)) -> UArray i e
runSTUArray
 do let top :: Int
top = Int
target Int -> Int -> Int
forall a. Integral a => a -> a -> a
`quot` Int
11
    a <- (Int, Int) -> Int -> ST s (STUArray s Int Int)
forall i. Ix i => (i, i) -> Int -> ST s (STUArray s i Int)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> e -> m (a i e)
newArray (Int
1,Int
top) Int
0
    a <$ exec_
     do elf   <- for   1 (<=top) (1+)
        house <- for elf (<= min top (elf*50)) (+elf)
        lift do old <- readArray a house
                writeArray a house (old + elf*11)