life

import Control.Applicative
import Control.Comonad
import Control.Comonad.Store
import Data.Finite
import Data.Foldable
import Data.Maybe
import qualified Data.Vector.Sized as Vec
import Data.Vector.Sized (Vector)
import GHC.TypeNats

type Index x y = (Finite x, Finite y)
type KnownSize x y = (KnownNat x, KnownNat y)

newtype Grid x y a = Grid { grid :: Vector x (Vector y a) }
  deriving Functor

instance (Show a, KnownSize x y) => Show (Grid x y a) where
  show = unlines . map unwords . map toList . toList . grid . fmap show

inc :: KnownNat n => Finite n -> Maybe (Finite n)
inc = strengthen . shift

dec :: KnownNat n => Finite n -> Maybe (Finite n)
dec = unshift . weaken

neighboringCells :: KnownSize x y => Index x y -> [Index x y]
neighboringCells (i,j) =
  catMaybes $ map (uncurry (liftA2 (,)))
    [ (dec i,  dec j),  (dec i, Just j), (dec i,  inc j)
    , (Just i, dec j),                   (Just i, inc j)
    , (inc i,  dec j),  (inc i, Just j), (inc i,  inc j)
    ]

everywhere :: KnownSize x y => Grid x y (Index x y)
everywhere = Grid $ Vec.generate $ \i -> Vec.generate $ \j -> (i,j)

type ComonadGrid x y w = (ComonadStore (Index x y) w, KnownSize x y)

neighbors :: ComonadGrid x y w => w a -> [a]
neighbors = experiment neighboringCells

retrieve :: ComonadGrid x y w => w a -> Grid x y a
retrieve = experiment (const everywhere)

data Cell = Alive | Dead
  deriving Eq

instance Show Cell where
  show Alive = "#"
  show Dead = "."

cell :: Bool -> Cell
cell True = Alive
cell False = Dead

alive :: Cell -> Bool
alive = (Alive ==)

population :: ComonadGrid x y w => w Cell -> Int
population = length . filter alive . neighbors

rule :: ComonadGrid x y w => w Cell -> Cell
rule w =
  case extract w of
    Alive -> cell (population w `elem` [2,3])
    Dead -> cell (population w == 3)

index :: Grid x y a -> Index x y -> a
index (Grid g) = uncurry (Vec.index . Vec.index g)

type GridStore x y = Store (Index x y)

gridStore :: Grid x y a -> GridStore x y a
gridStore g = store (index g) (error "nowhere")

underGrid ::
  KnownSize x y =>
  (GridStore x y a -> GridStore x y b) ->
  Grid x y a -> Grid x y b
underGrid f = retrieve . f . gridStore

generate :: KnownSize x y => (GridStore x y a -> a) -> Grid x y a -> [Grid x y a]
generate = iterate . underGrid . extend

generations :: KnownSize x y => Grid x y Cell -> [Grid x y Cell]
generations = generate rule

test =
  Grid $ Vec.fromTuple
    ( Vec.fromTuple (Dead, Dead, Alive)
    , Vec.fromTuple (Alive, Alive, Alive)
    , Vec.fromTuple (Dead, Dead, Dead)
    )