content: {-# LANGUAGE Rank2Types #-}
  
  module Stan.Simulate where
  
  import Stan.Data
  import Stan.AST hiding (normal, gamma)
  import Control.Monad.State
  import qualified Data.Map.Strict as Map
  import qualified Data.Vector as V
  import Data.Random
  import Data.Maybe
  import Data.Monoid
  import Data.Maybe
  import Data.List
  import Data.Random.Source.PureMT
  import Data.Random.Distribution.Normal
  import Data.Random.Distribution.Gamma
  
  eval :: StanEnv -> Expr -> StanValue
  eval _ (LitInt x) = VInt x
  eval _ (LitFloat x) = VDouble $ realToFrac x
  eval e (BinOp "+" x y) = numOp2 (+) (eval e x) (eval e y)
  eval e (BinOp "-" x y) = numOp2 (-) (eval e x) (eval e y)
  eval e (BinOp "*" x y) = numOp2 (*) (eval e x) (eval e y)
  --eval e (BinOp "/" x y) = numOp2 (/) (eval e x) (eval e y)
  eval e (Apply "sqrt" [x]) = fracOp sqrt (eval e x)
  eval e (Apply "exp" [x]) = fracOp exp (eval e x)
  eval e (Apply "dot_product" [ev1,ev2]) =
    let VArray v1 = eval e ev1
        VArray v2 = eval e ev2
    in esum $ V.toList $ V.zipWith (numOp2 (*)) v1 v2
  eval e (Var vnm) = case Map.lookup vnm e of
                       Nothing-> error $ "variable not found: "++vnm
                       Just v -> v
  eval e (Ix arrE dimEs ) =
    let arrv = eval e arrE
        dimVs = map (eval e) dimEs
    in case (arrv, dimVs) of
         (VArray v, [VInt ix]) -> v V.! (ix-1)
  eval _ ex = error $ "eval: "++show ex
  
  esum :: [StanValue] -> StanValue
  esum [] = VDouble 0
  esum (v:vs) = numOp2 (+) v $ esum vs
  
  numOp2 :: (forall a . Num a => a -> a -> a) -> StanValue -> StanValue -> StanValue
  numOp2 f (VDouble x) (VDouble y) = VDouble $ f x y
  numOp2 f (VInt x) (VInt y) = VInt $ f x y
  
  fracOp :: (forall a . Floating a => a -> a) -> StanValue -> StanValue
  fracOp f (VDouble x) = VDouble $ f x
  
  simulate :: Decl -> State StanEnv ()
  simulate (_ ::: _) = return ()
  simulate (Print _ _) = return ()
  simulate (For vnm loe hie ds) = do
    env <- get
    let VInt lo = eval env loe
        VInt hi = eval env hie
        oldv = Map.lookup vnm env
    forM_ [lo..hi] $ \i -> do
      modify $ Map.insert vnm (VInt i)
      mapM_ simulate ds
    case oldv of
      Nothing -> modify $ Map.delete vnm
      Just x -> modify $ Map.insert vnm x
  simulate ((vnm,ixs) := e) = do
    env <- get
    setVar vnm ixs $ eval env e
  simulate ((vnm,ixs) :~ (distnm, argEs)) = do
    env <- get
    let vixs = map (eval env) ixs
        mcmcNm = "__mcmc__"++vnm++postfixIxs vixs
    case Map.lookup mcmcNm env of
      Nothing -> do
        let argVs = map (eval env) argEs
            Just (VSeed oldSeed) = Map.lookup "__seed" env
        let (v, newSeed) = simDist distnm argVs oldSeed
        modify $ Map.insert "__seed" (VSeed newSeed)
        setVar vnm ixs v
      Just (VSamples dbls) -> do
        let Just (VInt n) = Map.lookup "__sampleix" env
            v = dbls V.! n
        setVar vnm ixs v
  
  postfixIxs :: [StanValue] -> [Char]
  postfixIxs [] = ""
  postfixIxs (VInt n : vs) = '.':show n++postfixIxs vs
  
  
  setVar :: Var -> [Expr] -> StanValue -> State StanEnv ()
  setVar vnm [] v = modify $ Map.insert vnm v
  setVar vnm [eix] v = do
    env <- get
    let VInt ix =  eval env eix
    modify $ Map.alter (vSet ix v) vnm
  
  simDist :: String -> [StanValue] -> PureMT-> (StanValue, PureMT)
  simDist "normal" [VDouble mu, VDouble sd] seed
    = sampleState (fmap VDouble $ normal mu sd) seed
  simDist "gamma" [VDouble a, VDouble b] seed
    = sampleState (fmap VDouble $ gamma a b) seed
  simDist d args _ = error $ "simDist: "++d++" "++show args
  
  vSet :: Int -> StanValue -> Maybe StanValue -> Maybe StanValue
  vSet ix newVal (Just (VArray v))
    | ix <= V.length v = Just $ VArray $ v V.// [(ix-1,newVal)]
    | ix == V.length v +1 = Just $ VArray $ v `V.snoc` newVal
  vSet 1 newVal Nothing = Just $ VArray $ V.singleton newVal
  vSet ix _ mv = error $ "vSet: "++show ix++" "++show mv
  
  
  runSimulateOnce :: [Stan] -> StanEnv -> StanEnv
  runSimulateOnce sts = execState (mapM_ simulate ds) where
    ds = (concat [ds' | Model ds' <- sts])
  
  runSimulate :: Int -> [Stan] -> StanEnv -> [StanEnv]
  runSimulate 0 _ _ = []
  runSimulate n sts initEnv =
    let samEnv = Map.insert "__sampleix" (VInt n) initEnv
        nextEnv = runSimulateOnce sts samEnv
        Just seed = Map.lookup "__seed" nextEnv
        recEnv = Map.insert "__seed" seed initEnv
    in nextEnv : runSimulate (n-1) sts recEnv
  
  seedEnv :: PureMT -> StanEnv
  seedEnv s = Map.fromList [("__seed", VSeed s), ("__sampleix", VInt 0)]
  
  
  
  mcmcToEnv :: Map.Map String [Double] -> StanEnv
  mcmcToEnv = Map.mapKeys ("__mcmc__"++) . Map.map f where
        f xs = VSamples $ V.fromList $ map VDouble xs
  
  avgVar :: [StanEnv] -> String -> StanValue
  avgVar senvs vnm =
    let n = realToFrac $ length senvs
        sm = foldl1' sumValue $ mapMaybe (Map.lookup vnm) senvs
    in scaleValue (1/n) sm