Sandpile.hs

{-# LANGUAGE FlexibleInstances, OverloadedStrings #-}
module Main where

import Data.List
import qualified Data.Map as M
import Data.IORef
import System.Random
import System.Directory
import Control.Monad
import Control.Concurrent
import Codec.Picture

import qualified Network.Wai as Wai
import qualified Network.HTTP.Types as Wai
import qualified Network.Wai.Handler.Warp as Warp

main :: IO ()
main = do
  cgr <- newIORef M.empty
  createDirectoryIfMissing True "plots/lg"
  forkIO $ Warp.run 4024 $ \req respond -> do
    cg <- readIORef cgr
    respond $ Wai.responseLBS Wai.status200 [("Content-Type","image/png")] $
      let Right f = encodeDynamicPng (ImageRGB8 $ plot cg) in f
  let
    six :: M.Map L256 Integer
    six = M.fromList $ (\n -> (n,(totalNeighbours n - 1) * 2)) <$> allNodes
    sm = stabilizeM (const $ writeIORef cgr)
  identity <- loadOrCompute "plots/lg/0.png" $ do
    d3 <- sm six
    sm $ M.unionWith (+) six $ fmap negate d3
  let
    single = M.fromList $ flip (,) 1 <$> allNodes
    go n m = do
      let
        n' = n + 1
        mu = M.unionWith (+) m single
        [pfn,nfn] = ["plots/lg/" ++ show x ++ ".png" | x <- [n', -n']]
      pm <- loadOrCompute pfn $ sm mu
      loadOrCompute nfn $ do
        d <- sm $ M.unionWith (+) six pm
        sm $ M.unionWith (+) six $ negate <$> d
      go n' pm
  go 0 identity

class (Ord g) => Graph g where
  totalNeighbours :: g -> Integer
  concreteNeighbours :: g -> [g]

class Graph g => Grid g where
  gridBounds :: Foldable f => f g -> ((Int,Int),(Int,Int))
  gridPoint :: Int -> Int -> g

class Graph g => FiniteGrid g where
  allNodes :: [g]

instance (Integral a) => Graph (a,a) where
  totalNeighbours _ = 4
  concreteNeighbours (x,y) = do
    f <- [\n -> (x + n,y), \n -> (x,y + n)]
    d <- [-1,1]
    return (f d)

data BT128 = BT128 !Int !Int deriving (Eq,Ord,Show)

instance Graph BT128 where
  totalNeighbours = const 4
  concreteNeighbours (BT128 x y) = let
    go a b = if b == 127 || b == 0
      then filter ((&&) <$> (>= 0) <*> (<= 127)) [a + 1, a - 1]
      else map (`mod` 128) [a + 1, a - 1]
    in BT128 <$> go x y <*> go y x

instance Grid BT128 where
  gridBounds = const ((0,0),(127,127))
  gridPoint = BT128

instance FiniteGrid BT128 where
  allNodes = BT128 <$> [0 .. 127] <*> [0 .. 127]

data L256 = L256 !Int !Int deriving (Eq,Ord,Show)

instance Graph L256 where
  totalNeighbours = const 4
  concreteNeighbours (L256 x y) = do
    (p,r) <- [
      (x + 256 * y, \p' -> let (y',x') = p' `divMod` 256 in L256 x' y'),
      (y + 256 * (255 - x), \p' -> let (x',y') = p' `divMod` 256 in L256 (255 - x') y')
     ]
    p' <- [p + 1, p - 1]
    if p' >= 65536 || p' < 0 then [] else [()]
    return (r p')

instance Grid L256 where
  gridBounds = const ((0,0),(255,255))
  gridPoint = L256

instance FiniteGrid L256 where
  allNodes = L256 <$> [0 .. 255] <*> [0 .. 255]

loadOrCompute :: Grid g => FilePath -> IO (M.Map g Integer) -> IO (M.Map g Integer)
loadOrCompute fn fb = let
  fallback =  do
    r <- fb
    savePlot r fn
    return r
  in doesFileExist fn >>= \ae -> if ae
    then readPng fn >>= \rr -> case rr of
      Right (ImageRGB8 img) -> return $ M.fromList $ do
        x <- [0 .. imageWidth img - 1]
        y <- [0 .. imageHeight img - 1]
        let
          node = gridPoint x y
          colour = pixelAt img x y
          Just n = find (\n' -> chooseColour n' == colour) [0 ..]
        return (node,n)
      _ -> fallback
    else fallback

-- As defined on Wikiepedia
step :: Graph g => M.Map g Integer -> M.Map g Integer
step m = M.fromListWith (+) $ do
  o@(p,h) <- M.toList m
  let t = totalNeighbours p
  if h < t
    then return o
    else let
      a = case h - t of
        0 -> id
        a' -> ((p,a') :)
      r = flip (,) 1 <$> concreteNeighbours p
      in a r

-- Likely converges faster than `step`; definitely more interesting to watch.
-- Converges to the same state due to Abelian properties of the sandpile model.
stepv2 :: Graph g => M.Map g Integer -> M.Map g Integer
stepv2 m = M.filter (/= 0) $ M.fromListWith (+) $ do
  o@(p,h) <- M.toList m
  if h > 0
    then do
      let (s,h') = h `divMod` (totalNeighbours p)
      (p,h') : (flip (,) s <$> concreteNeighbours p)
    else return o

stabilize :: Graph g => M.Map g Integer -> M.Map g Integer
stabilize m
  | M.foldrWithKey (\k h r -> if h >= totalNeighbours k then False else r) True m = stabilize (stepv2 m)
  | otherwise = m

stabilizeM :: (Graph g, Monad m) => (Integer -> M.Map g Integer -> m ()) -> M.Map g Integer -> m (M.Map g Integer)
stabilizeM df = go 1 where
  go n m
    | M.foldrWithKey (\k h r -> if h >= totalNeighbours k then False else r) True m = df n m >> return m
    | otherwise = df n m >> go (n + 1) (stepv2 m)

run :: Graph g => M.Map g Integer -> [M.Map g Integer]
run = go where
  go m = let
    n = step m
    in m : if n == m then [] else go n

runv2 :: Graph g => M.Map g Integer -> [M.Map g Integer]
runv2 = go where
  go m
    | M.foldr (\h r -> if h >= 4 then True else r) False m = m : go (stepv2 m)
    | otherwise = [m]

instance (Integral a) => Grid (a,a) where
  gridBounds = foldl' (\((x0,y0),(xn,yn)) (x,y) -> let
    x0' = min (fromIntegral x) x0
    y0' = min (fromIntegral y) y0
    xn' = max (fromIntegral x) xn
    yn' = max (fromIntegral y) yn
    in x0' `seq` y0' `seq` xn' `seq` yn' `seq` ((x0',y0'),(xn',yn'))
   ) ((0,0),(0,0))
  gridPoint x y = (fromIntegral x, fromIntegral y)

mBounds :: Grid g => M.Map g Integer -> ((Int,Int),(Int,Int))
mBounds = gridBounds . M.keys

plot :: Grid g => M.Map g Integer -> Image PixelRGB8
plot m = let
  ((x0,y0),(xn,yn)) = mBounds m
  cf x y = gridPoint (x + x0) (y + y0)
  in generateImage (\x y -> chooseColour $ fromIntegral $ case M.lookup (cf x y) m of
    Nothing -> 0
    Just v -> v
   ) (xn - x0 + 1) (yn - y0 + 1)

savePlot :: Grid g => M.Map g Integer -> FilePath -> IO ()
savePlot m fn = savePngImage fn (ImageRGB8 (plot m))

chooseColour n = PixelRGB8
  (fromIntegral $ n * 123 + 113)
  (fromIntegral $ n * 25 + 143)
  (fromIntegral $ n * 61 + 53)

testPalette = savePngImage "palette.png" $
  ImageRGB8 $
  generateImage (\x _ -> chooseColour $ x `div` 25)
  (25 * 7)
  25