{-# LANGUAGE
    MultiParamTypeClasses,
    FlexibleInstances, FlexibleContexts,
    UndecidableInstances
  #-}

{-# OPTIONS_GHC -fno-warn-simplifiable-class-constraints #-}

module Data.Random.Distribution.Beta where

import Data.Random.RVar
import Data.Random.Distribution
import Data.Random.Distribution.Gamma
import Data.Random.Distribution.Uniform

import Numeric.SpecFunctions

{-# SPECIALIZE fractionalBeta :: Float  -> Float  -> RVarT m Float #-}
{-# SPECIALIZE fractionalBeta :: Double -> Double -> RVarT m Double #-}
fractionalBeta :: (Fractional a, Eq a, Distribution Gamma a, Distribution StdUniform a) => a -> a -> RVarT m a
fractionalBeta :: forall a (m :: * -> *).
(Fractional a, Eq a, Distribution Gamma a,
 Distribution StdUniform a) =>
a -> a -> RVarT m a
fractionalBeta a
1 a
1 = forall a (m :: * -> *). Distribution StdUniform a => RVarT m a
stdUniformT
fractionalBeta a
a a
b = do
    a
x <- forall a (m :: * -> *). Distribution Gamma a => a -> a -> RVarT m a
gammaT a
a a
1
    a
y <- forall a (m :: * -> *). Distribution Gamma a => a -> a -> RVarT m a
gammaT a
b a
1
    forall (m :: * -> *) a. Monad m => a -> m a
return (a
x forall a. Fractional a => a -> a -> a
/ (a
x forall a. Num a => a -> a -> a
+ a
y))

{-# SPECIALIZE beta :: Float  -> Float  -> RVar Float #-}
{-# SPECIALIZE beta :: Double -> Double -> RVar Double #-}
beta :: Distribution Beta a => a -> a -> RVar a
beta :: forall a. Distribution Beta a => a -> a -> RVar a
beta a
a a
b = forall (d :: * -> *) t. Distribution d t => d t -> RVar t
rvar (forall a. a -> a -> Beta a
Beta a
a a
b)

{-# SPECIALIZE betaT :: Float  -> Float  -> RVarT m Float #-}
{-# SPECIALIZE betaT :: Double -> Double -> RVarT m Double #-}
betaT :: Distribution Beta a => a -> a -> RVarT m a
betaT :: forall a (m :: * -> *). Distribution Beta a => a -> a -> RVarT m a
betaT a
a a
b = forall (d :: * -> *) t (n :: * -> *).
Distribution d t =>
d t -> RVarT n t
rvarT (forall a. a -> a -> Beta a
Beta a
a a
b)

data Beta a = Beta a a

-- FIXME: I am far from convinced that NaNs are a good idea.
logBetaPdf :: Double -> Double -> Double -> Double
logBetaPdf :: Double -> Double -> Double -> Double
logBetaPdf Double
a Double
b Double
x
   | Double
a forall a. Ord a => a -> a -> Bool
<= Double
0 Bool -> Bool -> Bool
|| Double
b forall a. Ord a => a -> a -> Bool
<= Double
0 = Double
nan
   | Double
x forall a. Ord a => a -> a -> Bool
<= Double
0 = forall a. Floating a => a -> a
log Double
0
   | Double
x forall a. Ord a => a -> a -> Bool
>= Double
1 = forall a. Floating a => a -> a
log Double
0
   | Bool
otherwise = (Double
aforall a. Num a => a -> a -> a
-Double
1)forall a. Num a => a -> a -> a
*forall a. Floating a => a -> a
log Double
x forall a. Num a => a -> a -> a
+ (Double
bforall a. Num a => a -> a -> a
-Double
1)forall a. Num a => a -> a -> a
*forall a. Floating a => a -> a
log (Double
1forall a. Num a => a -> a -> a
-Double
x) forall a. Num a => a -> a -> a
- Double -> Double -> Double
logBeta Double
a Double
b
  where
    nan :: Double
nan = Double
0.0 forall a. Fractional a => a -> a -> a
/ Double
0.0

instance PDF Beta Double
  where
    pdf :: Beta Double -> Double -> Double
pdf (Beta Double
a Double
b) = forall a. Floating a => a -> a
exp forall b c a. (b -> c) -> (a -> b) -> a -> c
. Double -> Double -> Double -> Double
logBetaPdf Double
a Double
b

instance PDF Beta Float
  where
    pdf :: Beta Float -> Float -> Double
pdf (Beta Float
a Float
b) = forall a b. (Real a, Fractional b) => a -> b
realToFrac forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Floating a => a -> a
exp forall b c a. (b -> c) -> (a -> b) -> a -> c
. Double -> Double -> Double -> Double
logBetaPdf (forall a b. (Real a, Fractional b) => a -> b
realToFrac Float
a) (forall a b. (Real a, Fractional b) => a -> b
realToFrac Float
b) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (Real a, Fractional b) => a -> b
realToFrac

instance Distribution Beta Float
  where
    rvarT :: forall (n :: * -> *). Beta Float -> RVarT n Float
rvarT (Beta Float
a Float
b) = forall a (m :: * -> *).
(Fractional a, Eq a, Distribution Gamma a,
 Distribution StdUniform a) =>
a -> a -> RVarT m a
fractionalBeta Float
a Float
b

instance Distribution Beta Double
  where
    rvarT :: forall (n :: * -> *). Beta Double -> RVarT n Double
rvarT (Beta Double
a Double
b) = forall a (m :: * -> *).
(Fractional a, Eq a, Distribution Gamma a,
 Distribution StdUniform a) =>
a -> a -> RVarT m a
fractionalBeta Double
a Double
b