content: {-# LANGUAGE MultiParamTypeClasses, FunctionalDependencies #-}
  
  module Numeric.ADEV.Class (
    D(..), C(..), ADEV(..)) where
  
  import Numeric.Log
  
  -- | Type of density-carrying distributions.
  data D m r a = D (m a) (a -> Log r)
  
  -- | Type of CDF-carrying distributions over the reals, 
  -- for implicit reparameterization.
  data C m r = C (m Double) (Double -> Log Double) (r -> Log r)
  
  -- ---------------------------------------------------------------------------
  -- | A typeclass for ADEV programs, parameterized by:
  -- 
  --   * @r@ - the type used to represent real numbers
  --   * @m@ - the monad used to encode randomness (so that @m r@ is the type of
  --           unbiasedly estimated real numbers)
  --   * @p@ - the type used for monadic probabilistic programming (so 
  --           that @p m a@ is a probabilistic program returning @a@)
  class (RealFrac r, Monad (p m), Monad m) => ADEV p m r | p -> r, r -> p where
    -- | Sample a random uniform value between 0 and 1.
    sample           :: p m r
    -- | Add a real value into a running cost accumulator.
    -- When a @p m r@ is passed to @expect@, the result is
    -- an estimator of the expected cost *plus* the expected
    -- return value.
    add_cost         :: r -> p m ()
    -- | Flip a coin with a specified probability of heads.
    -- Uses enumeration (costly but low-variance) to estimate
    -- gradients.
    flip_enum        :: r -> p m Bool
    -- | Flip a coin with a specified probability of heads.
    -- Uses the REINFORCE estimator (cheaper but higher-variance) 
    -- for gradients.
    flip_reinforce   :: r -> p m Bool
    -- | Generate from a normal distribution. Uses the REPARAM gradient estimator.
    normal_reparam   :: r -> r -> p m r
    -- | Generate from a normal distribution. Uses the REINFORCE gradient estimator.
    normal_reinforce :: r -> r -> p m r
    -- | Estimate the expectation of a probabilistic computation.
    expect           :: p m r -> m r
    -- | Combinator DSL for estimators
    plus_            :: m r -> m r -> m r
    times_           :: m r -> m r -> m r
    exp_             :: m r -> m r
    minibatch_       :: Int -> Int -> (Int -> m r) -> m r
    exact_           :: r -> m r
    -- | Baselines for controlling variance
    baseline         :: p m r -> r -> m r
    -- | Automatic construction of new REINFORCE estimators
    reinforce        :: D m r a -> p m a
    -- | Storchastic leave_one_out estimator
    leave_one_out    :: Int -> D m r a -> p m a
    -- | Differentiable particle filter, accepting:
    --   * @p@: a density function for the target measure.
    --   * @q0@: an initial proposal for the particle filter.
    --   * @q@: a transition proposal for the particle filter.
    --   * @f@: an unbiased estimator of an integrand to estimate
    --   * @n@: the number of SMC steps to run
    --   * @k@: the number of particles to use
    -- Returns an SMC estimator of the integral
    smc :: ([a] -> Log r) -> D m r a -> (a -> D m r a) -> ([a] -> m r) -> Int -> Int -> m r
    -- | Importance sampling gradient estimator
    importance       :: D m r a -> D m r a -> p m a
    -- | Implicit reparameterization for real-valued distributions 
    -- differentiable with CDFs (e.g., mixtures of Gaussians)
    implicit_reparam :: C m r -> p m r
    -- | Sample from a Poisson distribution, using a measure-valued derivative.
    poisson_weak     :: Log r -> p m Int
    -- | Gradients through rejection sampling for density-carrying distributions.
    reparam_reject   :: D m r a -> (a -> b) -> (D m r b) -> (D m r b) -> Log r -> p m b
  
  
  