module Rosetta.AbelianSandpileModel.ST 
    ( simulate
    , test
    , toPGM
    ) where

import Control.Monad.Reader (asks, MonadReader (..), ReaderT, runReaderT)
import Control.Monad.ST (runST, ST)
import Control.Monad.State (evalStateT, forM_, lift, MonadState (..), StateT, modify, when)
import Data.Array.ST (freeze, readArray, STUArray, thaw, writeArray)
import Data.Array.Unboxed (array, assocs, bounds, UArray, (!))
import Data.Word (Word32)
import System.IO (hPutStr, hPutStrLn, IOMode (WriteMode), withFile)
import Text.Printf (printf)

type Point     = (Int, Int)
type ArrayST s = STUArray s Point Word32
type ArrayU    = UArray Point Word32

newtype M s a = M (ReaderT (S s) (StateT [Point] (ST s)) a)
    deriving (Functor, Applicative, Monad, MonadReader (S s), MonadState [Point])

data S s = S 
    { bMin :: !Point
    , bMax :: !Point
    , arr  :: !(ArrayST s)
    }

runM :: M s a -> S s -> [Point]-> ST s a
runM (M m) = evalStateT . runReaderT m

liftST :: ST s a -> M s a
liftST = M . lift . lift

simulate :: ArrayU -> ArrayU
simulate a = runST $ simulateST a

simulateST :: forall s. ArrayU -> ST s ArrayU
simulateST a = do
    let (p1, p2) = bounds a
        s = [p | (p, c) <- assocs a, c >= 4]
    b <- thaw a :: ST s (ArrayST s)
    let st = S { bMin = p1
               , bMax = p2
               , arr  = b
               }
    runM simulateM st s

simulateM :: forall s. M s ArrayU
simulateM = do
    ps <- get
    case ps of
        []      -> asks arr >>= liftST . freeze
        p: ps' -> do
            c <- changeArr p $ \x -> x - 4
            when (c < 4) $ put ps'
            forM_ [north, east, south, west] $ inc . ($ p)
            simulateM

changeArr :: Point -> (Word32 -> Word32) -> M s Word32
changeArr p f = do
    a    <- asks arr
    oldC <- liftST $ readArray a p
    let newC = f oldC
    liftST $ writeArray a p newC
    return newC

inc :: Point -> M s ()
inc p = do
    b <- inBounds p
    when b $ do
        c <- changeArr p succ
        when (c == 4) $ modify $ (p:)

inBounds :: Point -> M s Bool
inBounds p = do
    st <- ask
    return $ p >= bMin st && p <= bMax st

north, east, south, west :: Point -> Point
north (x, y) = (x, y + 1)
east  (x, y) = (x + 1, y)
south (x, y) = (x, y - 1)
west  (x, y) = (x - 1, y)

toPGM :: ArrayU -> FilePath -> IO ()
toPGM a fp = withFile fp WriteMode $ \h -> do
    let ((x1, y1), (x2, y2)) = bounds a
        width  = x2 - x1 + 1
        height = y2 - y1 + 1
    hPutStrLn h "P2"
    hPutStrLn h $ show width ++ " " ++ show height
    hPutStrLn h "3"
    forM_ [y1 .. y2] $ \y -> do
        forM_ [x1 .. x2] $ \x -> do
            let c = min 3 $ a ! (x, y)
            hPutStr h $ show c ++ " "
        hPutStrLn h ""

initArray :: Int -> Word32 -> ArrayU
initArray size height = array 
    ((-size, -size), (size, size))
    [((x, y), if x == 0 && y == 0 then height else 0) | x <- [-size .. size], y <- [-size .. size]]

test :: Int -> Word32 -> IO ()
test size height = do
    printf "size =%d, height =%d\n" size height
    let a  = initArray size height
        b  = simulate a
        fp = printf "sandpile_%d_%d.pgm" size height
    toPGM b fp
    putStrLn $ "wrote image to " ++ fp