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

module Data.Random.Distribution.Simplex
    ( StdSimplex(..)
    , stdSimplex
    , stdSimplexT
    , fractionalStdSimplex
    ) where

import Control.Applicative
import Control.Monad
import Data.List
import Data.Random.RVar
import Data.Random.Distribution
import Data.Random.Distribution.Uniform

-- |Uniform distribution over a standard simplex.
newtype StdSimplex as =
    -- | @StdSimplex k@ constructs a standard simplex of dimension @k@
    -- (standard /k/-simplex).
    -- An element of the simplex represents a vector variable @as = (a_0,
    -- a_1, ..., a_k)@. The elements of @as@ are more than or equal to @0@
    -- and @sum as@ is always equal to @1@.
    StdSimplex Int
    deriving (Eq, Show)

instance (Ord a, Fractional a, Distribution StdUniform a) => Distribution StdSimplex [a] where
    rvar (StdSimplex k) = fractionalStdSimplex k

-- |@stdSimplex k@ returns a random variable being uniformly distributed over
-- a standard simplex of dimension @k@.
stdSimplex :: Distribution StdSimplex [a] => Int -> RVar [a]
stdSimplex k = rvar (StdSimplex k)

stdSimplexT :: Distribution StdSimplex [a] => Int -> RVarT m [a]
stdSimplexT k = rvarT (StdSimplex k)

-- |An algorithm proposed by Rubinstein & Melamed (1998).
-- See, /e.g./, S. Onn, I. Weissman.
-- Generating uniform random vectors over a simplex with implications to
-- the volume of a certain polytope and to multivariate extremes.
-- /Ann Oper Res/ (2011) __189__:331-342.
fractionalStdSimplex :: (Ord a, Fractional a, Distribution StdUniform a) => Int -> RVar [a]
fractionalStdSimplex k = do us <- replicateM k stdUniform
                            let us' = sort us ++ [1]
                            return $ zipWith (-) us' (0 : us')