{-# LANGUAGE BangPatterns #-}

-- | Various internal utilities.
module Numeric.MCMC.Util 

where

import Control.Monad.Primitive (PrimMonad, PrimState)
import System.Random.MWC 
import Data.Function           (fix)
import qualified Data.Sequence as Seq
import Data.Foldable           (toList)
import Control.Monad
import Control.Monad.ST
import Data.STRef
import Data.List.Split         (chunksOf)

-- | Given Int, bounds, and generator, generate a different Int in the bound.
genDiffInt :: PrimMonad m => Int -> (Int, Int) -> Gen (PrimState m) -> m Int
genDiffInt a bounds gen = fix $ \loopB -> 
    do  b <- uniformR bounds gen
        if a == b then loopB else return b
{-# INLINE genDiffInt #-}

-- | Tail-recursive mean function.
mean :: [Double] -> Double
mean = go 0.0 0 
    where go :: Double -> Int -> [Double] -> Double
          go !s !l []     = s / fromIntegral l
          go !s !l (x:xs) = go (s + x) (l + 1) xs
{-# INLINE mean #-}

-- | Tail-recursive mean function for lists.
listMean :: [[Double]] -> [Double]
listMean = go (repeat 0) 0
    where go :: [Double] -> Int -> [[Double]] -> [Double]
          go !s !l []       = map (/ fromIntegral l) s
          go !s !l (xs:xss) = go (zipWith (+) s xs) (l + 1) xss
{-# INLINE listMean #-}

-- | Knuth-shuffle a list.  Uses 'Seq' internally.  
-- FIXME Would be nice if it was lazy, to work well with `sample`
shuffle :: [a] -> Seed -> [a]
shuffle xs seed = runST $ do
    xsref <- newSTRef $ Seq.fromList xs
    gen   <- restore seed
    let n = length xs
    forM_ [n-1,n-2..1] $ \i -> do
        j <- uniformR (0, i) gen

        xst <- readSTRef xsref
        let tmp = Seq.index xst j

        modifySTRef xsref $ Seq.update j (Seq.index xst i)
        modifySTRef xsref $ Seq.update i tmp

    result <- readSTRef xsref
    return $ toList result
{-# INLINE shuffle #-}

-- | Sample from a list without replacement.
sample :: Int -> [a] -> Seed -> [a]
sample k xs seed = take k $ shuffle xs seed
{-# INLINE sample #-}

-- | Autocovariance function.  `m` is the maximum lag.
acov ::  [Double] -> Int -> [Double]
acov xs m = go [] m xc
    where xc = take (l - m) $ map (subtract (mean xs)) xs
          l  = length xs
          go !s0 !j0 ys | j0 < 0    = map (/ fromIntegral (l - m)) (reverse s0)
                        | otherwise = go (sum (zipWith (*) xc ys) : s0) (j0 - 1) (drop 1 ys)
{-# INLINE acov #-}

-- | (mean, standard error, integrated autocorrelation time)
acor :: [Double] -> (Double, Double, Double)
acor xs = let t0 = d0 / head ac0 
          in  if   t0*fromIntegral winmult < fromIntegral maxlag 
              then (mean xs, sqrt (d0 / fromIntegral l0), t0)
              else (mean xs, sqrt (df / fromIntegral l0), df / head acf)
  where 
    -- Initial values and constants
    (l0, ac0, d0) = (length xs, acov xs maxlag, head ac0 + 2 * sum (tail ac0)) 
    (taumax, winmult) = (2, 5) :: (Int, Int)   
    maxlag = taumax*winmult

    -- Final autocovariance
    (sig1, acf) = go xs ac0 (sqrt (d0 / fromIntegral l0)) (d0 / head ac0)
    df          = 0.25 * sig1 * sig1 * fromIntegral l0

    -- The recursive worker
    go ys ac sig !tau 
      | tau*fromIntegral winmult < fromIntegral maxlag = (sig, ac)
      | otherwise = 
          let (ys1, ac1, d, t1) = (joiner ys, acov ys1 maxlag, head ac1 + (2 :: Double) * sum (tail ac1), d / head ac1)
          in  if ys1 == [] then (sig, ac) else go ys1 ac1 (sqrt (d / fromIntegral (length ys))) t1
        where 
          joiner zs = map sum $ take (truncate $ fromIntegral (length zs) / (2 :: Double)) (chunksOf 2 zs)
{-# INLINE acor #-}