{-# 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
  [size, cells] <- map read <$> getArgs
  let width = 256
      height = 256
      width' = div width size * size
  seed <- map (map (\c -> fromIntegral c `mod` cells))
        . take (size - 1)
        . chunk width'
        . B.unpack
        <$> B.readFile "seed.bin"
  let rs = zip [0 :: Integer ..] $ main' cells size seed
  forM_ rs $ \(i, (r, es)) -> do
    let stem = show size ++ "-" ++ show cells ++ "-" ++ show i
        p6 = "P6\n" ++ show width' ++ " " ++ show height ++ "\n255\n"
    putStrLn $ "<a href='" ++ stem ++ ".txt'><img src='" ++
        stem ++ ".png' alt='" ++ show i ++ "'></a>"
    writeFile (stem ++ ".txt") (show r)
    B.writeFile (stem ++ ".ppm") . B.pack $
      map (fromIntegral . ord) p6 ++
      (concatMap (palette !!) . concat . take height) es

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

type Cell = Int

type Block = [[Cell]]

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

type RuleSet = Set Block

type Symmetry = Block -> Block

main' :: Cell -> Int -> [[Cell]] -> [(Rule, [[Cell]])]
main' cells size seed
    = [(r, evolution r) | r <- enumerateSymmetric ]
  where

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

    allCells :: [Cell]
    allCells = [0 .. cells - 1]

    allTo :: [[Cell]]
    allTo = replicateM' size allCells

    allFrom :: [Block]
    allFrom = replicateM' (size - 1) allTo

    rotate :: Cell -> Cell
    rotate c = (c + 1) `mod` cells

    consistentSet :: Set Block -> Bool
    consistentSet bs = S.size bs == S.size (S.map tail bs)

    symmetry :: [Symmetry]
    symmetry =
      [ id
      , reverse
      , map reverse
      , transpose
      , map (map rotate)
      ]

    closure :: Set Block -> [Set Block]
    closure = go S.empty
      where
        go acc bs
          | not (consistentSet bs) = []
          | not (consistentSet acc) = []
          | S.null bs = return acc
          | otherwise = go (S.union acc bs) new
          where
            more = S.fromList
                     [ s b | s <- tail symmetry, b <- S.toList bs ]
            new = S.difference more acc

    enumerateSymmetric :: [Rule]
    enumerateSymmetric =
        nub . map (rule . canonical) $
          go (S.fromList allFrom) S.empty
      where
        go fs acc = case S.minView fs of
          Nothing
            | consistentSet acc -> return acc
            | otherwise -> []
          Just (f, next)
            | consistentSet acc -> do
                t <- allTo
                new <- closure (S.singleton (t : f))
                let fs' = S.difference next (S.map tail new)
                    acc' = S.union acc new
                go fs' acc'
            | otherwise -> []

    canonical :: RuleSet -> RuleSet
    canonical rs = minimum . map (permute rs) $
        permutations [0..cells - 1]

    permute :: RuleSet -> [Cell] -> RuleSet
    permute rs ps = S.map (map (map (ps !!))) rs

    rule :: RuleSet -> Rule
    rule = M.fromList . map f . S.toList
      where
        f (r:rs) = (rs, r)

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]
  ]

replicateM' :: Applicative m => Int -> m a -> m [a]
replicateM' m ls = go m (pure id)
  where
    go n fs
      | n <= 0 = do
          f <- fs
          pure (f [])
      | otherwise = go (n - 1) gs
          where
            gs = do
              f <- fs
              l <- ls
              pure (f . (l:))