{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE ViewPatterns #-}
module Main where

import qualified Data.Array.Unboxed as A
import Data.List (delete)
import qualified Data.Ix as Ix
import Data.Ix (Ix)
import Data.Maybe (catMaybes, fromMaybe)
import Data.Word (Word8)

import Input


generate :: Ix i => (i, i) -> (i -> e) -> A.Array i e
generate r f = A.array r [(i, f i) | i <- Ix.range r]

(!?) :: Ix i => A.Array i e -> i -> Maybe e
arr !? i | Ix.inRange (A.bounds arr) i = Just (arr A.! i)
         | otherwise = Nothing

class Ix i => Coordinate i where
    expand1 :: (i, i) -> (i, i)
    neighbourhood :: i -> [i]

instance Coordinate () where
    expand1 = id
    neighbourhood = pure

instance Coordinate i => Coordinate (i, Int) where
    expand1 ((idx1, i1), (idx2, i2)) =
        let (idx1', idx2') = expand1 (idx1, idx2)
        in ((idx1', i1 - 1), (idx2', i2 + 1))

    neighbourhood (idx, i) = [(idx', i + di) | idx' <- neighbourhood idx, di <- [-1..1]]


newtype Board i = Board (A.Array i Word8)
  deriving (Show)

fromInput :: [String] -> Board (((), Int), Int)
fromInput lns =
    let w = length (head lns)
        h = length lns
    in Board (A.array ((((), 0), 0), (((), w - 1), h - 1))
                      [((((), x), y), b2w (c == '#')) | (y, row) <- zip [0..] lns, (x, c) <- zip [0..] row])
  where b2w = fromIntegral . fromEnum :: Bool -> Word8

backpermute :: (Ix i, Ix j) => (j, j) -> (j -> i) -> Board i -> Board j
backpermute r f (Board arr) = Board (generate r (\j -> arr A.! f j))

addDimension :: Coordinate i => Board i -> Board (i, Int)
addDimension (Board arr@(A.bounds -> (i, j))) = backpermute ((i, 0), (j, 0)) fst (Board arr)

evolve :: Coordinate i => Board i -> Board i
evolve (Board prev) = Board vec
  where vec = generate (expand1 (A.bounds prev))
                       (\index ->
                           let cell = fromMaybe 0 (prev !? index)
                               nb = sum (catMaybes (map (prev !?) (delete index (neighbourhood index))))
                           in if (cell /= 0 && nb == 2) || nb == 3 then 1 else 0)

numActive :: Ix i => Board i -> Int
numActive (Board arr) = sum (map fromIntegral (A.elems arr))

main :: IO ()
main = do
    bd2 <- fromInput <$> getInput 17
    let bd3 = addDimension bd2
        bd4 = addDimension bd3
    print (numActive (iterate evolve bd3 !! 6))
    print (numActive (iterate evolve bd4 !! 6))