{-# LANGUAGE ApplicativeDo #-}

import Control.Monad (forM_)
import qualified Data.ByteString as B
import Data.Char (ord)
import Data.List (nub, permutations, transpose)
import Data.Map (Map)
import qualified Data.Map as M
import Data.Set (Set)
import qualified Data.Set as S
import Data.Word (Word8)
import System.Environment (getArgs)

main :: IO ()
main = do
  [r] <- map read <$> getArgs
  let size = length . snd . head . M.toList $ r
      cells = S.size . S.fromList $ concat
        [ concat (concat k) ++ concat v | (k, v) <- M.toList r]
      width = 256
      height = 256
      width' = div width size * size
      height' = div height size * size
  seed <- map (map (map (\c -> fromIntegral c `mod` cells)))
        . take (size - 1)
        . chunk height'
        . chunk width'
        . B.unpack
        <$> B.readFile "seed.bin"
  let iesss = zip [0 :: Integer ..] $ main' cells size seed r
  forM_ iesss $ \(i, ess) -> do
    let p6 = "P6\n" ++ show width' ++ " " ++ show height ++ "\n255\n"
    B.writeFile (show i ++ ".ppm") . B.pack $
      map (fromIntegral . ord) p6 ++
      (concatMap (palette !!) . concat) ess

chunk :: Int -> [a] -> [[a]]
chunk _ [] = []
chunk n xs = case splitAt n xs of (ys, zs) -> ys : chunk n zs

type Cell = Int

type Rule = Map [[[Cell]]] [[Cell]]

main' :: Cell -> Int -> [[[Cell]]] -> Rule -> [[[Cell]]]
main' cells size seed r = evolution r
  where

    evolution :: Rule -> [[[Cell]]]
    evolution r = map (last . snd) . ((0, seed) :) . takeWhile (/= (0, seed)) . tail . iterate step $ (0, seed)
      where
        step (o, s) = (mod (o + 1) size, take (size - 1) (new : s))
          where
            new = map unoffset
                . unoffset
                -- [w,h]
                . concat
                -- [w',s,h]
                . map (map concat)
                -- [w',s,h',s]
                . map transpose
                -- [w',h',s,s]
                . map (map (r M.!))
                -- [w',h',s-1,s,s]
                . map (transpose)
                -- [w',s-1,h',s,s]
                . transpose
                -- [s-1,w',h',s,s]
                . map (map transpose)
                -- [s-1,w',s,h',s]
                . map (chunk size)
                -- [s-1,w,h',s]
                . map (map (chunk size))
                -- [s-1,w,h]
                . map (offset . map offset)
                $ s
            offset xs = drop o xs ++ take o xs
            unoffset xs = drop o' xs ++ take o' xs
              where
                o' = length xs - o

palette :: [[Word8]]
palette =
  [ [255, 255, 255]
  , [0, 0, 0]
  , [255, 0, 0]
  , [0, 255, 0]
  , [0, 0, 255]
  , [128, 128, 128]
  , [255, 255, 0]
  , [0, 255, 255]
  , [255, 0, 255]
  ]