#! /usr/bin/env -S"ANSWER=42" nix-shell
#! nix-shell -p ghcid
#! nix-shell -p "haskellPackages.ghcWithPackages (p: with p; [pretty-simple vector])"
#! nix-shell -i "ghcid -c 'ghci' -T main"

{-# OPTIONS_GHC -Wall -Wincomplete-uni-patterns #-}
{-# OPTIONS_GHC -Wno-unused-top-binds -Wno-unused-imports -Wno-type-defaults #-}
{-# OPTIONS_GHC -Wno-unused-matches #-}
{-# LANGUAGE OverloadedStrings #-}

import Data.Maybe (catMaybes)
import Data.Function (fix)
import Data.Vector (Vector)
import Debug.Trace (trace,traceShow)
import Text.Pretty.Simple
import qualified Data.Vector as V

exampleData :: [String]
exampleData =
  [ "L.LL.LL.LL"
  , "LLLLLLL.LL"
  , "L.L.L..L.."
  , "LLLL.LL.LL"
  , "L.LL.LL.LL"
  , "L.LLLLL.LL"
  , "..L.L....."
  , "LLLLLLLLLL"
  , "L.LLLLLL.L"
  , "L.LLLLL.LL"
  ]

type Width = Int
type Pos   = Int

data State = StateEmpty | StateSeatFree | StateSeatOccupied
  deriving (Eq, Ord)

instance Show State
  where
    show StateEmpty         = "."
    show StateSeatFree      = "L"
    show StateSeatOccupied  = "#"

type SeatLayout = Vector State

drawSeatLayout :: Width -> SeatLayout -> String
drawSeatLayout w s = concat $ V.toList $ V.imap go s
  where
    go i s0 = show s0 ++ if (i `mod` w) == (w - 1) then "\n" else ""

parseInput :: String -> SeatLayout
parseInput = V.fromList . (map go)
  where
    go '.' = StateEmpty
    go 'L' = StateSeatFree
    go '#' = StateSeatOccupied
    go c   = trace (c:" is undefined") undefined

fStep :: Width -> SeatLayout -> SeatLayout
fStep w s = if f s == s then s else fStep w (f s)
  where
    f = step w

step :: Width -> SeatLayout -> SeatLayout
step w s = V.imap go s
  where
    -- If a seat is empty (L) and there are no occupied seats adjacent to it,
    -- the seat becomes occupied.
    go i StateSeatFree =
      if all (== StateSeatFree) $ filter (/= StateEmpty) $ neighbors i
      then StateSeatOccupied
      else StateSeatFree
    -- If a seat is occupied (#) and four or more seats adjacent to it are also
    -- occupied, the seat becomes empty.
    go i StateSeatOccupied =
      if ( (length . (filter (== StateSeatOccupied))) (neighbors i) ) >= 4
      then StateSeatFree
      else StateSeatOccupied
    -- Otherwise, the seat's state does not change.
    -- ie the rest remains as-is.
    go _ state = state

    -- Build a list of neighbors
    neighbors :: Pos -> [State]
    neighbors p = catMaybes [
        -- 1: east
        s V.!? toP (i+1,j)
        -- 2: west
      , s V.!? toP (i-1,j)
        -- 3: north
      , s V.!? toP (i,j-1)
        -- 4: south
      , s V.!? toP (i,j+1)
        -- 5: north east
      , s V.!? toP (i+1,j-1)
        -- 6: north west
      , s V.!? toP (i-1,j-1)
        -- 7: south east
      , s V.!? toP (i+1,j+1)
        -- 8: south west
      , s V.!? toP (i-1,j+1)
      ]
      where
        toP (i0,j0) | i0 < w && i0 >= 0 = j0 * w + i0
        toP (i0,j0) | otherwise         = -1
        (i,j) = (p `mod` w, p `div` w)

solvePart1 :: String -> Int
solvePart1 str =  length
                $ V.filter (== StateSeatOccupied)
                $ fStep w
                $ parseInput
                $ concat
                $ lines
                $ str
  where
    w = length $ head $ lines $ str

main :: IO ()
main = do
  putStrLn ":: Tests"
  putStrLn $ drawSeatLayout 10 $ fStep 10 $ parseInput $ concat exampleData
  putStrLn ":: Day 11 - Part 1"
  print $ solvePart1 $ unlines exampleData
  input <- readFile "day11/input"
  print $ 2329 -- print $ solvePart1 input
  putStrLn ":: Tests"

  putStrLn ":: Day 11 - Part 2"