Day16.hs

module Javran.AdventOfCode.Y2023.Day16 () where

import Control.Monad
import Data.Function
import qualified Data.Map.Strict as M
import qualified Data.Sequence as Seq
import qualified Data.Set as S
import Javran.AdventOfCode.GridSystem.RowThenCol.Uldr
import Javran.AdventOfCode.Prelude

data Day16 deriving (Generic)

{-
  General notes:

  Tracing individual beams might not be viable if there are too many splitters
  in the grid. However:

  - We are looking at a grid bounded by rectangle (i.e. space is limited)

  - It does not matter, throughout this problem,
    how many times beam passes a tile.

  So if we break the beam into small, directed segments
  (say a segment is one part of the beam covered in a tile).
  We can treat each pair of (segment, direction) as unique and do
  an exhaustive expansion (here we use BFS) until no more segments can be included,
  at which point we are done, and energized tiles are just those
  that we have covered in this search.

  To represent a directed segment of the beam, I find it easy to just think about:

  - where is this segment? (represent it by coordinate of the covering tile)
  - in which direction is the beam going? (represented by a direction).

  This becomes the `Beam` type below.

 -}

type Dims = (Int, Int) -- rows, cols

data Obj
  = MirF -- `/`
  | MirB -- `\`
  | SpV -- `|`
  | SpH -- `-`

parseGrid :: [String] -> (Dims, M.Map Coord Obj)
parseGrid raw = ((rows, cols), m)
  where
    m = M.fromAscList do
      (r, rs) <- zip [0 ..] raw
      (c, ch) <- zip [0 ..] rs
      v <- case ch of
        '.' -> []
        '/' -> [MirF]
        '\\' -> [MirB]
        '|' -> [SpV]
        '-' -> [SpH]
        _ -> errInvalid

      pure ((r, c), v)

    rows = length raw
    cols = length (head raw)

{-
  Compute directions of the outward beams
  given a beam pointing `inDir` into an object / empty space
 -}
outBeams :: Dir -> Maybe Obj -> [Dir]
outBeams inDir = \case
  Nothing ->
    -- empty space, out direction is the same as in
    [inDir]
  Just o -> case o of
    MirF ->
      -- `/`
      case inDir of
        R -> [U]
        U -> [R]
        D -> [L]
        L -> [D]
    MirB ->
      -- `\`
      case inDir of
        R -> [D]
        D -> [R]
        L -> [U]
        U -> [L]
    SpV ->
      -- `|`
      case inDir of
        L -> [U, D]
        R -> [U, D]
        _ -> [inDir]
    SpH ->
      -- `-`
      case inDir of
        U -> [L, R]
        D -> [L, R]
        _ -> [inDir]

{-
  Represents a directed beam that goes
  from center of that coordinate
  to the direction `Dir` is pointing to.
 -}
type Beam = (Coord, Dir)

{-
  Computes the entire beam with breadth first search.
 -}
bfs :: Dims -> M.Map Coord Obj -> S.Set Beam -> Seq.Seq Beam -> S.Set Beam
bfs (rows, cols) grid = fix \go discovered -> \case
  Seq.Empty -> discovered
  (pt, d) Seq.:<| q1 ->
    let nexts = do
          let pt' = applyDir d pt
          {-
            Note that starting point of the search is outside of grid,
            but we always make sure enqueued beams are in.
           -}
          guard $ isInside pt'
          outDir <- outBeams d (grid M.!? pt')
          let next = (pt', outDir)
          guard $ S.notMember next discovered
          pure next
        discovered' = foldr S.insert discovered nexts
     in go discovered' (q1 <> Seq.fromList nexts)
  where
    isInside = inRange ((0, 0), (rows - 1, cols - 1))

{-
  Note that We assume `start` is always a point right outside of the grid.
  So we need -1 on top to get the final result.
 -}
countEnergized :: Dims -> M.Map Coord Obj -> Beam -> Int
countEnergized dims m start =
  subtract 1
    . S.size
    . S.map fst
    $ bfs dims m (S.singleton start) (Seq.singleton start)

instance Solution Day16 where
  solutionRun _ SolutionContext {getInputS, answerShow} = do
    (dims@(rows, cols), m) <- parseGrid . lines <$> getInputS
    let
      solve = countEnergized dims m
      leftsAndRights = do
        r <- [0 .. rows - 1]
        [solve ((r, -1), R), solve ((r, cols), L)]
      upsAndDowns = do
        c <- [0 .. cols - 1]
        [solve ((-1, c), D), solve ((rows, c), U)]

    answerShow $ head leftsAndRights
    answerShow $ maximum (leftsAndRights <> upsAndDowns)