{-# LANGUAGE LambdaCase #-}
module Main where

import Data.List (find, findIndex, intercalate)
import Data.Maybe (catMaybes, fromJust)
import qualified Data.Vector as V
import qualified Data.Vector.Storable as U
import Data.Word (Word8)
import Foreign.Ptr (Ptr, castPtr)
import Foreign.Storable (Storable(..))

import Input


data Cell = Floor | Seat | Occupied
  deriving (Show, Eq, Enum)

instance Storable Cell where
    sizeOf _ = sizeOf (undefined :: Word8)
    alignment _ = alignment (undefined :: Word8)
    peek p = toEnum . fromIntegral <$> peek (castPtr p :: Ptr Word8)
    poke p x = poke (castPtr p :: Ptr Word8) (fromIntegral (fromEnum x))

parseCell :: Char -> Cell
parseCell '.' = Floor
parseCell 'L' = Seat
parseCell '#' = Occupied
parseCell _ = error "Cannot parse cell"

showCell :: Cell -> Char
showCell Floor = '.'
showCell Seat = 'L'
showCell Occupied = '#'

data Grid = Grid Int Int (U.Vector Cell)
  deriving (Show, Eq)

parseGrid :: [String] -> Grid
parseGrid lns =
    let w = length (head lns)
        h = length lns
    in Grid w h (U.fromListN (w * h) (concatMap (map parseCell) lns))

showGrid :: Grid -> String
showGrid (Grid w h v) =
    intercalate "\n" [[showCell (v U.! (w * y + x)) | x <- [0..w-1]]
                     | y <- [0..h-1]]

numOccupied :: Grid -> Int
numOccupied (Grid _ _ v) = sum (map (fromEnum . (== Occupied)) (U.toList v))

evolve :: Grid -> ([Cell] -> Maybe Int) -> Int -> Grid -> Grid
evolve (Grid w h template) nbFunc scareCount =
    \(Grid _ _ v) ->
        let newCellAt x y =
                let neighbours = [v U.! (w * y' + x') | (x', y') <- nbIndicesArr V.! (w * y + x)]
                    filledNbs = sum (map (fromEnum . (== Occupied)) neighbours)
                in case v U.! (w * y + x) of
                     Floor -> Floor
                     Seat | filledNbs == 0 -> Occupied
                          | otherwise -> Seat
                     Occupied | filledNbs >= scareCount -> Seat
                              | otherwise -> Occupied
        in Grid w h (U.generate (w * h) (\i -> newCellAt (i `rem` w) (i `quot` w)))
  where
    inBounds (x, y) = 0 <= x && x < w && 0 <= y && y < h
    nbIndices x y = catMaybes
                      [(sight !!) <$> nbFunc [template U.! (w * y' + x') | (x', y') <- sight]
                      | dy <- [-1..1], dx <- [-1..1]
                      , dx /= 0 || dy /= 0
                      , let sight = takeWhile inBounds [(x+i*dx, y+i*dy) | i <- [1..]]]
    nbIndicesArr = V.generate (w * h) (\i -> nbIndices (i `rem` w) (i `quot` w))


equilibriumFor :: Eq a => (a -> a) -> a -> a
equilibriumFor f x = fst (fromJust (find (uncurry (==)) (zip xs (tail xs))))
  where xs = iterate f x

main :: IO ()
main = do
    input <- getInput 11
    let initGrid = parseGrid input
        evolve1 = evolve initGrid (\case Seat:_ -> Just 0 ; _ -> Nothing) 4
        evolve2 = evolve initGrid (findIndex (== Seat)) 5
    print (numOccupied (equilibriumFor evolve1 initGrid))
    print (numOccupied (equilibriumFor evolve2 initGrid))