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

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

Apply light on/off/toggle commands to a grid.

-}
module Main where

import Advent (countBy, format, stageTH)
import Advent.Coord (Coord(C))
import Control.Monad.ST (ST, runST)
import Data.Array.ST
import Data.Foldable (for_, traverse_)

data C = Con | Coff | Ctoggle
data Command = Command !C Coord Coord

stageTH

-- | >>> :main
-- 377891
-- 14110788
main :: IO ()
IO ()
main =
 do input <- [format|2015 6 ((turn |)@C %d,%d through %d,%d%n)*|]
    let cmds = [C -> Coord -> Coord -> Command
Command C
c (Int -> Int -> Coord
C Int
y1 Int
x1) (Int -> Int -> Coord
C Int
y2 Int
x2) | (C
c,Int
x1,Int
y1,Int
x2,Int
y2) <- [(C, Int, Int, Int, Int)]
input]
    print (runST (part1 cmds))
    print (runST (part2 cmds))

part1 :: [Command] -> ST s Int
part1 :: forall s. [Command] -> ST s Int
part1 [Command]
cmds =
 do a <- ST s (STUArray s Coord Bool)
forall s. ST s (STUArray s Coord Bool)
newBitGrid
    traverse_ (bitCommand a) cmds
    xs <- getElems a
    return $! countBy id xs

part2 :: [Command] -> ST s Int
part2 :: forall s. [Command] -> ST s Int
part2 [Command]
cmds =
 do a <- ST s (STUArray s Coord Int)
forall s. ST s (STUArray s Coord Int)
newIntGrid
    traverse_ (intCommand a) cmds
    xs <- getElems a
    return $! sum xs

bitCommand :: STUArray s Coord Bool -> Command -> ST s ()
bitCommand :: forall s. STUArray s Coord Bool -> Command -> ST s ()
bitCommand STUArray s Coord Bool
a (Command C
op Coord
x Coord
y) =
  [Coord] -> (Coord -> ST s ()) -> ST s ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ ((Coord, Coord) -> [Coord]
forall a. Ix a => (a, a) -> [a]
range (Coord
x, Coord
y)) \Coord
p ->
    case C
op of
      C
Con     -> STUArray s Coord Bool -> Coord -> Bool -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray STUArray s Coord Bool
a Coord
p Bool
True
      C
Coff    -> STUArray s Coord Bool -> Coord -> Bool -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray STUArray s Coord Bool
a Coord
p Bool
False
      C
Ctoggle -> STUArray s Coord Bool -> Coord -> Bool -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray STUArray s Coord Bool
a Coord
p (Bool -> ST s ()) -> (Bool -> Bool) -> Bool -> ST s ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool -> Bool
not (Bool -> ST s ()) -> ST s Bool -> ST s ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< STUArray s Coord Bool -> Coord -> ST s Bool
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray STUArray s Coord Bool
a Coord
p

intCommand :: STUArray s Coord Int -> Command -> ST s ()
intCommand :: forall s. STUArray s Coord Int -> Command -> ST s ()
intCommand STUArray s Coord Int
a (Command C
op Coord
x Coord
y) =
  [Coord] -> (Coord -> ST s ()) -> ST s ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ ((Coord, Coord) -> [Coord]
forall a. Ix a => (a, a) -> [a]
range (Coord
x, Coord
y)) \Coord
p ->
    STUArray s Coord Int -> Coord -> Int -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray STUArray s Coord Int
a Coord
p (Int -> ST s ()) -> (Int -> Int) -> Int -> ST s ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int
upd (Int -> ST s ()) -> ST s Int -> ST s ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< STUArray s Coord Int -> Coord -> ST s Int
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray STUArray s Coord Int
a Coord
p
  where
  upd :: Int -> Int
upd = case C
op of
          C
Con     -> (Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
          C
Coff    -> Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 (Int -> Int) -> (Int -> Int) -> Int -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int -> Int
forall a. Num a => a -> a -> a
subtract Int
1
          C
Ctoggle -> (Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
2)

newBitGrid :: ST s (STUArray s Coord Bool)
newBitGrid :: forall s. ST s (STUArray s Coord Bool)
newBitGrid = (Coord, Coord) -> Bool -> ST s (STUArray s Coord Bool)
forall i. Ix i => (i, i) -> Bool -> ST s (STUArray s i Bool)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> e -> m (a i e)
newArray (Int -> Int -> Coord
C Int
0 Int
0, Int -> Int -> Coord
C Int
999 Int
999) Bool
False

newIntGrid :: ST s (STUArray s Coord Int)
newIntGrid :: forall s. ST s (STUArray s Coord Int)
newIntGrid = (Coord, Coord) -> Int -> ST s (STUArray s Coord 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 -> Int -> Coord
C Int
0 Int
0, Int -> Int -> Coord
C Int
999 Int
999) Int
0