Discussion on: Advent of Code 2020 Solution Megathread - Day 11: Seating System

module Main where

import qualified Data.Map.Strict as M
import Data.Maybe (mapMaybe)

type Coordinate = (Int, Int)

toCoordinateMap :: [[a]] -> M.Map (Int, Int) a
toCoordinateMap a = M.fromList $ do
  (y, row) <- zip [0 ..] a
  (x, v) <- zip [0 ..] row
  pure ((x, y), v)

findAdjSeats :: Coordinate -> M.Map (Int, Int) Char -> [Coordinate]
findAdjSeats (x, y) seats =
  filter (`M.member` seats)
    [ (x-1, y-1), (x, y-1), (x+1, y-1),
      (x-1, y),             (x+1, y),
      (x-1, y+1), (x, y+1), (x+1, y+1)
    ]

findFirstSeatInSight :: Coordinate -> M.Map Coordinate Char -> [Coordinate]
findFirstSeatInSight coord seatMap =
  mapMaybe (firstInSight coord)
    [ (-1, -1), (0, -1), (1, -1),
      (-1, 0),           (1, 0),
      (-1, 1),  (0, 1),  (1, 1)
    ]
  where
    firstInSight (x, y) (dx, dy) = (x-dx, y-dy) `M.lookup` seatMap >>= aux
      where
        aux seat =
          if seat == 'L' || seat == '#'
          then Just (x-dx, y-dy)
          else firstInSight (x-dx, y-dy) (dx, dy)

runSimulation :: M.Map Coordinate Char -> (Coordinate -> M.Map Coordinate Char -> [Coordinate]) -> Int -> Int
runSimulation seatMap findAdjSeatsF occupiedSeats = 
  if seatMap == nextCycle
  then length $ M.filter (== '#') nextCycle
  else runSimulation nextCycle findAdjSeatsF occupiedSeats
  where
    noOccupiedAdjSeats = not . any ((== Just '#') . (`M.lookup` seatMap))
    moreOrEqThanNOccupiedSeats n = (>= n) . length . filter ((== Just '#') . (`M.lookup` seatMap))

    nextCycle =
      M.mapWithKey (\coordinate c ->
        case c of
          '.' -> '.'
          'L' ->
            if noOccupiedAdjSeats $ findAdjSeatsF coordinate seatMap
            then '#'
            else 'L'

          '#' -> 
            if moreOrEqThanNOccupiedSeats occupiedSeats $ findAdjSeatsF coordinate seatMap
            then 'L'
            else '#'

          _ -> error $ "cannot do" ++ show c
      ) seatMap

main :: IO ()
main = do
  input <- lines <$> readFile "input.txt"
  let seatMap = toCoordinateMap input
  print $ runSimulation seatMap findAdjSeats 4
  print $ runSimulation seatMap findFirstSeatInSight 5