{-# LANGUAGE CPP
           , GADTs
           , EmptyCase
           , KindSignatures
           , DataKinds
           , TypeOperators
           , NoImplicitPrelude
           , ScopedTypeVariables
           , FlexibleContexts
           , ViewPatterns
           #-}

{-# OPTIONS_GHC -Wall -fwarn-tabs #-}
----------------------------------------------------------------
--                                                    2016.05.24
-- |
-- Module      :  Language.Hakaru.Expect
-- Copyright   :  Copyright (c) 2016 the Hakaru team
-- License     :  BSD3
-- Maintainer  :  wren@community.haskell.org
-- Stability   :  experimental
-- Portability :  GHC-only
--
--
----------------------------------------------------------------
module Language.Hakaru.Expect
    ( normalize
    , total
    , expect, expectInCtx, determineExpect
    ) where

import           Prelude               (($), (.), error, reverse, Maybe(..))
import qualified Data.Text             as Text
import           Data.Functor          ((<$>))
import qualified Data.Foldable         as F
import qualified Data.Traversable      as T
import qualified Data.List.NonEmpty    as NE
import           Control.Monad

import Language.Hakaru.Syntax.IClasses (Some2(..), Functor11(..))
import Language.Hakaru.Types.DataKind
import Language.Hakaru.Types.HClasses
import Language.Hakaru.Types.Sing
import Language.Hakaru.Syntax.ABT
import Language.Hakaru.Syntax.Datum
import Language.Hakaru.Syntax.DatumABT
import Language.Hakaru.Syntax.Transform (TransformCtx(..), minimalCtx)
import Language.Hakaru.Syntax.AST               hiding (Expect)
import qualified Language.Hakaru.Syntax.AST     as AST
import Language.Hakaru.Syntax.TypeOf            (typeOf)
import qualified Language.Hakaru.Syntax.Prelude as P
import Language.Hakaru.Evaluation.Types
import Language.Hakaru.Evaluation.ExpectMonad

#ifdef __TRACE_DISINTEGRATE__
import Prelude                          (show, (++))
import qualified Text.PrettyPrint       as PP
import Language.Hakaru.Pretty.Haskell   (pretty)
import Language.Hakaru.Evaluation.Types (ppStatement)
import Debug.Trace                      (trace)
#endif

----------------------------------------------------------------

-- | Convert an arbitrary measure into a probability measure; i.e.,
-- reweight things so that the total weight\/mass is 1.
normalize
    :: (ABT Term abt) => abt '[] ('HMeasure a) -> abt '[] ('HMeasure a)
normalize :: abt '[] ('HMeasure a) -> abt '[] ('HMeasure a)
normalize abt '[] ('HMeasure a)
m = abt '[] 'HProb -> abt '[] ('HMeasure a) -> abt '[] ('HMeasure a)
forall (abt :: [Hakaru] -> Hakaru -> *) (w :: Hakaru).
ABT Term abt =>
abt '[] 'HProb -> abt '[] ('HMeasure w) -> abt '[] ('HMeasure w)
P.withWeight (abt '[] 'HProb -> abt '[] 'HProb
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
(ABT Term abt, HFractional_ a) =>
abt '[] a -> abt '[] a
P.recip (abt '[] 'HProb -> abt '[] 'HProb)
-> abt '[] 'HProb -> abt '[] 'HProb
forall a b. (a -> b) -> a -> b
$ abt '[] ('HMeasure a) -> abt '[] 'HProb
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] ('HMeasure a) -> abt '[] 'HProb
total abt '[] ('HMeasure a)
m) abt '[] ('HMeasure a)
m


-- | Compute the total weight\/mass of a measure.
total :: (ABT Term abt) => abt '[] ('HMeasure a) -> abt '[] 'HProb
total :: abt '[] ('HMeasure a) -> abt '[] 'HProb
total abt '[] ('HMeasure a)
m =
    abt '[] ('HMeasure a) -> abt '[a] 'HProb -> abt '[] 'HProb
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] ('HMeasure a) -> abt '[a] 'HProb -> abt '[] 'HProb
expect abt '[] ('HMeasure a)
m (abt '[a] 'HProb -> abt '[] 'HProb)
-> ((abt '[] a -> abt '[] 'HProb) -> abt '[a] 'HProb)
-> (abt '[] a -> abt '[] 'HProb)
-> abt '[] 'HProb
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Sing a -> (abt '[] a -> abt '[] 'HProb) -> abt '[a] 'HProb
forall a1 (syn :: ([a1] -> a1 -> *) -> a1 -> *)
       (abt :: [a1] -> a1 -> *) (a2 :: a1) (xs :: [a1]) (b :: a1).
ABT syn abt =>
Text -> Sing a2 -> (abt '[] a2 -> abt xs b) -> abt (a2 : xs) b
binder Text
Text.empty (Sing ('HMeasure a) -> Sing a
forall (a :: Hakaru). Sing ('HMeasure a) -> Sing a
sUnMeasure (Sing ('HMeasure a) -> Sing a) -> Sing ('HMeasure a) -> Sing a
forall a b. (a -> b) -> a -> b
$ abt '[] ('HMeasure a) -> Sing ('HMeasure a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> Sing a
typeOf abt '[] ('HMeasure a)
m) ((abt '[] a -> abt '[] 'HProb) -> abt '[] 'HProb)
-> (abt '[] a -> abt '[] 'HProb) -> abt '[] 'HProb
forall a b. (a -> b) -> a -> b
$ \abt '[] a
_ -> abt '[] 'HProb
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
(ABT Term abt, HSemiring_ a) =>
abt '[] a
P.one

-- TODO: is it actually a _measurable_ function from measurable-functions
-- to probs? If so, shouldn't we also capture that in the types?
--
-- | Convert a measure into its integrator. N.B., the second argument
-- is (a representation of) a measurable function from @a@ to
-- 'HProb@. We represent it as a binding form rather than as @abt
-- '[] (a ':-> 'HProb)@ in order to avoid introducing administrative
-- redexes. We could, instead, have used a Haskell function @abt
-- '[] a -> abt '[] 'HProb@ to eliminate the administrative redexes,
-- but that would introduce other implementation difficulties we'd
-- rather avoid.
expect
    :: (ABT Term abt)
    => abt '[] ('HMeasure a)
    -> abt '[a] 'HProb
    -> abt '[] 'HProb
expect :: abt '[] ('HMeasure a) -> abt '[a] 'HProb -> abt '[] 'HProb
expect = TransformCtx
-> abt '[] ('HMeasure a) -> abt '[a] 'HProb -> abt '[] 'HProb
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
TransformCtx
-> abt '[] ('HMeasure a) -> abt '[a] 'HProb -> abt '[] 'HProb
expectInCtx TransformCtx
minimalCtx

expectInCtx
    :: (ABT Term abt)
    => TransformCtx
    -> abt '[] ('HMeasure a)
    -> abt '[a] 'HProb
    -> abt '[] 'HProb
expectInCtx :: TransformCtx
-> abt '[] ('HMeasure a) -> abt '[a] 'HProb -> abt '[] 'HProb
expectInCtx TransformCtx
ctx abt '[] ('HMeasure a)
e abt '[a] 'HProb
f = Expect abt (abt '[] a)
-> TransformCtx -> abt '[a] 'HProb -> [Some2 abt] -> abt '[] 'HProb
forall (abt :: [Hakaru] -> Hakaru -> *) (f :: * -> *)
       (a :: Hakaru).
(ABT Term abt, Foldable f) =>
Expect abt (abt '[] a)
-> TransformCtx
-> abt '[a] 'HProb
-> f (Some2 abt)
-> abt '[] 'HProb
runExpect (abt '[] ('HMeasure a) -> Expect abt (abt '[] a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] ('HMeasure a) -> Expect abt (abt '[] a)
expectTerm abt '[] ('HMeasure a)
e) TransformCtx
ctx abt '[a] 'HProb
f [abt '[] ('HMeasure a) -> Some2 abt
forall k1 k2 (a :: k1 -> k2 -> *) (i :: k1) (j :: k2).
a i j -> Some2 a
Some2 abt '[] ('HMeasure a)
e, abt '[a] 'HProb -> Some2 abt
forall k1 k2 (a :: k1 -> k2 -> *) (i :: k1) (j :: k2).
a i j -> Some2 a
Some2 abt '[a] 'HProb
f]

-- | A helper which converts residualized `expect' to a `Nothing' instead.
determineExpect
    :: (ABT Term abt)
    => abt '[] 'HProb
    -> Maybe (abt '[] 'HProb)
determineExpect :: abt '[] 'HProb -> Maybe (abt '[] 'HProb)
determineExpect abt '[] 'HProb
e =
  case abt '[] 'HProb
e of
    (abt '[] 'HProb -> View (Term abt) '[] 'HProb
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (xs :: [k]) (a :: k).
ABT syn abt =>
abt xs a -> View (syn abt) xs a
viewABT -> Syn (AST.Transform_ AST.Expect :$ _)) -> Maybe (abt '[] 'HProb)
forall a. Maybe a
Nothing
    abt '[] 'HProb
r -> abt '[] 'HProb -> Maybe (abt '[] 'HProb)
forall a. a -> Maybe a
Just abt '[] 'HProb
r


residualizeExpect
    :: (ABT Term abt)
    => abt '[] ('HMeasure a)
    -> Expect abt (abt '[] a)
residualizeExpect :: abt '[] ('HMeasure a) -> Expect abt (abt '[] a)
residualizeExpect abt '[] ('HMeasure a)
e = do
    -- BUG: is this what we really mean? or do we actually mean the old 'emit' version?
    Variable a
x <- Text -> Sing a -> Expect abt (Variable a)
forall (abt :: [Hakaru] -> Hakaru -> *) (m :: * -> *) (p :: Purity)
       (a :: Hakaru).
EvaluationMonad abt m p =>
Text -> Sing a -> m (Variable a)
freshVar Text
Text.empty (Sing ('HMeasure a) -> Sing a
forall (a :: Hakaru). Sing ('HMeasure a) -> Sing a
sUnMeasure (Sing ('HMeasure a) -> Sing a) -> Sing ('HMeasure a) -> Sing a
forall a b. (a -> b) -> a -> b
$ abt '[] ('HMeasure a) -> Sing ('HMeasure a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> Sing a
typeOf abt '[] ('HMeasure a)
e)
    Statement abt Location 'ExpectP -> Expect abt ()
forall (abt :: [Hakaru] -> Hakaru -> *) (m :: * -> *)
       (p :: Purity).
EvaluationMonad abt m p =>
Statement abt Location p -> m ()
unsafePush (Location a
-> (abt '[] 'HProb -> abt '[] 'HProb)
-> [Index (abt '[])]
-> Statement abt Location 'ExpectP
forall (abt :: [Hakaru] -> Hakaru -> *) (v :: Hakaru -> *)
       (a :: Hakaru).
v a
-> (abt '[] 'HProb -> abt '[] 'HProb)
-> [Index (abt '[])]
-> Statement abt v 'ExpectP
SStuff1 (Variable a -> Location a
forall k (a :: k). Variable a -> Location a
Location Variable a
x) (\abt '[] 'HProb
c ->
        Term abt 'HProb -> abt '[] 'HProb
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
syn abt a -> abt '[] a
syn (Transform '[LC ('HMeasure a), '( '[a], 'HProb)] 'HProb
-> SCon '[LC ('HMeasure a), '( '[a], 'HProb)] 'HProb
forall (as :: [([Hakaru], Hakaru)]) (x :: Hakaru).
Transform as x -> SCon as x
AST.Transform_ Transform '[LC ('HMeasure a), '( '[a], 'HProb)] 'HProb
forall (a :: Hakaru).
Transform '[LC ('HMeasure a), '( '[a], 'HProb)] 'HProb
AST.Expect SCon '[LC ('HMeasure a), '( '[a], 'HProb)] 'HProb
-> SArgs abt '[LC ('HMeasure a), '( '[a], 'HProb)]
-> Term abt 'HProb
forall (args :: [([Hakaru], Hakaru)]) (a :: Hakaru)
       (abt :: [Hakaru] -> Hakaru -> *).
SCon args a -> SArgs abt args -> Term abt a
:$ abt '[] ('HMeasure a)
e abt '[] ('HMeasure a)
-> SArgs abt '[ '( '[a], 'HProb)]
-> SArgs abt '[LC ('HMeasure a), '( '[a], 'HProb)]
forall (abt :: [Hakaru] -> Hakaru -> *) (vars :: [Hakaru])
       (a :: Hakaru) (args :: [([Hakaru], Hakaru)]).
abt vars a -> SArgs abt args -> SArgs abt ('(vars, a) : args)
:* Variable a -> abt '[] 'HProb -> abt '[a] 'HProb
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k) (xs :: [k]) (b :: k).
ABT syn abt =>
Variable a -> abt xs b -> abt (a : xs) b
bind Variable a
x abt '[] 'HProb
c abt '[a] 'HProb -> SArgs abt '[] -> SArgs abt '[ '( '[a], 'HProb)]
forall (abt :: [Hakaru] -> Hakaru -> *) (vars :: [Hakaru])
       (a :: Hakaru) (args :: [([Hakaru], Hakaru)]).
abt vars a -> SArgs abt args -> SArgs abt ('(vars, a) : args)
:* SArgs abt '[]
forall (abt :: [Hakaru] -> Hakaru -> *). SArgs abt '[]
End)) [])
    abt '[] a -> Expect abt (abt '[] a)
forall (m :: * -> *) a. Monad m => a -> m a
return (abt '[] a -> Expect abt (abt '[] a))
-> abt '[] a -> Expect abt (abt '[] a)
forall a b. (a -> b) -> a -> b
$ Variable a -> abt '[] a
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
Variable a -> abt '[] a
var Variable a
x
{-
residualizeExpect e = do
    var <$> emit Text.empty (sUnMeasure $ typeOf e)
        (\c -> syn (AST.Expect :$ e :* c :* End))
-}

-- This version checks whether the first argument is a variable or not, avoiding the extraneous let binding as appropriate. We also avoid using 'binder', which is good because it constructs the term more directly, but is bad because we make no guarantees about hygiene! We expect callers to handle that.
-- TODO: move this elsewhere, so that 'runExpect' can use it.
-- TODO: make even smarter so that we drop the let binding in case @f@ doesn't actually use it?
let_ :: (ABT Term abt) => abt '[] a -> abt '[a] b -> abt '[] b
let_ :: abt '[] a -> abt '[a] b -> abt '[] b
let_ abt '[] a
e abt '[a] b
f =
    abt '[] a
-> (Variable a -> abt '[] b)
-> (Term abt a -> abt '[] b)
-> abt '[] b
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k) r.
ABT syn abt =>
abt '[] a -> (Variable a -> r) -> (syn abt a -> r) -> r
caseVarSyn abt '[] a
e
        (\Variable a
x -> abt '[a] b -> (Variable a -> abt '[] b -> abt '[] b) -> abt '[] b
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (x :: k) (xs :: [k]) (a :: k) r.
ABT syn abt =>
abt (x : xs) a -> (Variable x -> abt xs a -> r) -> r
caseBind abt '[a] b
f ((Variable a -> abt '[] b -> abt '[] b) -> abt '[] b)
-> (Variable a -> abt '[] b -> abt '[] b) -> abt '[] b
forall a b. (a -> b) -> a -> b
$ \Variable a
y abt '[] b
f' -> Variable a -> abt '[] a -> abt '[] b -> abt '[] b
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k) (xs :: [k]) (b :: k).
(JmEq1 Sing, Show1 Sing, Traversable21 syn, ABT syn abt) =>
Variable a -> abt '[] a -> abt xs b -> abt xs b
subst Variable a
y (Variable a -> abt '[] a
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
Variable a -> abt '[] a
var Variable a
x) abt '[] b
f')
        (\Term abt a
_ -> Term abt b -> abt '[] b
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
syn abt a -> abt '[] a
syn (SCon '[LC a, '( '[a], b)] b
forall (a :: Hakaru) (b :: Hakaru). SCon '[LC a, '( '[a], b)] b
Let_ SCon '[LC a, '( '[a], b)] b
-> SArgs abt '[LC a, '( '[a], b)] -> Term abt b
forall (args :: [([Hakaru], Hakaru)]) (a :: Hakaru)
       (abt :: [Hakaru] -> Hakaru -> *).
SCon args a -> SArgs abt args -> Term abt a
:$ abt '[] a
e abt '[] a
-> SArgs abt '[ '( '[a], b)] -> SArgs abt '[LC a, '( '[a], b)]
forall (abt :: [Hakaru] -> Hakaru -> *) (vars :: [Hakaru])
       (a :: Hakaru) (args :: [([Hakaru], Hakaru)]).
abt vars a -> SArgs abt args -> SArgs abt ('(vars, a) : args)
:* abt '[a] b
f abt '[a] b -> SArgs abt '[] -> SArgs abt '[ '( '[a], b)]
forall (abt :: [Hakaru] -> Hakaru -> *) (vars :: [Hakaru])
       (a :: Hakaru) (args :: [([Hakaru], Hakaru)]).
abt vars a -> SArgs abt args -> SArgs abt ('(vars, a) : args)
:* SArgs abt '[]
forall (abt :: [Hakaru] -> Hakaru -> *). SArgs abt '[]
End))

    
expectCase
    :: (ABT Term abt)
    => abt '[] a
    -> [Branch a abt ('HMeasure b)]
    -> Expect abt (abt '[] b)
expectCase :: abt '[] a -> [Branch a abt ('HMeasure b)] -> Expect abt (abt '[] b)
expectCase abt '[] a
scrutinee [Branch a abt ('HMeasure b)]
bs = do
    -- Get the current context and then clear it.
    ListContext abt 'ExpectP
ctx <- ((ListContext abt 'ExpectP -> ExpectAns abt) -> ExpectAns abt)
-> Expect abt (ListContext abt 'ExpectP)
forall (abt :: [Hakaru] -> Hakaru -> *) x.
((x -> ExpectAns abt) -> ExpectAns abt) -> Expect abt x
Expect (((ListContext abt 'ExpectP -> ExpectAns abt) -> ExpectAns abt)
 -> Expect abt (ListContext abt 'ExpectP))
-> ((ListContext abt 'ExpectP -> ExpectAns abt) -> ExpectAns abt)
-> Expect abt (ListContext abt 'ExpectP)
forall a b. (a -> b) -> a -> b
$ \ListContext abt 'ExpectP -> ExpectAns abt
c ListContext abt 'ExpectP
h -> ListContext abt 'ExpectP -> ExpectAns abt
c ListContext abt 'ExpectP
h (ListContext abt 'ExpectP
h {statements :: [Statement abt Location 'ExpectP]
statements = []})
    -- Emit the old "current" context.
    ((() -> ExpectAns abt) -> ExpectAns abt) -> Expect abt ()
forall (abt :: [Hakaru] -> Hakaru -> *) x.
((x -> ExpectAns abt) -> ExpectAns abt) -> Expect abt x
Expect (((() -> ExpectAns abt) -> ExpectAns abt) -> Expect abt ())
-> ((() -> ExpectAns abt) -> ExpectAns abt) -> Expect abt ()
forall a b. (a -> b) -> a -> b
$ \() -> ExpectAns abt
c ListContext abt 'ExpectP
h -> abt '[] 'HProb -> ExpectAns abt
forall (abt :: [Hakaru] -> Hakaru -> *).
ABT Term abt =>
abt '[] 'HProb -> ListContext abt 'ExpectP -> abt '[] 'HProb
residualizeExpectListContext (() -> ExpectAns abt
c () ListContext abt 'ExpectP
h) ListContext abt 'ExpectP
ctx
    -- @emitCaseWith@
    [GBranch a (Expect abt (abt '[] b))]
gms <- [Branch a abt ('HMeasure b)]
-> (Branch a abt ('HMeasure b)
    -> Expect abt (GBranch a (Expect abt (abt '[] b))))
-> Expect abt [GBranch a (Expect abt (abt '[] b))]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
t a -> (a -> f b) -> f (t b)
T.for [Branch a abt ('HMeasure b)]
bs ((Branch a abt ('HMeasure b)
  -> Expect abt (GBranch a (Expect abt (abt '[] b))))
 -> Expect abt [GBranch a (Expect abt (abt '[] b))])
-> (Branch a abt ('HMeasure b)
    -> Expect abt (GBranch a (Expect abt (abt '[] b))))
-> Expect abt [GBranch a (Expect abt (abt '[] b))]
forall a b. (a -> b) -> a -> b
$ \(Branch Pattern xs a
pat abt xs ('HMeasure b)
body) ->
        let (List1 Variable xs
vars, abt '[] ('HMeasure b)
body') = abt xs ('HMeasure b) -> (List1 Variable xs, abt '[] ('HMeasure b))
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (xs :: [k]) (a :: k).
ABT syn abt =>
abt xs a -> (List1 Variable xs, abt '[] a)
caseBinds abt xs ('HMeasure b)
body
        in  (\List1 Variable xs
vars' ->
                let rho :: Assocs (abt '[])
rho = List1 Variable xs -> List1 (abt '[]) xs -> Assocs (abt '[])
forall k (xs :: [k]) (ast :: k -> *).
List1 Variable xs -> List1 ast xs -> Assocs ast
toAssocs1 List1 Variable xs
vars ((forall (i :: Hakaru). Variable i -> abt '[] i)
-> List1 Variable xs -> List1 (abt '[]) xs
forall k1 k2 (f :: (k1 -> *) -> k2 -> *) (a :: k1 -> *)
       (b :: k1 -> *) (j :: k2).
Functor11 f =>
(forall (i :: k1). a i -> b i) -> f a j -> f b j
fmap11 forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
Variable a -> abt '[] a
forall (i :: Hakaru). Variable i -> abt '[] i
var List1 Variable xs
vars')
                in  Pattern xs a
-> List1 Variable xs
-> Expect abt (abt '[] b)
-> GBranch a (Expect abt (abt '[] b))
forall (a :: Hakaru) r (xs :: [Hakaru]).
Pattern xs a -> List1 Variable xs -> r -> GBranch a r
GBranch Pattern xs a
pat List1 Variable xs
vars' (abt '[] ('HMeasure b) -> Expect abt (abt '[] b)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] ('HMeasure a) -> Expect abt (abt '[] a)
expectTerm (abt '[] ('HMeasure b) -> Expect abt (abt '[] b))
-> abt '[] ('HMeasure b) -> Expect abt (abt '[] b)
forall a b. (a -> b) -> a -> b
$ Assocs (abt '[]) -> abt '[] ('HMeasure b) -> abt '[] ('HMeasure b)
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (xs :: [k]) (a :: k).
(ABT syn abt, JmEq1 Sing, Show1 Sing, Traversable21 syn) =>
Assocs (abt '[]) -> abt xs a -> abt xs a
substs Assocs (abt '[])
rho abt '[] ('HMeasure b)
body')
            ) (List1 Variable xs -> GBranch a (Expect abt (abt '[] b)))
-> Expect abt (List1 Variable xs)
-> Expect abt (GBranch a (Expect abt (abt '[] b)))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> List1 Variable xs -> Expect abt (List1 Variable xs)
forall (abt :: [Hakaru] -> Hakaru -> *) (m :: * -> *) (p :: Purity)
       (xs :: [Hakaru]).
EvaluationMonad abt m p =>
List1 Variable xs -> m (List1 Variable xs)
freshenVars List1 Variable xs
vars
    ((abt '[] b -> ExpectAns abt) -> ExpectAns abt)
-> Expect abt (abt '[] b)
forall (abt :: [Hakaru] -> Hakaru -> *) x.
((x -> ExpectAns abt) -> ExpectAns abt) -> Expect abt x
Expect (((abt '[] b -> ExpectAns abt) -> ExpectAns abt)
 -> Expect abt (abt '[] b))
-> ((abt '[] b -> ExpectAns abt) -> ExpectAns abt)
-> Expect abt (abt '[] b)
forall a b. (a -> b) -> a -> b
$ \abt '[] b -> ExpectAns abt
c ListContext abt 'ExpectP
h ->
        Term abt 'HProb -> abt '[] 'HProb
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
syn abt a -> abt '[] a
syn (Term abt 'HProb -> abt '[] 'HProb)
-> Term abt 'HProb -> abt '[] 'HProb
forall a b. (a -> b) -> a -> b
$ abt '[] a -> [Branch a abt 'HProb] -> Term abt 'HProb
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru)
       (b :: Hakaru).
abt '[] a -> [Branch a abt b] -> Term abt b
Case_ abt '[] a
scrutinee
            [ GBranch a (abt '[] 'HProb) -> Branch a abt 'HProb
forall (syn :: ([Hakaru] -> Hakaru -> *) -> Hakaru -> *)
       (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru) (b :: Hakaru).
ABT syn abt =>
GBranch a (abt '[] b) -> Branch a abt b
fromGBranch (GBranch a (abt '[] 'HProb) -> Branch a abt 'HProb)
-> GBranch a (abt '[] 'HProb) -> Branch a abt 'HProb
forall a b. (a -> b) -> a -> b
$ (Expect abt (abt '[] b) -> abt '[] 'HProb)
-> GBranch a (Expect abt (abt '[] b)) -> GBranch a (abt '[] 'HProb)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\Expect abt (abt '[] b)
m -> Expect abt (abt '[] b)
-> (abt '[] b -> ExpectAns abt) -> ExpectAns abt
forall (abt :: [Hakaru] -> Hakaru -> *) x.
Expect abt x -> (x -> ExpectAns abt) -> ExpectAns abt
unExpect Expect abt (abt '[] b)
m abt '[] b -> ExpectAns abt
c ListContext abt 'ExpectP
h) GBranch a (Expect abt (abt '[] b))
gm
            | GBranch a (Expect abt (abt '[] b))
gm <- [GBranch a (Expect abt (abt '[] b))]
gms
            ]

----------------------------------------------------------------
-- BUG: really rather than using 'pureEvaluate' itself, we should
-- have our own similar version which pushes the @expect _ c@ under
-- the branches; in lieu of allowing 'defaultCaseEvaluator' to
-- return a 'Neutral' term. How can we get this to work right? Seems
-- like a common problem to have since the backwards disintegrator(s)
-- have to do it too.


-- N.B., in the ICFP 2015 pearl paper, we took the expectation of
-- bound variables prior to taking the expectation of their scope.
-- E.g., @expect(let_ v e1 e2) xs c = expect e1 xs $ \x -> expect
-- e2 (insertAssoc v x xs) c@. Whereas here, I'm being lazier and
-- performing the expectation on variable lookup. This delayed
-- evaluation preserves the expectation semantics (ICFP 2015, §5.6.0)
-- whenever (1) the variable is never used (by wasted computation),
-- or (2) used exactly once (by Tonelli's theorem); so we only need
-- to worry if (3) the variable is used more than once, in which
-- case we'll have to worry about whether the various integrals
-- commute/exchange with one another. More generally, cf. Bhat et
-- al.'s \"active variables\"

-- TODO: do we want to move this to the public API of
-- "Language.Hakaru.Evaluation.DisintegrationMonad"?
#ifdef __TRACE_DISINTEGRATE__
getStatements :: Expect abt [Statement abt Location 'ExpectP]
getStatements = Expect $ \c h -> c (statements h) h
#endif


expectTerm
    :: (ABT Term abt)
    => abt '[] ('HMeasure a)
    -> Expect abt (abt '[] a)
expectTerm :: abt '[] ('HMeasure a) -> Expect abt (abt '[] a)
expectTerm abt '[] ('HMeasure a)
e = do
#ifdef __TRACE_DISINTEGRATE__
    ss <- getStatements
    trace ("\n-- expectTerm --\n"
        ++ show (pretty_Statements_withTerm ss e)
        ++ "\n") $ return ()
#endif
    Expect abt ()
forall (abt :: [Hakaru] -> Hakaru -> *).
ABT Term abt =>
Expect abt ()
emitExpectListContext
    Whnf abt ('HMeasure a)
w <- abt '[] ('HMeasure a) -> Expect abt (Whnf abt ('HMeasure a))
forall (abt :: [Hakaru] -> Hakaru -> *).
ABT Term abt =>
TermEvaluator abt (Expect abt)
pureEvaluate abt '[] ('HMeasure a)
e
    case Whnf abt ('HMeasure a)
w of
        -- TODO: if the neutral term is a 'Case_' then we want to go under it
        Neutral abt '[] ('HMeasure a)
e'              ->
            abt '[] ('HMeasure a)
-> (Variable ('HMeasure a) -> Expect abt (abt '[] a))
-> (Term abt ('HMeasure a) -> Expect abt (abt '[] a))
-> Expect abt (abt '[] a)
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k) r.
ABT syn abt =>
abt '[] a -> (Variable a -> r) -> (syn abt a -> r) -> r
caseVarSyn abt '[] ('HMeasure a)
e' (abt '[] ('HMeasure a) -> Expect abt (abt '[] a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] ('HMeasure a) -> Expect abt (abt '[] a)
residualizeExpect (abt '[] ('HMeasure a) -> Expect abt (abt '[] a))
-> (Variable ('HMeasure a) -> abt '[] ('HMeasure a))
-> Variable ('HMeasure a)
-> Expect abt (abt '[] a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Variable ('HMeasure a) -> abt '[] ('HMeasure a)
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
Variable a -> abt '[] a
var) ((Term abt ('HMeasure a) -> Expect abt (abt '[] a))
 -> Expect abt (abt '[] a))
-> (Term abt ('HMeasure a) -> Expect abt (abt '[] a))
-> Expect abt (abt '[] a)
forall a b. (a -> b) -> a -> b
$ \Term abt ('HMeasure a)
t ->
                case Term abt ('HMeasure a)
t of
                Case_ abt '[] a
e1 [Branch a abt ('HMeasure a)]
bs -> abt '[] a -> [Branch a abt ('HMeasure a)] -> Expect abt (abt '[] a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru)
       (b :: Hakaru).
ABT Term abt =>
abt '[] a -> [Branch a abt ('HMeasure b)] -> Expect abt (abt '[] b)
expectCase abt '[] a
e1 [Branch a abt ('HMeasure a)]
bs
                Term abt ('HMeasure a)
_           -> abt '[] ('HMeasure a) -> Expect abt (abt '[] a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] ('HMeasure a) -> Expect abt (abt '[] a)
residualizeExpect abt '[] ('HMeasure a)
e'
        Head_ (WLiteral    Literal ('HMeasure a)
_)    -> [Char] -> Expect abt (abt '[] a)
forall a. HasCallStack => [Char] -> a
error [Char]
"expect: the impossible happened"
        Head_ (WCoerceTo   Coercion a ('HMeasure a)
_  Head abt a
_) -> [Char] -> Expect abt (abt '[] a)
forall a. HasCallStack => [Char] -> a
error [Char]
"expect: the impossible happened"
        Head_ (WUnsafeFrom Coercion ('HMeasure a) b
_  Head abt b
_) -> [Char] -> Expect abt (abt '[] a)
forall a. HasCallStack => [Char] -> a
error [Char]
"expect: the impossible happened"
        Head_ (WMeasureOp MeasureOp typs a
o SArgs abt args
es) -> MeasureOp typs a -> SArgs abt args -> Expect abt (abt '[] a)
forall (abt :: [Hakaru] -> Hakaru -> *) (typs :: [Hakaru])
       (args :: [([Hakaru], Hakaru)]) (a :: Hakaru).
(ABT Term abt, typs ~ UnLCs args, args ~ LCs typs) =>
MeasureOp typs a -> SArgs abt args -> Expect abt (abt '[] a)
expectMeasureOp MeasureOp typs a
o SArgs abt args
es
        Head_ (WDirac abt '[] a
e1)       -> abt '[] a -> Expect abt (abt '[] a)
forall (m :: * -> *) a. Monad m => a -> m a
return abt '[] a
e1
        Head_ (WMBind abt '[] ('HMeasure a)
e1 abt '[a] ('HMeasure b)
e2)    -> do
            abt '[] a
v1 <- abt '[] ('HMeasure a) -> Expect abt (abt '[] a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] ('HMeasure a) -> Expect abt (abt '[] a)
expectTerm abt '[] ('HMeasure a)
e1
            abt '[] ('HMeasure b) -> Expect abt (abt '[] b)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] ('HMeasure a) -> Expect abt (abt '[] a)
expectTerm (abt '[] a -> abt '[a] ('HMeasure b) -> abt '[] ('HMeasure b)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru)
       (b :: Hakaru).
ABT Term abt =>
abt '[] a -> abt '[a] b -> abt '[] b
let_ abt '[] a
v1 abt '[a] ('HMeasure b)
e2)
        Head_ (WPlate abt '[] 'HNat
_ abt '[ 'HNat] ('HMeasure a)
_)     -> [Char] -> Expect abt (abt '[] a)
forall a. HasCallStack => [Char] -> a
error [Char]
"TODO: expect{Plate}"
        Head_ (WChain abt '[] 'HNat
_ abt '[] s
_ abt '[s] ('HMeasure (HPair a s))
_)   -> [Char] -> Expect abt (abt '[] a)
forall a. HasCallStack => [Char] -> a
error [Char]
"TODO: expect{Chain}"
        Head_ (WReject    Sing ('HMeasure a)
_)   -> [(abt '[] 'HProb, abt '[] ('HMeasure a))] -> Expect abt (abt '[] a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
[(abt '[] 'HProb, abt '[] ('HMeasure a))] -> Expect abt (abt '[] a)
expectSuperpose []
        Head_ (WSuperpose NonEmpty (abt '[] 'HProb, abt '[] ('HMeasure a))
pes) -> [(abt '[] 'HProb, abt '[] ('HMeasure a))] -> Expect abt (abt '[] a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
[(abt '[] 'HProb, abt '[] ('HMeasure a))] -> Expect abt (abt '[] a)
expectSuperpose (NonEmpty (abt '[] 'HProb, abt '[] ('HMeasure a))
-> [(abt '[] 'HProb, abt '[] ('HMeasure a))]
forall a. NonEmpty a -> [a]
NE.toList NonEmpty (abt '[] 'HProb, abt '[] ('HMeasure a))
pes)


-- N.B., we guarantee that each @e@ is called with the same heap
-- @h@ and continuation @c@.
expectSuperpose
    :: (ABT Term abt)
    => [(abt '[] 'HProb, abt '[] ('HMeasure a))]
    -> Expect abt (abt '[] a)
expectSuperpose :: [(abt '[] 'HProb, abt '[] ('HMeasure a))] -> Expect abt (abt '[] a)
expectSuperpose [(abt '[] 'HProb, abt '[] ('HMeasure a))]
pes = do
#ifdef __TRACE_DISINTEGRATE__
    ss <- getStatements
    trace ("\n-- expectSuperpose --\n"
        ++ show (pretty_Statements_withTerm ss (syn $ Superpose_ (NE.fromList pes)))
        ++ "\n") $ return ()
#endif
    -- First, emit the current heap (so that each @p@ is emissible)
    Expect abt ()
forall (abt :: [Hakaru] -> Hakaru -> *).
ABT Term abt =>
Expect abt ()
emitExpectListContext
    -- Then emit the 'sum', and call the same continuation on each @e@
    ((abt '[] a -> ExpectAns abt) -> ExpectAns abt)
-> Expect abt (abt '[] a)
forall (abt :: [Hakaru] -> Hakaru -> *) x.
((x -> ExpectAns abt) -> ExpectAns abt) -> Expect abt x
Expect (((abt '[] a -> ExpectAns abt) -> ExpectAns abt)
 -> Expect abt (abt '[] a))
-> ((abt '[] a -> ExpectAns abt) -> ExpectAns abt)
-> Expect abt (abt '[] a)
forall a b. (a -> b) -> a -> b
$ \abt '[] a -> ExpectAns abt
c ListContext abt 'ExpectP
h ->
        [abt '[] 'HProb] -> abt '[] 'HProb
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
(ABT Term abt, HSemiring_ a) =>
[abt '[] a] -> abt '[] a
P.sum [ abt '[] 'HProb
p abt '[] 'HProb -> abt '[] 'HProb -> abt '[] 'HProb
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
(ABT Term abt, HSemiring_ a) =>
abt '[] a -> abt '[] a -> abt '[] a
P.* Expect abt (abt '[] a)
-> (abt '[] a -> ExpectAns abt) -> ExpectAns abt
forall (abt :: [Hakaru] -> Hakaru -> *) x.
Expect abt x -> (x -> ExpectAns abt) -> ExpectAns abt
unExpect (abt '[] ('HMeasure a) -> Expect abt (abt '[] a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] ('HMeasure a) -> Expect abt (abt '[] a)
expectTerm abt '[] ('HMeasure a)
e) abt '[] a -> ExpectAns abt
c ListContext abt 'ExpectP
h | (abt '[] 'HProb
p,abt '[] ('HMeasure a)
e) <- [(abt '[] 'HProb, abt '[] ('HMeasure a))]
pes]
    -- TODO: if @pes@ is null, then automatically simplify to just 0
    -- TODO: in the Lazy.tex paper, we guard against that interpretation being negative...

emitExpectListContext :: forall abt. (ABT Term abt) => Expect abt ()
emitExpectListContext :: Expect abt ()
emitExpectListContext = do
    [Statement abt Location 'ExpectP]
ss <- (([Statement abt Location 'ExpectP] -> ExpectAns abt)
 -> ExpectAns abt)
-> Expect abt [Statement abt Location 'ExpectP]
forall (abt :: [Hakaru] -> Hakaru -> *) x.
((x -> ExpectAns abt) -> ExpectAns abt) -> Expect abt x
Expect ((([Statement abt Location 'ExpectP] -> ExpectAns abt)
  -> ExpectAns abt)
 -> Expect abt [Statement abt Location 'ExpectP])
-> (([Statement abt Location 'ExpectP] -> ExpectAns abt)
    -> ExpectAns abt)
-> Expect abt [Statement abt Location 'ExpectP]
forall a b. (a -> b) -> a -> b
$ \[Statement abt Location 'ExpectP] -> ExpectAns abt
c ListContext abt 'ExpectP
h -> [Statement abt Location 'ExpectP] -> ExpectAns abt
c (ListContext abt 'ExpectP -> [Statement abt Location 'ExpectP]
forall (abt :: [Hakaru] -> Hakaru -> *) (p :: Purity).
ListContext abt p -> [Statement abt Location p]
statements ListContext abt 'ExpectP
h) (ListContext abt 'ExpectP
h {statements :: [Statement abt Location 'ExpectP]
statements = []})
    (Statement abt Location 'ExpectP -> Expect abt ())
-> [Statement abt Location 'ExpectP] -> Expect abt ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
F.traverse_ Statement abt Location 'ExpectP -> Expect abt ()
step ([Statement abt Location 'ExpectP]
-> [Statement abt Location 'ExpectP]
forall a. [a] -> [a]
reverse [Statement abt Location 'ExpectP]
ss) -- TODO: use composition tricks to avoid reversing @ss@
    where
    step :: Statement abt Location 'ExpectP -> Expect abt ()
    step :: Statement abt Location 'ExpectP -> Expect abt ()
step Statement abt Location 'ExpectP
s =
#ifdef __TRACE_DISINTEGRATE__
        trace ("\n-- emitExpectListContext: " ++ show (ppStatement 0 s)) $
#endif
        case Statement abt Location 'ExpectP
s of
        SLet Location a
l Lazy abt a
body [Index (abt '[])]
_ ->
            -- TODO: be smart about dropping unused let-bindings and inlining trivial let-bindings
            ((() -> ExpectAns abt) -> ExpectAns abt) -> Expect abt ()
forall (abt :: [Hakaru] -> Hakaru -> *) x.
((x -> ExpectAns abt) -> ExpectAns abt) -> Expect abt x
Expect (((() -> ExpectAns abt) -> ExpectAns abt) -> Expect abt ())
-> ((() -> ExpectAns abt) -> ExpectAns abt) -> Expect abt ()
forall a b. (a -> b) -> a -> b
$ \() -> ExpectAns abt
c ListContext abt 'ExpectP
h ->
                Term abt 'HProb -> abt '[] 'HProb
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
syn abt a -> abt '[] a
syn (SCon '[LC a, '( '[a], 'HProb)] 'HProb
forall (a :: Hakaru) (b :: Hakaru). SCon '[LC a, '( '[a], b)] b
Let_ SCon '[LC a, '( '[a], 'HProb)] 'HProb
-> SArgs abt '[LC a, '( '[a], 'HProb)] -> Term abt 'HProb
forall (args :: [([Hakaru], Hakaru)]) (a :: Hakaru)
       (abt :: [Hakaru] -> Hakaru -> *).
SCon args a -> SArgs abt args -> Term abt a
:$ Lazy abt a -> abt '[] a
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
Lazy abt a -> abt '[] a
fromLazy Lazy abt a
body abt '[] a
-> SArgs abt '[ '( '[a], 'HProb)]
-> SArgs abt '[LC a, '( '[a], 'HProb)]
forall (abt :: [Hakaru] -> Hakaru -> *) (vars :: [Hakaru])
       (a :: Hakaru) (args :: [([Hakaru], Hakaru)]).
abt vars a -> SArgs abt args -> SArgs abt ('(vars, a) : args)
:* Variable a -> abt '[] 'HProb -> abt '[a] 'HProb
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k) (xs :: [k]) (b :: k).
ABT syn abt =>
Variable a -> abt xs b -> abt (a : xs) b
bind (Location a -> Variable a
forall k (a :: k). Location a -> Variable a
fromLocation Location a
l) (() -> ExpectAns abt
c () ListContext abt 'ExpectP
h) abt '[a] 'HProb -> SArgs abt '[] -> SArgs abt '[ '( '[a], 'HProb)]
forall (abt :: [Hakaru] -> Hakaru -> *) (vars :: [Hakaru])
       (a :: Hakaru) (args :: [([Hakaru], Hakaru)]).
abt vars a -> SArgs abt args -> SArgs abt ('(vars, a) : args)
:* SArgs abt '[]
forall (abt :: [Hakaru] -> Hakaru -> *). SArgs abt '[]
End)
        SStuff0   abt '[] 'HProb -> abt '[] 'HProb
f [Index (abt '[])]
_ -> ((() -> ExpectAns abt) -> ExpectAns abt) -> Expect abt ()
forall (abt :: [Hakaru] -> Hakaru -> *) x.
((x -> ExpectAns abt) -> ExpectAns abt) -> Expect abt x
Expect (((() -> ExpectAns abt) -> ExpectAns abt) -> Expect abt ())
-> ((() -> ExpectAns abt) -> ExpectAns abt) -> Expect abt ()
forall a b. (a -> b) -> a -> b
$ \() -> ExpectAns abt
c ListContext abt 'ExpectP
h -> abt '[] 'HProb -> abt '[] 'HProb
f (() -> ExpectAns abt
c () ListContext abt 'ExpectP
h)
        SStuff1 Location a
_ abt '[] 'HProb -> abt '[] 'HProb
f [Index (abt '[])]
_ -> ((() -> ExpectAns abt) -> ExpectAns abt) -> Expect abt ()
forall (abt :: [Hakaru] -> Hakaru -> *) x.
((x -> ExpectAns abt) -> ExpectAns abt) -> Expect abt x
Expect (((() -> ExpectAns abt) -> ExpectAns abt) -> Expect abt ())
-> ((() -> ExpectAns abt) -> ExpectAns abt) -> Expect abt ()
forall a b. (a -> b) -> a -> b
$ \() -> ExpectAns abt
c ListContext abt 'ExpectP
h -> abt '[] 'HProb -> abt '[] 'HProb
f (() -> ExpectAns abt
c () ListContext abt 'ExpectP
h)


pushIntegrate
    :: (ABT Term abt)
    => abt '[] 'HReal
    -> abt '[] 'HReal
    -> Expect abt (Variable 'HReal)
pushIntegrate :: abt '[] 'HReal -> abt '[] 'HReal -> Expect abt (Variable 'HReal)
pushIntegrate abt '[] 'HReal
lo abt '[] 'HReal
hi = do
    Variable 'HReal
x <- Text -> Sing 'HReal -> Expect abt (Variable 'HReal)
forall (abt :: [Hakaru] -> Hakaru -> *) (m :: * -> *) (p :: Purity)
       (a :: Hakaru).
EvaluationMonad abt m p =>
Text -> Sing a -> m (Variable a)
freshVar Text
Text.empty Sing 'HReal
SReal
    Statement abt Location 'ExpectP -> Expect abt ()
forall (abt :: [Hakaru] -> Hakaru -> *) (m :: * -> *)
       (p :: Purity).
EvaluationMonad abt m p =>
Statement abt Location p -> m ()
unsafePush (Location 'HReal
-> (abt '[] 'HProb -> abt '[] 'HProb)
-> [Index (abt '[])]
-> Statement abt Location 'ExpectP
forall (abt :: [Hakaru] -> Hakaru -> *) (v :: Hakaru -> *)
       (a :: Hakaru).
v a
-> (abt '[] 'HProb -> abt '[] 'HProb)
-> [Index (abt '[])]
-> Statement abt v 'ExpectP
SStuff1 (Variable 'HReal -> Location 'HReal
forall k (a :: k). Variable a -> Location a
Location Variable 'HReal
x) (\abt '[] 'HProb
c ->
        Term abt 'HProb -> abt '[] 'HProb
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
syn abt a -> abt '[] a
syn (SCon '[LC 'HReal, LC 'HReal, '( '[ 'HReal], 'HProb)] 'HProb
Integrate SCon '[LC 'HReal, LC 'HReal, '( '[ 'HReal], 'HProb)] 'HProb
-> SArgs abt '[LC 'HReal, LC 'HReal, '( '[ 'HReal], 'HProb)]
-> Term abt 'HProb
forall (args :: [([Hakaru], Hakaru)]) (a :: Hakaru)
       (abt :: [Hakaru] -> Hakaru -> *).
SCon args a -> SArgs abt args -> Term abt a
:$ abt '[] 'HReal
lo abt '[] 'HReal
-> SArgs abt '[LC 'HReal, '( '[ 'HReal], 'HProb)]
-> SArgs abt '[LC 'HReal, LC 'HReal, '( '[ 'HReal], 'HProb)]
forall (abt :: [Hakaru] -> Hakaru -> *) (vars :: [Hakaru])
       (a :: Hakaru) (args :: [([Hakaru], Hakaru)]).
abt vars a -> SArgs abt args -> SArgs abt ('(vars, a) : args)
:* abt '[] 'HReal
hi abt '[] 'HReal
-> SArgs abt '[ '( '[ 'HReal], 'HProb)]
-> SArgs abt '[LC 'HReal, '( '[ 'HReal], 'HProb)]
forall (abt :: [Hakaru] -> Hakaru -> *) (vars :: [Hakaru])
       (a :: Hakaru) (args :: [([Hakaru], Hakaru)]).
abt vars a -> SArgs abt args -> SArgs abt ('(vars, a) : args)
:* Variable 'HReal -> abt '[] 'HProb -> abt '[ 'HReal] 'HProb
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k) (xs :: [k]) (b :: k).
ABT syn abt =>
Variable a -> abt xs b -> abt (a : xs) b
bind Variable 'HReal
x abt '[] 'HProb
c abt '[ 'HReal] 'HProb
-> SArgs abt '[] -> SArgs abt '[ '( '[ 'HReal], 'HProb)]
forall (abt :: [Hakaru] -> Hakaru -> *) (vars :: [Hakaru])
       (a :: Hakaru) (args :: [([Hakaru], Hakaru)]).
abt vars a -> SArgs abt args -> SArgs abt ('(vars, a) : args)
:* SArgs abt '[]
forall (abt :: [Hakaru] -> Hakaru -> *). SArgs abt '[]
End)) [])
    Variable 'HReal -> Expect abt (Variable 'HReal)
forall (m :: * -> *) a. Monad m => a -> m a
return Variable 'HReal
x
{-
-- BUG: we assume the arguments are emissible!
emitIntegrate lo hi =
    emit Text.empty SReal (\c ->
        syn (Integrate :$ lo :* hi :* c :* End))
-}

-- Needs to be more polymorphic
pushSummate
    :: (ABT Term abt, HDiscrete_ a, SingI a)
    => abt '[] a
    -> abt '[] a
    -> Expect abt (Variable a)
pushSummate :: abt '[] a -> abt '[] a -> Expect abt (Variable a)
pushSummate abt '[] a
lo abt '[] a
hi = do
    Variable a
x <- Text -> Sing a -> Expect abt (Variable a)
forall (abt :: [Hakaru] -> Hakaru -> *) (m :: * -> *) (p :: Purity)
       (a :: Hakaru).
EvaluationMonad abt m p =>
Text -> Sing a -> m (Variable a)
freshVar Text
Text.empty Sing a
forall k (a :: k). SingI a => Sing a
sing
    Statement abt Location 'ExpectP -> Expect abt ()
forall (abt :: [Hakaru] -> Hakaru -> *) (m :: * -> *)
       (p :: Purity).
EvaluationMonad abt m p =>
Statement abt Location p -> m ()
unsafePush (Location a
-> (abt '[] 'HProb -> abt '[] 'HProb)
-> [Index (abt '[])]
-> Statement abt Location 'ExpectP
forall (abt :: [Hakaru] -> Hakaru -> *) (v :: Hakaru -> *)
       (a :: Hakaru).
v a
-> (abt '[] 'HProb -> abt '[] 'HProb)
-> [Index (abt '[])]
-> Statement abt v 'ExpectP
SStuff1 (Variable a -> Location a
forall k (a :: k). Variable a -> Location a
Location Variable a
x) (\abt '[] 'HProb
c ->
        Term abt 'HProb -> abt '[] 'HProb
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
syn abt a -> abt '[] a
syn (HDiscrete a
-> HSemiring 'HProb -> SCon '[LC a, LC a, '( '[a], 'HProb)] 'HProb
forall (a :: Hakaru) (b :: Hakaru).
HDiscrete a -> HSemiring b -> SCon '[LC a, LC a, '( '[a], b)] b
Summate HDiscrete a
forall (a :: Hakaru). HDiscrete_ a => HDiscrete a
hDiscrete HSemiring 'HProb
forall (a :: Hakaru). HSemiring_ a => HSemiring a
hSemiring
             SCon '[LC a, LC a, '( '[a], 'HProb)] 'HProb
-> SArgs abt '[LC a, LC a, '( '[a], 'HProb)] -> Term abt 'HProb
forall (args :: [([Hakaru], Hakaru)]) (a :: Hakaru)
       (abt :: [Hakaru] -> Hakaru -> *).
SCon args a -> SArgs abt args -> Term abt a
:$ abt '[] a
lo abt '[] a
-> SArgs abt '[LC a, '( '[a], 'HProb)]
-> SArgs abt '[LC a, LC a, '( '[a], 'HProb)]
forall (abt :: [Hakaru] -> Hakaru -> *) (vars :: [Hakaru])
       (a :: Hakaru) (args :: [([Hakaru], Hakaru)]).
abt vars a -> SArgs abt args -> SArgs abt ('(vars, a) : args)
:* abt '[] a
hi abt '[] a
-> SArgs abt '[ '( '[a], 'HProb)]
-> SArgs abt '[LC a, '( '[a], 'HProb)]
forall (abt :: [Hakaru] -> Hakaru -> *) (vars :: [Hakaru])
       (a :: Hakaru) (args :: [([Hakaru], Hakaru)]).
abt vars a -> SArgs abt args -> SArgs abt ('(vars, a) : args)
:* Variable a -> abt '[] 'HProb -> abt '[a] 'HProb
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k) (xs :: [k]) (b :: k).
ABT syn abt =>
Variable a -> abt xs b -> abt (a : xs) b
bind Variable a
x abt '[] 'HProb
c abt '[a] 'HProb -> SArgs abt '[] -> SArgs abt '[ '( '[a], 'HProb)]
forall (abt :: [Hakaru] -> Hakaru -> *) (vars :: [Hakaru])
       (a :: Hakaru) (args :: [([Hakaru], Hakaru)]).
abt vars a -> SArgs abt args -> SArgs abt ('(vars, a) : args)
:* SArgs abt '[]
forall (abt :: [Hakaru] -> Hakaru -> *). SArgs abt '[]
End)) [])
    Variable a -> Expect abt (Variable a)
forall (m :: * -> *) a. Monad m => a -> m a
return Variable a
x
{-
-- BUG: we assume the arguments are emissible!
emitSummate lo hi =
    emit Text.empty SInt (\c ->
        syn (Summate :$ lo :* hi :* c :* End))
-}

-- TODO: can we / would it help to, reuse 'let_'?
-- BUG: we assume the argument is emissible!
pushLet :: (ABT Term abt) => abt '[] a -> Expect abt (Variable a)
pushLet :: abt '[] a -> Expect abt (Variable a)
pushLet abt '[] a
e =
    abt '[] a
-> (Variable a -> Expect abt (Variable a))
-> (Term abt a -> Expect abt (Variable a))
-> Expect abt (Variable a)
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k) r.
ABT syn abt =>
abt '[] a -> (Variable a -> r) -> (syn abt a -> r) -> r
caseVarSyn abt '[] a
e Variable a -> Expect abt (Variable a)
forall (m :: * -> *) a. Monad m => a -> m a
return ((Term abt a -> Expect abt (Variable a))
 -> Expect abt (Variable a))
-> (Term abt a -> Expect abt (Variable a))
-> Expect abt (Variable a)
forall a b. (a -> b) -> a -> b
$ \Term abt a
_ -> do
        Variable a
x <- Text -> Sing a -> Expect abt (Variable a)
forall (abt :: [Hakaru] -> Hakaru -> *) (m :: * -> *) (p :: Purity)
       (a :: Hakaru).
EvaluationMonad abt m p =>
Text -> Sing a -> m (Variable a)
freshVar Text
Text.empty (abt '[] a -> Sing a
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> Sing a
typeOf abt '[] a
e)
        Statement abt Location 'ExpectP -> Expect abt ()
forall (abt :: [Hakaru] -> Hakaru -> *) (m :: * -> *)
       (p :: Purity).
EvaluationMonad abt m p =>
Statement abt Location p -> m ()
unsafePush (Location a
-> (abt '[] 'HProb -> abt '[] 'HProb)
-> [Index (abt '[])]
-> Statement abt Location 'ExpectP
forall (abt :: [Hakaru] -> Hakaru -> *) (v :: Hakaru -> *)
       (a :: Hakaru).
v a
-> (abt '[] 'HProb -> abt '[] 'HProb)
-> [Index (abt '[])]
-> Statement abt v 'ExpectP
SStuff1 (Variable a -> Location a
forall k (a :: k). Variable a -> Location a
Location Variable a
x) (\abt '[] 'HProb
c ->
            Term abt 'HProb -> abt '[] 'HProb
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
syn abt a -> abt '[] a
syn (SCon '[LC a, '( '[a], 'HProb)] 'HProb
forall (a :: Hakaru) (b :: Hakaru). SCon '[LC a, '( '[a], b)] b
Let_ SCon '[LC a, '( '[a], 'HProb)] 'HProb
-> SArgs abt '[LC a, '( '[a], 'HProb)] -> Term abt 'HProb
forall (args :: [([Hakaru], Hakaru)]) (a :: Hakaru)
       (abt :: [Hakaru] -> Hakaru -> *).
SCon args a -> SArgs abt args -> Term abt a
:$ abt '[] a
e abt '[] a
-> SArgs abt '[ '( '[a], 'HProb)]
-> SArgs abt '[LC a, '( '[a], 'HProb)]
forall (abt :: [Hakaru] -> Hakaru -> *) (vars :: [Hakaru])
       (a :: Hakaru) (args :: [([Hakaru], Hakaru)]).
abt vars a -> SArgs abt args -> SArgs abt ('(vars, a) : args)
:* Variable a -> abt '[] 'HProb -> abt '[a] 'HProb
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k) (xs :: [k]) (b :: k).
ABT syn abt =>
Variable a -> abt xs b -> abt (a : xs) b
bind Variable a
x abt '[] 'HProb
c abt '[a] 'HProb -> SArgs abt '[] -> SArgs abt '[ '( '[a], 'HProb)]
forall (abt :: [Hakaru] -> Hakaru -> *) (vars :: [Hakaru])
       (a :: Hakaru) (args :: [([Hakaru], Hakaru)]).
abt vars a -> SArgs abt args -> SArgs abt ('(vars, a) : args)
:* SArgs abt '[]
forall (abt :: [Hakaru] -> Hakaru -> *). SArgs abt '[]
End)) [])
        Variable a -> Expect abt (Variable a)
forall (m :: * -> *) a. Monad m => a -> m a
return Variable a
x
{-
emitLet e =
    caseVarSyn e return $ \_ ->
        emit Text.empty (typeOf e) $ \f ->
            syn (Let_ :$ e :* f :* End)
-}


-- TODO: introduce HProb variants of integrate\/summate so we can avoid the need for 'unsafeProb' here
expectMeasureOp
    :: forall abt typs args a
    .  (ABT Term abt, typs ~ UnLCs args, args ~ LCs typs)
    => MeasureOp typs a
    -> SArgs abt args
    -> Expect abt (abt '[] a)
expectMeasureOp :: MeasureOp typs a -> SArgs abt args -> Expect abt (abt '[] a)
expectMeasureOp MeasureOp typs a
Lebesgue = \(abt vars a
lo :* abt vars a
hi :* SArgs abt args
End) -> do 
    abt '[] a
lo' <- Variable a -> abt '[] a
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
Variable a -> abt '[] a
var (Variable a -> abt '[] a)
-> Expect abt (Variable a) -> Expect abt (abt '[] a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> abt '[] a -> Expect abt (Variable a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> Expect abt (Variable a)
pushLet abt vars a
abt '[] a
lo
    abt '[] a
hi' <- Variable a -> abt '[] a
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
Variable a -> abt '[] a
var (Variable a -> abt '[] a)
-> Expect abt (Variable a) -> Expect abt (abt '[] a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> abt '[] a -> Expect abt (Variable a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> Expect abt (Variable a)
pushLet abt vars a
abt '[] a
hi
    Variable 'HReal -> abt '[] 'HReal
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
Variable a -> abt '[] a
var (Variable 'HReal -> abt '[] 'HReal)
-> Expect abt (Variable 'HReal) -> Expect abt (abt '[] 'HReal)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> abt '[] 'HReal -> abt '[] 'HReal -> Expect abt (Variable 'HReal)
forall (abt :: [Hakaru] -> Hakaru -> *).
ABT Term abt =>
abt '[] 'HReal -> abt '[] 'HReal -> Expect abt (Variable 'HReal)
pushIntegrate abt '[] a
abt '[] 'HReal
lo' abt '[] a
abt '[] 'HReal
hi' 
expectMeasureOp MeasureOp typs a
Counting = \SArgs abt args
End ->
    Variable a -> abt '[] a
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
Variable a -> abt '[] a
var (Variable a -> abt '[] a)
-> Expect abt (Variable a) -> Expect abt (abt '[] a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> abt '[] a -> abt '[] a -> Expect abt (Variable a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
(ABT Term abt, HDiscrete_ a, SingI a) =>
abt '[] a -> abt '[] a -> Expect abt (Variable a)
pushSummate abt '[] a
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
(ABT Term abt, HRing_ a, Integrable a) =>
abt '[] a
P.negativeInfinity abt '[] a
forall (a :: Hakaru) (abt :: [Hakaru] -> Hakaru -> *).
(Integrable a, ABT Term abt) =>
abt '[] a
P.infinity
expectMeasureOp MeasureOp typs a
Categorical = \(abt vars a
ps :* SArgs abt args
End) -> do
    abt '[] a
ps' <- Variable a -> abt '[] a
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
Variable a -> abt '[] a
var (Variable a -> abt '[] a)
-> Expect abt (Variable a) -> Expect abt (abt '[] a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> abt '[] a -> Expect abt (Variable a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> Expect abt (Variable a)
pushLet abt vars a
abt '[] a
ps
    abt '[] 'HProb
tot <- Variable 'HProb -> abt '[] 'HProb
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
Variable a -> abt '[] a
var (Variable 'HProb -> abt '[] 'HProb)
-> Expect abt (Variable 'HProb) -> Expect abt (abt '[] 'HProb)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> abt '[] 'HProb -> Expect abt (Variable 'HProb)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> Expect abt (Variable a)
pushLet (abt '[] ('HArray 'HProb) -> abt '[] 'HProb
forall (abt :: [Hakaru] -> Hakaru -> *).
ABT Term abt =>
abt '[] ('HArray 'HProb) -> abt '[] 'HProb
P.summateV abt '[] a
abt '[] ('HArray 'HProb)
ps')
    Statement abt Location 'ExpectP -> Expect abt ()
forall (abt :: [Hakaru] -> Hakaru -> *) (m :: * -> *)
       (p :: Purity).
EvaluationMonad abt m p =>
Statement abt Location p -> m ()
unsafePush ((abt '[] 'HProb -> abt '[] 'HProb)
-> [Index (abt '[])] -> Statement abt Location 'ExpectP
forall (abt :: [Hakaru] -> Hakaru -> *) (v :: Hakaru -> *).
(abt '[] 'HProb -> abt '[] 'HProb)
-> [Index (abt '[])] -> Statement abt v 'ExpectP
SStuff0 (\abt '[] 'HProb
c -> abt '[] HBool -> abt '[] 'HProb -> abt '[] 'HProb -> abt '[] 'HProb
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] HBool -> abt '[] a -> abt '[] a -> abt '[] a
P.if_ (abt '[] 'HProb
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
(ABT Term abt, HSemiring_ a) =>
abt '[] a
P.zero abt '[] 'HProb -> abt '[] 'HProb -> abt '[] HBool
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
(ABT Term abt, HOrd_ a) =>
abt '[] a -> abt '[] a -> abt '[] HBool
P.< abt '[] 'HProb
tot) abt '[] 'HProb
c abt '[] 'HProb
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
(ABT Term abt, HSemiring_ a) =>
abt '[] a
P.zero) [])
    Variable 'HNat
i <- Text -> Sing 'HNat -> Expect abt (Variable 'HNat)
forall (abt :: [Hakaru] -> Hakaru -> *) (m :: * -> *) (p :: Purity)
       (a :: Hakaru).
EvaluationMonad abt m p =>
Text -> Sing a -> m (Variable a)
freshVar Text
Text.empty Sing 'HNat
SNat
    ((abt '[] 'HNat -> ExpectAns abt) -> ExpectAns abt)
-> Expect abt (abt '[] 'HNat)
forall (abt :: [Hakaru] -> Hakaru -> *) x.
((x -> ExpectAns abt) -> ExpectAns abt) -> Expect abt x
Expect (((abt '[] 'HNat -> ExpectAns abt) -> ExpectAns abt)
 -> Expect abt (abt '[] 'HNat))
-> ((abt '[] 'HNat -> ExpectAns abt) -> ExpectAns abt)
-> Expect abt (abt '[] 'HNat)
forall a b. (a -> b) -> a -> b
$ \abt '[] 'HNat -> ExpectAns abt
c ListContext abt 'ExpectP
h ->
        abt '[] ('HArray 'HProb) -> abt '[] 'HProb
forall (abt :: [Hakaru] -> Hakaru -> *).
ABT Term abt =>
abt '[] ('HArray 'HProb) -> abt '[] 'HProb
P.summateV
            (Term abt ('HArray 'HProb) -> abt '[] ('HArray 'HProb)
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
syn abt a -> abt '[] a
syn (abt '[] 'HNat -> abt '[ 'HNat] 'HProb -> Term abt ('HArray 'HProb)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
abt '[] 'HNat -> abt '[ 'HNat] a -> Term abt ('HArray a)
Array_ (abt '[] ('HArray 'HProb) -> abt '[] 'HNat
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] ('HArray a) -> abt '[] 'HNat
P.size abt '[] a
abt '[] ('HArray 'HProb)
ps') (Variable 'HNat -> abt '[] 'HProb -> abt '[ 'HNat] 'HProb
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k) (xs :: [k]) (b :: k).
ABT syn abt =>
Variable a -> abt xs b -> abt (a : xs) b
bind Variable 'HNat
i ((abt '[] a
abt '[] ('HArray 'HProb)
ps' abt '[] ('HArray 'HProb) -> abt '[] 'HNat -> abt '[] 'HProb
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] ('HArray a) -> abt '[] 'HNat -> abt '[] a
P.! Variable 'HNat -> abt '[] 'HNat
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
Variable a -> abt '[] a
var Variable 'HNat
i) abt '[] 'HProb -> abt '[] 'HProb -> abt '[] 'HProb
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
(ABT Term abt, HSemiring_ a) =>
abt '[] a -> abt '[] a -> abt '[] a
P.* abt '[] 'HNat -> ExpectAns abt
c (Variable 'HNat -> abt '[] 'HNat
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
Variable a -> abt '[] a
var Variable 'HNat
i) ListContext abt 'ExpectP
h))))
            abt '[] 'HProb -> abt '[] 'HProb -> abt '[] 'HProb
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
(ABT Term abt, HFractional_ a) =>
abt '[] a -> abt '[] a -> abt '[] a
P./ abt '[] 'HProb
tot
    {-
    let_ ps $ \ps' ->
    let_ (summateV ps') $ \tot ->
    if_ (zero < tot)
        (summateV (mapWithIndex (\i p -> p * inst c i) ps') / tot)
        zero
    -}
expectMeasureOp MeasureOp typs a
Uniform = \(abt vars a
lo :* abt vars a
hi :* SArgs abt args
End) -> do
    -- BUG: @(let_ zero $ \y -> uniform y one)@ doesn't work as desired; *drops* the @SLet y zero@ binding entirely!!
    abt '[] a
lo' <- Variable a -> abt '[] a
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
Variable a -> abt '[] a
var (Variable a -> abt '[] a)
-> Expect abt (Variable a) -> Expect abt (abt '[] a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> abt '[] a -> Expect abt (Variable a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> Expect abt (Variable a)
pushLet abt vars a
abt '[] a
lo
    abt '[] a
hi' <- Variable a -> abt '[] a
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
Variable a -> abt '[] a
var (Variable a -> abt '[] a)
-> Expect abt (Variable a) -> Expect abt (abt '[] a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> abt '[] a -> Expect abt (Variable a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> Expect abt (Variable a)
pushLet abt vars a
abt '[] a
hi
    abt '[] 'HReal
x   <- Variable 'HReal -> abt '[] 'HReal
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
Variable a -> abt '[] a
var (Variable 'HReal -> abt '[] 'HReal)
-> Expect abt (Variable 'HReal) -> Expect abt (abt '[] 'HReal)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> abt '[] 'HReal -> abt '[] 'HReal -> Expect abt (Variable 'HReal)
forall (abt :: [Hakaru] -> Hakaru -> *).
ABT Term abt =>
abt '[] 'HReal -> abt '[] 'HReal -> Expect abt (Variable 'HReal)
pushIntegrate abt '[] a
abt '[] 'HReal
lo' abt '[] a
abt '[] 'HReal
hi'
    Statement abt Location 'ExpectP -> Expect abt ()
forall (abt :: [Hakaru] -> Hakaru -> *) (m :: * -> *)
       (p :: Purity).
EvaluationMonad abt m p =>
Statement abt Location p -> m ()
unsafePush ((abt '[] 'HProb -> abt '[] 'HProb)
-> [Index (abt '[])] -> Statement abt Location 'ExpectP
forall (abt :: [Hakaru] -> Hakaru -> *) (v :: Hakaru -> *).
(abt '[] 'HProb -> abt '[] 'HProb)
-> [Index (abt '[])] -> Statement abt v 'ExpectP
SStuff0 (\abt '[] 'HProb
c -> abt '[] 'HReal
-> abt '[] 'HReal -> abt '[] 'HReal -> abt '[] 'HProb
forall (abt :: [Hakaru] -> Hakaru -> *).
ABT Term abt =>
abt '[] 'HReal
-> abt '[] 'HReal -> abt '[] 'HReal -> abt '[] 'HProb
P.densityUniform abt '[] a
abt '[] 'HReal
lo' abt '[] a
abt '[] 'HReal
hi' abt '[] 'HReal
x abt '[] 'HProb -> abt '[] 'HProb -> abt '[] 'HProb
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
(ABT Term abt, HSemiring_ a) =>
abt '[] a -> abt '[] a -> abt '[] a
P.* abt '[] 'HProb
c) [])
    abt '[] 'HReal -> Expect abt (abt '[] 'HReal)
forall (m :: * -> *) a. Monad m => a -> m a
return abt '[] 'HReal
x
    {-
    let_ lo $ \lo' ->
    let_ hi $ \hi' ->
    integrate lo' hi' $ \x ->
        densityUniform lo' hi' x * inst c x
    -}
expectMeasureOp MeasureOp typs a
Normal = \(abt vars a
mu :* abt vars a
sd :* SArgs abt args
End) -> do
    -- HACK: for some reason w need to break apart the 'emit' and the 'emit_' or else we get a "<<loop>>" exception. Not entirely sure why, but it prolly indicates a bug somewhere.
    abt '[] 'HReal
x <- Variable 'HReal -> abt '[] 'HReal
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
Variable a -> abt '[] a
var (Variable 'HReal -> abt '[] 'HReal)
-> Expect abt (Variable 'HReal) -> Expect abt (abt '[] 'HReal)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> abt '[] 'HReal -> abt '[] 'HReal -> Expect abt (Variable 'HReal)
forall (abt :: [Hakaru] -> Hakaru -> *).
ABT Term abt =>
abt '[] 'HReal -> abt '[] 'HReal -> Expect abt (Variable 'HReal)
pushIntegrate abt '[] 'HReal
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
(ABT Term abt, HRing_ a, Integrable a) =>
abt '[] a
P.negativeInfinity abt '[] 'HReal
forall (a :: Hakaru) (abt :: [Hakaru] -> Hakaru -> *).
(Integrable a, ABT Term abt) =>
abt '[] a
P.infinity
    Statement abt Location 'ExpectP -> Expect abt ()
forall (abt :: [Hakaru] -> Hakaru -> *) (m :: * -> *)
       (p :: Purity).
EvaluationMonad abt m p =>
Statement abt Location p -> m ()
unsafePush ((abt '[] 'HProb -> abt '[] 'HProb)
-> [Index (abt '[])] -> Statement abt Location 'ExpectP
forall (abt :: [Hakaru] -> Hakaru -> *) (v :: Hakaru -> *).
(abt '[] 'HProb -> abt '[] 'HProb)
-> [Index (abt '[])] -> Statement abt v 'ExpectP
SStuff0 (\abt '[] 'HProb
c -> abt '[] 'HReal
-> abt '[] 'HProb -> abt '[] 'HReal -> abt '[] 'HProb
forall (abt :: [Hakaru] -> Hakaru -> *).
ABT Term abt =>
abt '[] 'HReal
-> abt '[] 'HProb -> abt '[] 'HReal -> abt '[] 'HProb
P.densityNormal abt vars a
abt '[] 'HReal
mu abt vars a
abt '[] 'HProb
sd abt '[] 'HReal
x abt '[] 'HProb -> abt '[] 'HProb -> abt '[] 'HProb
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
(ABT Term abt, HSemiring_ a) =>
abt '[] a -> abt '[] a -> abt '[] a
P.* abt '[] 'HProb
c) [])
    abt '[] 'HReal -> Expect abt (abt '[] 'HReal)
forall (m :: * -> *) a. Monad m => a -> m a
return abt '[] 'HReal
x
    {-
    \c ->
        P.integrate P.negativeInfinity P.infinity $ \x ->
            P.densityNormal mu sd x P.* let_ x c)
    -}
expectMeasureOp MeasureOp typs a
Poisson = \(abt vars a
l :* SArgs abt args
End) -> do
    abt '[] a
l' <- Variable a -> abt '[] a
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
Variable a -> abt '[] a
var (Variable a -> abt '[] a)
-> Expect abt (Variable a) -> Expect abt (abt '[] a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> abt '[] a -> Expect abt (Variable a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> Expect abt (Variable a)
pushLet abt vars a
abt '[] a
l
    Statement abt Location 'ExpectP -> Expect abt ()
forall (abt :: [Hakaru] -> Hakaru -> *) (m :: * -> *)
       (p :: Purity).
EvaluationMonad abt m p =>
Statement abt Location p -> m ()
unsafePush ((abt '[] 'HProb -> abt '[] 'HProb)
-> [Index (abt '[])] -> Statement abt Location 'ExpectP
forall (abt :: [Hakaru] -> Hakaru -> *) (v :: Hakaru -> *).
(abt '[] 'HProb -> abt '[] 'HProb)
-> [Index (abt '[])] -> Statement abt v 'ExpectP
SStuff0 (\abt '[] 'HProb
c -> abt '[] HBool -> abt '[] 'HProb -> abt '[] 'HProb -> abt '[] 'HProb
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] HBool -> abt '[] a -> abt '[] a -> abt '[] a
P.if_ (abt '[] a
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
(ABT Term abt, HSemiring_ a) =>
abt '[] a
P.zero abt '[] a -> abt '[] a -> abt '[] HBool
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
(ABT Term abt, HOrd_ a) =>
abt '[] a -> abt '[] a -> abt '[] HBool
P.< abt '[] a
l') abt '[] 'HProb
c abt '[] 'HProb
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
(ABT Term abt, HSemiring_ a) =>
abt '[] a
P.zero) [])
    abt '[] 'HNat
x  <- Variable 'HNat -> abt '[] 'HNat
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
Variable a -> abt '[] a
var (Variable 'HNat -> abt '[] 'HNat)
-> Expect abt (Variable 'HNat) -> Expect abt (abt '[] 'HNat)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> abt '[] 'HNat -> abt '[] 'HNat -> Expect abt (Variable 'HNat)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
(ABT Term abt, HDiscrete_ a, SingI a) =>
abt '[] a -> abt '[] a -> Expect abt (Variable a)
pushSummate abt '[] 'HNat
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
(ABT Term abt, HSemiring_ a) =>
abt '[] a
P.zero abt '[] 'HNat
forall (a :: Hakaru) (abt :: [Hakaru] -> Hakaru -> *).
(Integrable a, ABT Term abt) =>
abt '[] a
P.infinity
    Statement abt Location 'ExpectP -> Expect abt ()
forall (abt :: [Hakaru] -> Hakaru -> *) (m :: * -> *)
       (p :: Purity).
EvaluationMonad abt m p =>
Statement abt Location p -> m ()
unsafePush ((abt '[] 'HProb -> abt '[] 'HProb)
-> [Index (abt '[])] -> Statement abt Location 'ExpectP
forall (abt :: [Hakaru] -> Hakaru -> *) (v :: Hakaru -> *).
(abt '[] 'HProb -> abt '[] 'HProb)
-> [Index (abt '[])] -> Statement abt v 'ExpectP
SStuff0 (\abt '[] 'HProb
c -> abt '[] 'HProb -> abt '[] 'HNat -> abt '[] 'HProb
forall (abt :: [Hakaru] -> Hakaru -> *).
ABT Term abt =>
abt '[] 'HProb -> abt '[] 'HNat -> abt '[] 'HProb
P.densityPoisson abt '[] a
abt '[] 'HProb
l' abt '[] 'HNat
x abt '[] 'HProb -> abt '[] 'HProb -> abt '[] 'HProb
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
(ABT Term abt, HSemiring_ a) =>
abt '[] a -> abt '[] a -> abt '[] a
P.* abt '[] 'HProb
c) [])
    abt '[] 'HNat -> Expect abt (abt '[] 'HNat)
forall (m :: * -> *) a. Monad m => a -> m a
return abt '[] 'HNat
x
    {-
    let_ l $ \l' ->
    if_ (zero < l')
        (summate zero infinity $ \x ->
            let x_ = unsafeFrom_ signed x in
            densityPoisson l' x_ * inst c x_)
        zero
    -}
expectMeasureOp MeasureOp typs a
Gamma = \(abt vars a
shape :* abt vars a
scale :* SArgs abt args
End) -> do
    abt '[] 'HReal
x  <- Variable 'HReal -> abt '[] 'HReal
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
Variable a -> abt '[] a
var (Variable 'HReal -> abt '[] 'HReal)
-> Expect abt (Variable 'HReal) -> Expect abt (abt '[] 'HReal)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> abt '[] 'HReal -> abt '[] 'HReal -> Expect abt (Variable 'HReal)
forall (abt :: [Hakaru] -> Hakaru -> *).
ABT Term abt =>
abt '[] 'HReal -> abt '[] 'HReal -> Expect abt (Variable 'HReal)
pushIntegrate abt '[] 'HReal
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
(ABT Term abt, HSemiring_ a) =>
abt '[] a
P.zero abt '[] 'HReal
forall (a :: Hakaru) (abt :: [Hakaru] -> Hakaru -> *).
(Integrable a, ABT Term abt) =>
abt '[] a
P.infinity
    abt '[] 'HProb
x_ <- Variable 'HProb -> abt '[] 'HProb
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
Variable a -> abt '[] a
var (Variable 'HProb -> abt '[] 'HProb)
-> Expect abt (Variable 'HProb) -> Expect abt (abt '[] 'HProb)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> abt '[] 'HProb -> Expect abt (Variable 'HProb)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> Expect abt (Variable a)
pushLet (abt '[] 'HReal -> abt '[] 'HProb
forall (abt :: [Hakaru] -> Hakaru -> *).
ABT Term abt =>
abt '[] 'HReal -> abt '[] 'HProb
P.unsafeProb abt '[] 'HReal
x) -- TODO: Or is this small enough that we'd be fine using Haskell's "let" and so duplicating the coercion of a variable however often?
    Statement abt Location 'ExpectP -> Expect abt ()
forall (abt :: [Hakaru] -> Hakaru -> *) (m :: * -> *)
       (p :: Purity).
EvaluationMonad abt m p =>
Statement abt Location p -> m ()
unsafePush ((abt '[] 'HProb -> abt '[] 'HProb)
-> [Index (abt '[])] -> Statement abt Location 'ExpectP
forall (abt :: [Hakaru] -> Hakaru -> *) (v :: Hakaru -> *).
(abt '[] 'HProb -> abt '[] 'HProb)
-> [Index (abt '[])] -> Statement abt v 'ExpectP
SStuff0 (\abt '[] 'HProb
c -> abt '[] 'HProb
-> abt '[] 'HProb -> abt '[] 'HProb -> abt '[] 'HProb
forall (abt :: [Hakaru] -> Hakaru -> *).
ABT Term abt =>
abt '[] 'HProb
-> abt '[] 'HProb -> abt '[] 'HProb -> abt '[] 'HProb
P.densityGamma abt vars a
abt '[] 'HProb
shape abt vars a
abt '[] 'HProb
scale abt '[] 'HProb
x_ abt '[] 'HProb -> abt '[] 'HProb -> abt '[] 'HProb
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
(ABT Term abt, HSemiring_ a) =>
abt '[] a -> abt '[] a -> abt '[] a
P.* abt '[] 'HProb
c) [])
    abt '[] 'HProb -> Expect abt (abt '[] 'HProb)
forall (m :: * -> *) a. Monad m => a -> m a
return abt '[] 'HProb
x_
    {-
    integrate zero infinity $ \x ->
        let x_ = unsafeProb x in
        densityGamma shape scale x_ * inst c x_
    -}
expectMeasureOp MeasureOp typs a
Beta = \(abt vars a
a :* abt vars a
b :* SArgs abt args
End) -> do
    abt '[] 'HReal
x  <- Variable 'HReal -> abt '[] 'HReal
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
Variable a -> abt '[] a
var (Variable 'HReal -> abt '[] 'HReal)
-> Expect abt (Variable 'HReal) -> Expect abt (abt '[] 'HReal)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> abt '[] 'HReal -> abt '[] 'HReal -> Expect abt (Variable 'HReal)
forall (abt :: [Hakaru] -> Hakaru -> *).
ABT Term abt =>
abt '[] 'HReal -> abt '[] 'HReal -> Expect abt (Variable 'HReal)
pushIntegrate abt '[] 'HReal
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
(ABT Term abt, HSemiring_ a) =>
abt '[] a
P.zero abt '[] 'HReal
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
(ABT Term abt, HSemiring_ a) =>
abt '[] a
P.one
    abt '[] 'HProb
x_ <- Variable 'HProb -> abt '[] 'HProb
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
Variable a -> abt '[] a
var (Variable 'HProb -> abt '[] 'HProb)
-> Expect abt (Variable 'HProb) -> Expect abt (abt '[] 'HProb)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> abt '[] 'HProb -> Expect abt (Variable 'HProb)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> Expect abt (Variable a)
pushLet (abt '[] 'HReal -> abt '[] 'HProb
forall (abt :: [Hakaru] -> Hakaru -> *).
ABT Term abt =>
abt '[] 'HReal -> abt '[] 'HProb
P.unsafeProb abt '[] 'HReal
x) -- TODO: Or is this small enough that we'd be fine using Haskell's "let" and so duplicating the coercion of a variable however often?
    Statement abt Location 'ExpectP -> Expect abt ()
forall (abt :: [Hakaru] -> Hakaru -> *) (m :: * -> *)
       (p :: Purity).
EvaluationMonad abt m p =>
Statement abt Location p -> m ()
unsafePush ((abt '[] 'HProb -> abt '[] 'HProb)
-> [Index (abt '[])] -> Statement abt Location 'ExpectP
forall (abt :: [Hakaru] -> Hakaru -> *) (v :: Hakaru -> *).
(abt '[] 'HProb -> abt '[] 'HProb)
-> [Index (abt '[])] -> Statement abt v 'ExpectP
SStuff0 (\abt '[] 'HProb
c -> abt '[] 'HProb
-> abt '[] 'HProb -> abt '[] 'HProb -> abt '[] 'HProb
forall (abt :: [Hakaru] -> Hakaru -> *).
ABT Term abt =>
abt '[] 'HProb
-> abt '[] 'HProb -> abt '[] 'HProb -> abt '[] 'HProb
P.densityBeta abt vars a
abt '[] 'HProb
a abt vars a
abt '[] 'HProb
b abt '[] 'HProb
x_ abt '[] 'HProb -> abt '[] 'HProb -> abt '[] 'HProb
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
(ABT Term abt, HSemiring_ a) =>
abt '[] a -> abt '[] a -> abt '[] a
P.* abt '[] 'HProb
c) [])
    abt '[] 'HProb -> Expect abt (abt '[] 'HProb)
forall (m :: * -> *) a. Monad m => a -> m a
return abt '[] 'HProb
x_
    {-
    integrate zero one $ \x ->
        let x_ = unsafeProb x in
        densityBeta a b x_ * inst c x_
    -}

----------------------------------------------------------------
----------------------------------------------------------- fin.