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

module Data.Random.Distribution.Dirichlet where

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

import Data.List

fractionalDirichlet :: (Fractional a, Distribution Gamma a) => [a] -> RVarT m [a]
fractionalDirichlet :: forall a (m :: * -> *).
(Fractional a, Distribution Gamma a) =>
[a] -> RVarT m [a]
fractionalDirichlet []  = forall (m :: * -> *) a. Monad m => a -> m a
return []
fractionalDirichlet [a
_] = forall (m :: * -> *) a. Monad m => a -> m a
return [a
1]
fractionalDirichlet [a]
as = do
    [a]
xs <- forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence [forall a (m :: * -> *). Distribution Gamma a => a -> a -> RVarT m a
gammaT a
a a
1 | a
a <- [a]
as]
    let total :: a
total = forall a. (a -> a -> a) -> [a] -> a
foldl1' forall a. Num a => a -> a -> a
(+) [a]
xs

    forall (m :: * -> *) a. Monad m => a -> m a
return (forall a b. (a -> b) -> [a] -> [b]
map (forall a. Num a => a -> a -> a
* forall a. Fractional a => a -> a
recip a
total) [a]
xs)

dirichlet :: Distribution Dirichlet [a] => [a] -> RVar [a]
dirichlet :: forall a. Distribution Dirichlet [a] => [a] -> RVar [a]
dirichlet [a]
as = forall (d :: * -> *) t. Distribution d t => d t -> RVar t
rvar (forall a. a -> Dirichlet a
Dirichlet [a]
as)

dirichletT :: Distribution Dirichlet [a] => [a] -> RVarT m [a]
dirichletT :: forall a (m :: * -> *).
Distribution Dirichlet [a] =>
[a] -> RVarT m [a]
dirichletT [a]
as = forall (d :: * -> *) t (n :: * -> *).
Distribution d t =>
d t -> RVarT n t
rvarT (forall a. a -> Dirichlet a
Dirichlet [a]
as)

newtype Dirichlet a = Dirichlet a deriving (Dirichlet a -> Dirichlet a -> Bool
forall a. Eq a => Dirichlet a -> Dirichlet a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Dirichlet a -> Dirichlet a -> Bool
$c/= :: forall a. Eq a => Dirichlet a -> Dirichlet a -> Bool
== :: Dirichlet a -> Dirichlet a -> Bool
$c== :: forall a. Eq a => Dirichlet a -> Dirichlet a -> Bool
Eq, Int -> Dirichlet a -> ShowS
forall a. Show a => Int -> Dirichlet a -> ShowS
forall a. Show a => [Dirichlet a] -> ShowS
forall a. Show a => Dirichlet a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Dirichlet a] -> ShowS
$cshowList :: forall a. Show a => [Dirichlet a] -> ShowS
show :: Dirichlet a -> String
$cshow :: forall a. Show a => Dirichlet a -> String
showsPrec :: Int -> Dirichlet a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> Dirichlet a -> ShowS
Show)

instance (Fractional a, Distribution Gamma a) => Distribution Dirichlet [a] where
    rvarT :: forall (n :: * -> *). Dirichlet [a] -> RVarT n [a]
rvarT (Dirichlet [a]
as) = forall a (m :: * -> *).
(Fractional a, Distribution Gamma a) =>
[a] -> RVarT m [a]
fractionalDirichlet [a]
as