import Data.Function (fix)
import Data.List (find)
import Data.Maybe (isJust)
import Data.Array
import Debug.Trace

import System.IO.Unsafe

type State = Array (Int, Int) Char

dirs :: [(Int, Int)]
dirs = [(x, y) | x <- [-1..1], y <- [-1..1], (x, y) /= (0, 0)]

traverseDirection :: State -> (Int, Int) -> (Int, Int) -> Maybe Char
traverseDirection state from (dx, dy) =
    find (const True) $ filter g $ map (state !) $ filter (inRange (bounds state)) lineCoords
    where
        g c = c /= '.'
        lineCoords = iterate (\(x, y) -> (x + dx, y + dy)) from
    
adjacentOccupied :: State -> (Int, Int) -> Int
adjacentOccupied state from =
    length . filter occupied $ map f dirs
    where
        occupied (Just '#') = True
        occupied _ = False
        f = traverseDirection state from

nextChar :: State -> (Int, Int) -> Char -> Char
nextChar state from c = 
    let n = adjacentOccupied state from in
    if c == 'L' && (traceShow ("from=", from, "c=", c, "n=", n) n) == 0 then '#' else
    if c == '#' && n >= 5 then 'L' else c

nextState :: State -> State
nextState state = state // map f (assocs (traceShow state state))
    where f (idx, ch) = (idx, nextChar state idx ch)

advanceState :: (State -> State) -> State -> State
advanceState rec state = let next = nextState state in
    if state == next then state
    else rec next

countOccupied :: State -> Int
countOccupied = length . filter (== '#') . elems

convertToState :: [String] -> State
convertToState lines =
    let lr = length lines in
    let lc = length (lines !! 0) in
    listArray ((0, 0), (lr - 1, lc - 1)) (concat lines)

solution2 :: [String] -> Int
solution2 = countOccupied . fix advanceState . convertToState

st = unsafePerformIO $ convertToState <$> lines <$> readFile "11ex.txt"

main :: IO Int
main = do
    input <- lines <$> readFile "11ex.txt"
    return (solution2 input)