{-# LANGUAGE CPP
           , GADTs
           , DataKinds
           , TypeFamilies
           , FlexibleContexts
           , UndecidableInstances
           , LambdaCase
           , OverloadedStrings
           , Rank2Types
           #-}

{-# OPTIONS_GHC -Wall -fwarn-tabs -fsimpl-tick-factor=1000 #-}
module Language.Hakaru.Runtime.Prelude where

#if __GLASGOW_HASKELL__ < 710
import           Data.Functor                    ((<$>))
import           Control.Applicative             (Applicative(..))
#endif
import           Data.Foldable                   as F
import qualified System.Random.MWC               as MWC
import qualified System.Random.MWC.Distributions as MWCD
import           Data.Number.Natural
import           Data.STRef
import qualified Data.Vector                     as V
import qualified Data.Vector.Unboxed             as U
import qualified Data.Vector.Generic             as G
import           Control.Monad
import           Control.Monad.ST
import           Prelude                         hiding (product, init)
import           Language.Hakaru.Runtime.CmdLine (Measure(..), makeMeasure)

type family MinBoxVec (v1 :: * -> *) (v2 :: * -> *) :: * -> *
type instance MinBoxVec V.Vector v        = V.Vector
type instance MinBoxVec v        V.Vector = V.Vector
type instance MinBoxVec U.Vector U.Vector = U.Vector

type family MayBoxVec a :: * -> *
type instance MayBoxVec ()           = U.Vector
type instance MayBoxVec Int          = U.Vector
type instance MayBoxVec Double       = U.Vector
type instance MayBoxVec Bool         = U.Vector
type instance MayBoxVec (U.Vector a) = V.Vector
type instance MayBoxVec (V.Vector a) = V.Vector
type instance MayBoxVec (a,b)        = MinBoxVec (MayBoxVec a) (MayBoxVec b)

type Prob = Double

lam :: (a -> b) -> a -> b
lam :: (a -> b) -> a -> b
lam = (a -> b) -> a -> b
forall a. a -> a
id
{-# INLINE lam #-}

app :: (a -> b) -> a -> b
app :: (a -> b) -> a -> b
app a -> b
f a
x = a -> b
f a
x
{-# INLINE app #-}

let_ :: a -> (a -> b) -> b
let_ :: a -> (a -> b) -> b
let_ a
x a -> b
f = let x1 :: a
x1 = a
x in a -> b
f a
x1
{-# INLINE let_ #-}

ann_ :: a -> b -> b
ann_ :: a -> b -> b
ann_ a
_ b
a = b
a
{-# INLINE ann_ #-}

uniform :: Double -> Double -> Measure Double
uniform :: Double -> Double -> Measure Double
uniform Double
lo Double
hi = (GenIO -> IO Double) -> Measure Double
forall a. (GenIO -> IO a) -> Measure a
makeMeasure ((GenIO -> IO Double) -> Measure Double)
-> (GenIO -> IO Double) -> Measure Double
forall a b. (a -> b) -> a -> b
$ (Double, Double) -> GenIO -> IO Double
forall a (m :: * -> *).
(Variate a, PrimMonad m) =>
(a, a) -> Gen (PrimState m) -> m a
MWC.uniformR (Double
lo, Double
hi)
{-# INLINE uniform #-}

normal :: Double -> Prob -> Measure Double
normal :: Double -> Double -> Measure Double
normal Double
mu Double
sd = (GenIO -> IO Double) -> Measure Double
forall a. (GenIO -> IO a) -> Measure a
makeMeasure ((GenIO -> IO Double) -> Measure Double)
-> (GenIO -> IO Double) -> Measure Double
forall a b. (a -> b) -> a -> b
$ Double -> Double -> GenIO -> IO Double
forall (m :: * -> *).
PrimMonad m =>
Double -> Double -> Gen (PrimState m) -> m Double
MWCD.normal Double
mu (Double -> Double
fromProb Double
sd)
{-# INLINE normal #-}

beta :: Prob -> Prob -> Measure Prob
beta :: Double -> Double -> Measure Double
beta Double
a Double
b = (GenIO -> IO Double) -> Measure Double
forall a. (GenIO -> IO a) -> Measure a
makeMeasure ((GenIO -> IO Double) -> Measure Double)
-> (GenIO -> IO Double) -> Measure Double
forall a b. (a -> b) -> a -> b
$ \GenIO
g ->
  Double -> Double
unsafeProb (Double -> Double) -> IO Double -> IO Double
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Double -> Double -> GenIO -> IO Double
forall (m :: * -> *).
PrimMonad m =>
Double -> Double -> Gen (PrimState m) -> m Double
MWCD.beta (Double -> Double
fromProb Double
a) (Double -> Double
fromProb Double
b) GenIO
g
{-# INLINE beta #-}

gamma :: Prob -> Prob -> Measure Prob
gamma :: Double -> Double -> Measure Double
gamma Double
a Double
b = (GenIO -> IO Double) -> Measure Double
forall a. (GenIO -> IO a) -> Measure a
makeMeasure ((GenIO -> IO Double) -> Measure Double)
-> (GenIO -> IO Double) -> Measure Double
forall a b. (a -> b) -> a -> b
$ \GenIO
g ->
  Double -> Double
unsafeProb (Double -> Double) -> IO Double -> IO Double
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Double -> Double -> GenIO -> IO Double
forall (m :: * -> *).
PrimMonad m =>
Double -> Double -> Gen (PrimState m) -> m Double
MWCD.gamma (Double -> Double
fromProb Double
a) (Double -> Double
fromProb Double
b) GenIO
g
{-# INLINE gamma #-}

categorical :: MayBoxVec Prob Prob -> Measure Int
categorical :: MayBoxVec Double Double -> Measure Int
categorical MayBoxVec Double Double
a = (GenIO -> IO Int) -> Measure Int
forall a. (GenIO -> IO a) -> Measure a
makeMeasure ((GenIO -> IO Int) -> Measure Int)
-> (GenIO -> IO Int) -> Measure Int
forall a b. (a -> b) -> a -> b
$ Vector Double -> GenIO -> IO Int
forall (m :: * -> *) (v :: * -> *).
(PrimMonad m, Vector v Double) =>
v Double -> Gen (PrimState m) -> m Int
MWCD.categorical Vector Double
MayBoxVec Double Double
a
{-# INLINE categorical #-}

plate :: (G.Vector (MayBoxVec a) a) =>
         Int -> (Int -> Measure a) -> Measure (MayBoxVec a a)
plate :: Int -> (Int -> Measure a) -> Measure (MayBoxVec a a)
plate Int
n Int -> Measure a
f = Int -> (Int -> Measure a) -> Measure (MayBoxVec a a)
forall (m :: * -> *) (v :: * -> *) a.
(Monad m, Vector v a) =>
Int -> (Int -> m a) -> m (v a)
G.generateM (Int -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n) ((Int -> Measure a) -> Measure (MayBoxVec a a))
-> (Int -> Measure a) -> Measure (MayBoxVec a a)
forall a b. (a -> b) -> a -> b
$ \Int
x ->
             Int -> Measure a
f (Int -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
x)
{-# INLINE plate #-}

bucket :: Int -> Int -> (forall s. Reducer () s a) -> a
bucket :: Int -> Int -> (forall s. Reducer () s a) -> a
bucket Int
b Int
e forall s. Reducer () s a
r = (forall s. ST s a) -> a
forall a. (forall s. ST s a) -> a
runST
             ((forall s. ST s a) -> a) -> (forall s. ST s a) -> a
forall a b. (a -> b) -> a -> b
$ case Reducer () s a
forall s. Reducer () s a
r of Reducer{init :: ()
init=() -> ST s cell
initR,accum :: ()
accum=() -> Int -> cell -> ST s ()
accumR,done :: ()
done=cell -> ST s a
doneR} -> do
                          cell
s' <- () -> ST s cell
initR ()
                          (Int -> ST s ()) -> [Int] -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
F.mapM_ (\Int
i -> () -> Int -> cell -> ST s ()
accumR () Int
i cell
s') [Int
b .. Int
e Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
                          cell -> ST s a
doneR cell
s'
{-# INLINE bucket #-}

data Reducer xs s a =
    forall cell.
    Reducer { ()
init  :: xs -> ST s cell
            , ()
accum :: xs -> Int -> cell -> ST s ()
            , ()
done  :: cell -> ST s a
            }

r_fanout :: Reducer xs s a
         -> Reducer xs s b
         -> Reducer xs s (a,b)
r_fanout :: Reducer xs s a -> Reducer xs s b -> Reducer xs s (a, b)
r_fanout Reducer{init :: ()
init=xs -> ST s cell
initA,accum :: ()
accum=xs -> Int -> cell -> ST s ()
accumA,done :: ()
done=cell -> ST s a
doneA}
         Reducer{init :: ()
init=xs -> ST s cell
initB,accum :: ()
accum=xs -> Int -> cell -> ST s ()
accumB,done :: ()
done=cell -> ST s b
doneB} = Reducer :: forall xs s a cell.
(xs -> ST s cell)
-> (xs -> Int -> cell -> ST s ())
-> (cell -> ST s a)
-> Reducer xs s a
Reducer
   { init :: xs -> ST s (cell, cell)
init  = \xs
xs       -> (cell -> cell -> (cell, cell))
-> ST s cell -> ST s cell -> ST s (cell, cell)
forall (m :: * -> *) a1 a2 r.
Monad m =>
(a1 -> a2 -> r) -> m a1 -> m a2 -> m r
liftM2 (,) (xs -> ST s cell
initA xs
xs) (xs -> ST s cell
initB xs
xs)
   , accum :: xs -> Int -> (cell, cell) -> ST s ()
accum = \xs
bs Int
i (cell
s1, cell
s2) ->
             xs -> Int -> cell -> ST s ()
accumA xs
bs Int
i cell
s1 ST s () -> ST s () -> ST s ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> xs -> Int -> cell -> ST s ()
accumB xs
bs Int
i cell
s2
   , done :: (cell, cell) -> ST s (a, b)
done  = \(cell
s1, cell
s2) -> (a -> b -> (a, b)) -> ST s a -> ST s b -> ST s (a, b)
forall (m :: * -> *) a1 a2 r.
Monad m =>
(a1 -> a2 -> r) -> m a1 -> m a2 -> m r
liftM2 (,) (cell -> ST s a
doneA cell
s1) (cell -> ST s b
doneB cell
s2)
   }
{-# INLINE r_fanout #-}

r_index :: (G.Vector (MayBoxVec a) a)
        => (xs -> Int)
        -> ((Int, xs) -> Int)
        -> Reducer (Int, xs) s a
        -> Reducer xs s (MayBoxVec a a)
r_index :: (xs -> Int)
-> ((Int, xs) -> Int)
-> Reducer (Int, xs) s a
-> Reducer xs s (MayBoxVec a a)
r_index xs -> Int
n (Int, xs) -> Int
f Reducer{init :: ()
init=(Int, xs) -> ST s cell
initR,accum :: ()
accum=(Int, xs) -> Int -> cell -> ST s ()
accumR,done :: ()
done=cell -> ST s a
doneR} = Reducer :: forall xs s a cell.
(xs -> ST s cell)
-> (xs -> Int -> cell -> ST s ())
-> (cell -> ST s a)
-> Reducer xs s a
Reducer
   { init :: xs -> ST s (Vector cell)
init  = \xs
xs -> Int -> (Int -> ST s cell) -> ST s (Vector cell)
forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> m a) -> m (Vector a)
V.generateM (xs -> Int
n xs
xs) (\Int
b -> (Int, xs) -> ST s cell
initR (Int
b, xs
xs))
   , accum :: xs -> Int -> Vector cell -> ST s ()
accum = \xs
bs Int
i Vector cell
v ->
             let ov :: Int
ov = (Int, xs) -> Int
f (Int
i, xs
bs) in
             (Int, xs) -> Int -> cell -> ST s ()
accumR (Int
ov,xs
bs) Int
i (Vector cell
v Vector cell -> Int -> cell
forall a. Vector a -> Int -> a
V.! Int
ov)
   , done :: Vector cell -> ST s (MayBoxVec a a)
done  = \Vector cell
v -> (Vector a -> MayBoxVec a a)
-> ST s (Vector a) -> ST s (MayBoxVec a a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Vector a -> MayBoxVec a a
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
G.convert ((cell -> ST s a) -> Vector cell -> ST s (Vector a)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Vector a -> m (Vector b)
V.mapM cell -> ST s a
doneR Vector cell
v)
   }
{-# INLINE r_index #-}

r_split :: ((Int, xs) -> Bool)
        -> Reducer xs s a
        -> Reducer xs s b
        -> Reducer xs s (a,b)
r_split :: ((Int, xs) -> Bool)
-> Reducer xs s a -> Reducer xs s b -> Reducer xs s (a, b)
r_split (Int, xs) -> Bool
b Reducer{init :: ()
init=xs -> ST s cell
initA,accum :: ()
accum=xs -> Int -> cell -> ST s ()
accumA,done :: ()
done=cell -> ST s a
doneA}
          Reducer{init :: ()
init=xs -> ST s cell
initB,accum :: ()
accum=xs -> Int -> cell -> ST s ()
accumB,done :: ()
done=cell -> ST s b
doneB} = Reducer :: forall xs s a cell.
(xs -> ST s cell)
-> (xs -> Int -> cell -> ST s ())
-> (cell -> ST s a)
-> Reducer xs s a
Reducer
   { init :: xs -> ST s (cell, cell)
init  = \xs
xs -> (cell -> cell -> (cell, cell))
-> ST s cell -> ST s cell -> ST s (cell, cell)
forall (m :: * -> *) a1 a2 r.
Monad m =>
(a1 -> a2 -> r) -> m a1 -> m a2 -> m r
liftM2 (,) (xs -> ST s cell
initA xs
xs) (xs -> ST s cell
initB xs
xs)
   , accum :: xs -> Int -> (cell, cell) -> ST s ()
accum = \xs
bs Int
i (cell
s1, cell
s2) ->
             if ((Int, xs) -> Bool
b (Int
i,xs
bs)) then xs -> Int -> cell -> ST s ()
accumA xs
bs Int
i cell
s1 else xs -> Int -> cell -> ST s ()
accumB xs
bs Int
i cell
s2
   , done :: (cell, cell) -> ST s (a, b)
done  = \(cell
s1, cell
s2) -> (a -> b -> (a, b)) -> ST s a -> ST s b -> ST s (a, b)
forall (m :: * -> *) a1 a2 r.
Monad m =>
(a1 -> a2 -> r) -> m a1 -> m a2 -> m r
liftM2 (,) (cell -> ST s a
doneA cell
s1) (cell -> ST s b
doneB cell
s2)
   }
{-# INLINE r_split #-}

r_add :: Num a => ((Int, xs) -> a) -> Reducer xs s a
r_add :: ((Int, xs) -> a) -> Reducer xs s a
r_add (Int, xs) -> a
e = Reducer :: forall xs s a cell.
(xs -> ST s cell)
-> (xs -> Int -> cell -> ST s ())
-> (cell -> ST s a)
-> Reducer xs s a
Reducer
   { init :: xs -> ST s (STRef s a)
init  = \xs
_ -> a -> ST s (STRef s a)
forall a s. a -> ST s (STRef s a)
newSTRef a
0
   , accum :: xs -> Int -> STRef s a -> ST s ()
accum = \xs
bs Int
i STRef s a
s ->
             STRef s a -> (a -> a) -> ST s ()
forall s a. STRef s a -> (a -> a) -> ST s ()
modifySTRef' STRef s a
s (a -> a -> a
forall a. Num a => a -> a -> a
+ ((Int, xs) -> a
e (Int
i,xs
bs)))
   , done :: STRef s a -> ST s a
done  = STRef s a -> ST s a
forall s a. STRef s a -> ST s a
readSTRef
   }
{-# INLINE r_add #-}

r_nop :: Reducer xs s ()
r_nop :: Reducer xs s ()
r_nop = Reducer :: forall xs s a cell.
(xs -> ST s cell)
-> (xs -> Int -> cell -> ST s ())
-> (cell -> ST s a)
-> Reducer xs s a
Reducer
   { init :: xs -> ST s ()
init  = \xs
_ -> () -> ST s ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
   , accum :: xs -> Int -> () -> ST s ()
accum = \xs
_ Int
_ ()
_ -> () -> ST s ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
   , done :: () -> ST s ()
done  = \()
_ -> () -> ST s ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
   }
{-# INLINE r_nop #-}

pair :: a -> b -> (a, b)
pair :: a -> b -> (a, b)
pair = (,)
{-# INLINE pair #-}

true, false :: Bool
true :: Bool
true  = Bool
True
false :: Bool
false = Bool
False

nothing :: Maybe a
nothing :: Maybe a
nothing = Maybe a
forall a. Maybe a
Nothing

just :: a -> Maybe a
just :: a -> Maybe a
just = a -> Maybe a
forall a. a -> Maybe a
Just

left :: a -> Either a b
left :: a -> Either a b
left = a -> Either a b
forall a b. a -> Either a b
Left

right :: b -> Either a b
right :: b -> Either a b
right = b -> Either a b
forall a b. b -> Either a b
Right

unit :: ()
unit :: ()
unit = ()

data Pattern = PVar | PWild
newtype Branch a b =
    Branch { Branch a b -> a -> Maybe b
extract :: a -> Maybe b }

ptrue, pfalse :: a -> Branch Bool a
ptrue :: a -> Branch Bool a
ptrue  a
b = Branch :: forall a b. (a -> Maybe b) -> Branch a b
Branch { extract :: Bool -> Maybe a
extract = Bool -> a -> Bool -> Maybe a
forall a. Bool -> a -> Bool -> Maybe a
extractBool Bool
True  a
b }
pfalse :: a -> Branch Bool a
pfalse a
b = Branch :: forall a b. (a -> Maybe b) -> Branch a b
Branch { extract :: Bool -> Maybe a
extract = Bool -> a -> Bool -> Maybe a
forall a. Bool -> a -> Bool -> Maybe a
extractBool Bool
False a
b }
{-# INLINE ptrue  #-}
{-# INLINE pfalse #-}

extractBool :: Bool -> a -> Bool -> Maybe a
extractBool :: Bool -> a -> Bool -> Maybe a
extractBool Bool
b a
a Bool
p | Bool
p Bool -> Bool -> Bool
forall a. Eq a => a -> a -> Bool
== Bool
b     = a -> Maybe a
forall a. a -> Maybe a
Just a
a
                  | Bool
otherwise  = Maybe a
forall a. Maybe a
Nothing
{-# INLINE extractBool #-}

pnothing :: b -> Branch (Maybe a) b
pnothing :: b -> Branch (Maybe a) b
pnothing b
b = Branch :: forall a b. (a -> Maybe b) -> Branch a b
Branch { extract :: Maybe a -> Maybe b
extract = \Maybe a
ma -> case Maybe a
ma of
                                         Maybe a
Nothing -> b -> Maybe b
forall a. a -> Maybe a
Just b
b
                                         Just a
_  -> Maybe b
forall a. Maybe a
Nothing }

pjust :: Pattern -> (a -> b) -> Branch (Maybe a) b
pjust :: Pattern -> (a -> b) -> Branch (Maybe a) b
pjust Pattern
PVar a -> b
c = Branch :: forall a b. (a -> Maybe b) -> Branch a b
Branch { extract :: Maybe a -> Maybe b
extract = \Maybe a
ma -> case Maybe a
ma of
                                           Maybe a
Nothing -> Maybe b
forall a. Maybe a
Nothing
                                           Just a
x  -> b -> Maybe b
forall a. a -> Maybe a
Just (a -> b
c a
x) }
pjust Pattern
_ a -> b
_ = [Char] -> Branch (Maybe a) b
forall a. HasCallStack => [Char] -> a
error [Char]
"TODO: Runtime.Prelude{pjust}"

pleft :: Pattern -> (a -> c) -> Branch (Either a b) c
pleft :: Pattern -> (a -> c) -> Branch (Either a b) c
pleft Pattern
PVar a -> c
f = Branch :: forall a b. (a -> Maybe b) -> Branch a b
Branch { extract :: Either a b -> Maybe c
extract = \Either a b
ma -> case Either a b
ma of
                                           Right b
_ -> Maybe c
forall a. Maybe a
Nothing
                                           Left a
x -> c -> Maybe c
forall a. a -> Maybe a
Just (a -> c
f a
x) }
pleft Pattern
_ a -> c
_ = [Char] -> Branch (Either a b) c
forall a. HasCallStack => [Char] -> a
error [Char]
"TODO: Runtime.Prelude{pLeft}"

pright :: Pattern -> (b -> c) -> Branch (Either a b) c
pright :: Pattern -> (b -> c) -> Branch (Either a b) c
pright Pattern
PVar b -> c
f = Branch :: forall a b. (a -> Maybe b) -> Branch a b
Branch { extract :: Either a b -> Maybe c
extract = \Either a b
ma -> case Either a b
ma of
                                            Left a
_ -> Maybe c
forall a. Maybe a
Nothing
                                            Right b
x -> c -> Maybe c
forall a. a -> Maybe a
Just (b -> c
f b
x) }
pright Pattern
_ b -> c
_ = [Char] -> Branch (Either a b) c
forall a. HasCallStack => [Char] -> a
error [Char]
"TODO: Runtime.Prelude{pRight}"


ppair :: Pattern -> Pattern -> (x -> y -> b) -> Branch (x,y) b
ppair :: Pattern -> Pattern -> (x -> y -> b) -> Branch (x, y) b
ppair Pattern
PVar  Pattern
PVar x -> y -> b
c = Branch :: forall a b. (a -> Maybe b) -> Branch a b
Branch { extract :: (x, y) -> Maybe b
extract = (\(x
x,y
y) -> b -> Maybe b
forall a. a -> Maybe a
Just (x -> y -> b
c x
x y
y)) }
ppair Pattern
_     Pattern
_    x -> y -> b
_ = [Char] -> Branch (x, y) b
forall a. HasCallStack => [Char] -> a
error [Char]
"ppair: TODO"

uncase_ :: Maybe a -> a
uncase_ :: Maybe a -> a
uncase_ (Just a
a) = a
a
uncase_ Maybe a
Nothing  = [Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"case_: unable to match any branches"
{-# INLINE uncase_ #-}

case_ :: a -> [Branch a b] -> b
case_ :: a -> [Branch a b] -> b
case_ a
e [Branch a b
c1]     = Maybe b -> b
forall a. Maybe a -> a
uncase_ (Branch a b -> a -> Maybe b
forall a b. Branch a b -> a -> Maybe b
extract Branch a b
c1 a
e)
case_ a
e [Branch a b
c1, Branch a b
c2] = Maybe b -> b
forall a. Maybe a -> a
uncase_ (Branch a b -> a -> Maybe b
forall a b. Branch a b -> a -> Maybe b
extract Branch a b
c1 a
e Maybe b -> Maybe b -> Maybe b
forall (m :: * -> *) a. MonadPlus m => m a -> m a -> m a
`mplus` Branch a b -> a -> Maybe b
forall a b. Branch a b -> a -> Maybe b
extract Branch a b
c2 a
e)
case_ a
e [Branch a b]
bs_      = [Branch a b] -> b
go [Branch a b]
bs_
  where go :: [Branch a b] -> b
go []     = [Char] -> b
forall a. HasCallStack => [Char] -> a
error [Char]
"case_: unable to match any branches"
        go (Branch a b
b:[Branch a b]
bs) = case Branch a b -> a -> Maybe b
forall a b. Branch a b -> a -> Maybe b
extract Branch a b
b a
e of
                      Just b
b' -> b
b'
                      Maybe b
Nothing -> [Branch a b] -> b
go [Branch a b]
bs
{-# INLINE case_ #-}

branch :: (c -> Branch a b) -> c -> Branch a b
branch :: (c -> Branch a b) -> c -> Branch a b
branch c -> Branch a b
pat c
body = c -> Branch a b
pat c
body
{-# INLINE branch #-}

dirac :: a -> Measure a
dirac :: a -> Measure a
dirac = a -> Measure a
forall (m :: * -> *) a. Monad m => a -> m a
return
{-# INLINE dirac #-}

pose :: Prob -> Measure a -> Measure a
pose :: Double -> Measure a -> Measure a
pose Double
_ Measure a
a = Measure a
a
{-# INLINE pose #-}

superpose :: [(Prob, Measure a)]
          -> Measure a
superpose :: [(Double, Measure a)] -> Measure a
superpose [(Double, Measure a)]
pms = do
  Int
i <- MayBoxVec Double Double -> Measure Int
categorical ([Double] -> Vector Double
forall (v :: * -> *) a. Vector v a => [a] -> v a
G.fromList ([Double] -> Vector Double) -> [Double] -> Vector Double
forall a b. (a -> b) -> a -> b
$ ((Double, Measure a) -> Double)
-> [(Double, Measure a)] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (Double, Measure a) -> Double
forall a b. (a, b) -> a
fst [(Double, Measure a)]
pms)
  (Double, Measure a) -> Measure a
forall a b. (a, b) -> b
snd ([(Double, Measure a)]
pms [(Double, Measure a)] -> Int -> (Double, Measure a)
forall a. [a] -> Int -> a
!! Int
i)
{-# INLINE superpose #-}

reject :: Measure a
reject :: Measure a
reject = (GenIO -> IO (Maybe a)) -> Measure a
forall a. (GenIO -> IO (Maybe a)) -> Measure a
Measure ((GenIO -> IO (Maybe a)) -> Measure a)
-> (GenIO -> IO (Maybe a)) -> Measure a
forall a b. (a -> b) -> a -> b
$ \GenIO
_ -> Maybe a -> IO (Maybe a)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe a
forall a. Maybe a
Nothing

nat_ :: Int -> Int
nat_ :: Int -> Int
nat_ = Int -> Int
forall a. a -> a
id

int_ :: Int -> Int
int_ :: Int -> Int
int_ = Int -> Int
forall a. a -> a
id

unsafeNat :: Int -> Int
unsafeNat :: Int -> Int
unsafeNat = Int -> Int
forall a. a -> a
id

nat2prob :: Int -> Prob
nat2prob :: Int -> Double
nat2prob = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral

fromInt  :: Int -> Double
fromInt :: Int -> Double
fromInt  = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral

nat2int  :: Int -> Int
nat2int :: Int -> Int
nat2int  = Int -> Int
forall a. a -> a
id

nat2real :: Int -> Double
nat2real :: Int -> Double
nat2real = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral

fromProb :: Prob -> Double
fromProb :: Double -> Double
fromProb = Double -> Double
forall a. a -> a
id

unsafeProb :: Double -> Prob
unsafeProb :: Double -> Double
unsafeProb = Double -> Double
forall a. a -> a
id

real_ :: Rational -> Double
real_ :: Rational -> Double
real_ = Rational -> Double
forall a. Fractional a => Rational -> a
fromRational

prob_ :: NonNegativeRational -> Prob
prob_ :: NonNegativeRational -> Double
prob_ = Rational -> Double
forall a. Fractional a => Rational -> a
fromRational (Rational -> Double)
-> (NonNegativeRational -> Rational)
-> NonNegativeRational
-> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NonNegativeRational -> Rational
fromNonNegativeRational

infinity :: Double
infinity :: Double
infinity = Double
1Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/Double
0

abs_ :: Num a => a -> a
abs_ :: a -> a
abs_ = a -> a
forall a. Num a => a -> a
abs

thRootOf :: Int -> Prob -> Prob
thRootOf :: Int -> Double -> Double
thRootOf Int
a Double
b = Double
b Double -> Double -> Double
forall a. Floating a => a -> a -> a
** (Double -> Double
forall a. Fractional a => a -> a
recip (Double -> Double) -> Double -> Double
forall a b. (a -> b) -> a -> b
$ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
a)
{-# INLINE thRootOf #-}

array
    :: (G.Vector (MayBoxVec a) a)
    => Int
    -> (Int -> a)
    -> MayBoxVec a a
array :: Int -> (Int -> a) -> MayBoxVec a a
array Int
n Int -> a
f = Int -> (Int -> a) -> MayBoxVec a a
forall (v :: * -> *) a. Vector v a => Int -> (Int -> a) -> v a
G.generate (Int -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n) (Int -> a
f (Int -> a) -> (Int -> Int) -> Int -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral)
{-# INLINE array #-}

arrayLit :: (G.Vector (MayBoxVec a) a) => [a] -> MayBoxVec a a
arrayLit :: [a] -> MayBoxVec a a
arrayLit = [a] -> MayBoxVec a a
forall (v :: * -> *) a. Vector v a => [a] -> v a
G.fromList
{-# INLINE arrayLit #-}

(!) :: (G.Vector (MayBoxVec a) a) => MayBoxVec a a -> Int -> a
MayBoxVec a a
a ! :: MayBoxVec a a -> Int -> a
! Int
b = MayBoxVec a a
a MayBoxVec a a -> Int -> a
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
G.! (Int -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
b)
{-# INLINE (!) #-}

size :: (G.Vector (MayBoxVec a) a) => MayBoxVec a a -> Int
size :: MayBoxVec a a -> Int
size MayBoxVec a a
v = Int -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (MayBoxVec a a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
G.length MayBoxVec a a
v)
{-# INLINE size #-}

reduce
    :: (G.Vector (MayBoxVec a) a)
    => (a -> a -> a)
    -> a
    -> MayBoxVec a a
    -> a
reduce :: (a -> a -> a) -> a -> MayBoxVec a a -> a
reduce a -> a -> a
f a
n MayBoxVec a a
v = (a -> a -> a) -> a -> MayBoxVec a a -> a
forall (v :: * -> *) a b.
Vector v a =>
(a -> b -> b) -> b -> v a -> b
G.foldr a -> a -> a
f a
n MayBoxVec a a
v
{-# INLINE reduce #-}

product
    :: Num a
    => Int
    -> Int
    -> (Int -> a)
    -> a
product :: Int -> Int -> (Int -> a) -> a
product Int
a Int
b Int -> a
f = (a -> Int -> a) -> a -> [Int] -> a
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
F.foldl' (\a
x Int
y -> a
x a -> a -> a
forall a. Num a => a -> a -> a
* Int -> a
f Int
y) a
1 [Int
a .. Int
bInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1]
{-# INLINE product #-}

summate
    :: Num a
    => Int
    -> Int
    -> (Int -> a)
    -> a
summate :: Int -> Int -> (Int -> a) -> a
summate Int
a Int
b Int -> a
f = (a -> Int -> a) -> a -> [Int] -> a
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
F.foldl' (\a
x Int
y -> a
x a -> a -> a
forall a. Num a => a -> a -> a
+ Int -> a
f Int
y) a
0 [Int
a .. Int
bInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1]
{-# INLINE summate #-}

run :: Show a
    => MWC.GenIO
    -> Measure a
    -> IO ()
run :: GenIO -> Measure a -> IO ()
run GenIO
g Measure a
k = Measure a -> GenIO -> IO (Maybe a)
forall a. Measure a -> GenIO -> IO (Maybe a)
unMeasure Measure a
k GenIO
g IO (Maybe a) -> (Maybe a -> IO ()) -> IO ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
           Just a
a  -> a -> IO ()
forall a. Show a => a -> IO ()
print a
a
           Maybe a
Nothing -> () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()

iterateM_
    :: Monad m
    => (a -> m a)
    -> a
    -> m b
iterateM_ :: (a -> m a) -> a -> m b
iterateM_ a -> m a
f = a -> m b
g
    where g :: a -> m b
g a
x = a -> m a
f a
x m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= a -> m b
g

withPrint :: Show a => (a -> IO b) -> a -> IO b
withPrint :: (a -> IO b) -> a -> IO b
withPrint a -> IO b
f a
x = a -> IO ()
forall a. Show a => a -> IO ()
print a
x IO () -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> a -> IO b
f a
x