module MCMC.Actions ( Action (..)
                    -- * Predefined actions
                    , collect
                    , display
                    -- ** Batch actions
                    , Batch
                    , BatchAct
                    , BatchAction
                    , inBatches
                    , pack
                    , PrintF
                    , batchViz
                    , batchPrint
                    ) where

import Control.Monad
import MCMC.Types

collect :: Action x IO [x] [x]
collect = makeAction (\ x ls -> return (x:ls)) return []

display :: Show s => (x -> s) -> Action x IO () ()
display f = makeAction (\x _ -> print $ f x) print ()

-- Batch actions --  

type Batch x = ([x], Int)
type BatchAct x m = Act x m (Batch x)
type BatchAction x m b = Action x m (Batch x) b

inBatches :: Monad m => (Batch x -> m b) -> Int -> BatchAct x m
inBatches f n x a@(l, i)
    | i == n = f a >> return ([x], 1)
    | otherwise = return (x:l, i+1)

pack :: (Batch x -> m b) -> BatchAct x m -> BatchAction x m b
pack f act = makeAction act f ([], 0)

-- Batch Visualization --

type PrintF x s = [x] -> [s]

vizJSON :: Show s => [s] -> String
vizJSON samplelist = "{\"rvars\": {\"x\": " ++ show samplelist ++ "}}"

closeJSON :: Show s => [s] -> String
closeJSON samplelist = "{\"close\": true, \"rvars\": {\"x\": " ++ show samplelist ++ "}}"

visualize :: Show s => PrintF x s -> Batch x -> IO ()
visualize f (ls,_) = unless (null ls) $ putStrLn.vizJSON $ f ls

vizClose :: Show s => PrintF x s -> Batch x -> IO ()
vizClose f (ls,_) = unless (null ls) $ putStrLn.closeJSON $ f ls

batchViz :: Show s => PrintF x s -> Int -> BatchAction x IO ()
batchViz f n = 
    let viz = visualize f 
        close = vizClose f
    in pack close $ inBatches viz n

batchPrint :: Show s => PrintF x s -> Int -> BatchAction x IO ()
batchPrint f n =
    let p fn (ls,_) = unless (null ls) $ print $ fn ls
    in pack (p f) $ inBatches (p f) n