    {-# LANGUAGE OverloadedStrings #-}
    {-# LANGUAGE DeriveDataTypeable, DeriveGeneric #-}
    -----------------------------------------------------------------------------
    -- |
    -- Module      :  Statistics.Distribution.Beta
    -- Copyright   :  (C) 2012 Edward Kmett,
    -- License     :  BSD-style (see the file LICENSE)
    --
    -- Maintainer  :  Edward Kmett <ekmett@gmail.com>
    -- Stability   :  provisional
    -- Portability :  DeriveDataTypeable
    --
    ----------------------------------------------------------------------------
    module Statistics.Distribution.Beta
      ( BetaDistribution
        -- * Constructor
      , betaDistr
      , betaDistrE
      , improperBetaDistr
      , improperBetaDistrE
        -- * Accessors
      , bdAlpha
      , bdBeta
      ) where

    import Control.Applicative
    import Data.Aeson            (FromJSON(..), ToJSON, Value(..), (.:))
    import Data.Binary           (Binary(..))
    import Data.Data             (Data, Typeable)
    import GHC.Generics          (Generic)
    import Numeric.SpecFunctions (
      incompleteBeta, invIncompleteBeta, logBeta, digamma, log1p)
    import Numeric.MathFunctions.Constants (m_NaN,m_neg_inf)
    import qualified Statistics.Distribution as D
    import Statistics.Internal

    -- | The beta distribution
    data BetaDistribution = BD
     { bdAlpha :: {-# UNPACK #-} !Double
       -- ^ Alpha shape parameter
     , bdBeta  :: {-# UNPACK #-} !Double
       -- ^ Beta shape parameter
     } deriving (Eq, Typeable, Data, Generic)

    instance Show BetaDistribution where
      showsPrec n (BD a b) = defaultShow2 "improperBetaDistr" a b n
    instance Read BetaDistribution where
      readPrec = defaultReadPrecM2 "improperBetaDistr" improperBetaDistrE

    instance ToJSON BetaDistribution
    instance FromJSON BetaDistribution where
      parseJSON (Object v) = do
        a <- v .: "bdAlpha"
        b <- v .: "bdBeta"
        maybe (fail $ errMsgI a b) return $ improperBetaDistrE a b
      parseJSON _ = empty

    instance Binary BetaDistribution where
      put (BD a b) = put a >> put b
      get = do
        a <- get
        b <- get
        maybe (fail $ errMsgI a b) return $ improperBetaDistrE a b

    -- | Create beta distribution. Both shape parameters must be positive.
    betaDistr :: Double             -- ^ Shape parameter alpha
              -> Double             -- ^ Shape parameter beta
              -> BetaDistribution
    betaDistr a b = maybe (error $ errMsg a b) id $ betaDistrE a b

    -- | Create beta distribution. Both shape parameters must be positive.
    betaDistrE :: Double             -- ^ Shape parameter alpha
              -> Double             -- ^ Shape parameter beta
              -> Maybe BetaDistribution
    betaDistrE a b
      | a > 0 && b > 0 = Just (BD a b)
      | otherwise      = Nothing

    errMsg :: Double -> Double -> String
    errMsg a b = "Statistics.Distribution.Beta.betaDistr: "
              ++ "shape parameters must be positive. Got a = "
              ++ show a
              ++ " b = "
              ++ show b

    -- | Create beta distribution. Both shape parameters must be
    -- non-negative. So it allows to construct improper beta distribution
    -- which could be used as improper prior.
    improperBetaDistr :: Double             -- ^ Shape parameter alpha
                      -> Double             -- ^ Shape parameter beta
                      -> BetaDistribution
    improperBetaDistr a b
      = maybe (error $ errMsgI a b) id $ improperBetaDistrE a b

    -- | Create beta distribution. Both shape parameters must be
    -- non-negative. So it allows to construct improper beta distribution
    -- which could be used as improper prior.
    improperBetaDistrE :: Double             -- ^ Shape parameter alpha
                       -> Double             -- ^ Shape parameter beta
                       -> Maybe BetaDistribution
    improperBetaDistrE a b
      | a >= 0 && b >= 0 = Just (BD a b)
      | otherwise        = Nothing

    errMsgI :: Double -> Double -> String
    errMsgI a b
      =  "Statistics.Distribution.Beta.betaDistr: "
      ++ "shape parameters must be non-negative. Got a = " ++ show a
      ++ " b = " ++ show b

    instance D.Distribution BetaDistribution where
      cumulative (BD a b) x
        | x <= 0    = 0
        | x >= 1    = 1
        | otherwise = incompleteBeta a b x
      complCumulative (BD a b) x
        | x <= 0    = 1
        | x >= 1    = 0
        -- For small x we use direct computation to avoid precision loss
        -- when computing (1-x)
        | x <  0.5  = 1 - incompleteBeta a b x
        -- Otherwise we use property of incomplete beta:
        --  > I(x,a,b) = 1 - I(1-x,b,a)
        | otherwise = incompleteBeta b a (1-x)

    instance D.Mean BetaDistribution where
      mean (BD a b) = a / (a + b)

    instance D.MaybeMean BetaDistribution where
      maybeMean = Just . D.mean

    instance D.Variance BetaDistribution where
      variance (BD a b) = a*b / (apb*apb*(apb+1))
        where apb = a + b

    instance D.MaybeVariance BetaDistribution where
      maybeVariance = Just . D.variance

    instance D.Entropy BetaDistribution where
      entropy (BD a b) =
        logBeta a b
        - (a-1) * digamma a
        - (b-1) * digamma b
        + (a+b-2) * digamma (a+b)

    instance D.MaybeEntropy BetaDistribution where
      maybeEntropy = Just . D.entropy

    instance D.ContDistr BetaDistribution where
      density (BD a b) x
        | a <= 0 || b <= 0 = m_NaN
        | x <= 0 = 0
        | x >= 1 = 0
        | otherwise = exp $ (a-1)*log x + (b-1) * log1p (-x) - logBeta a b
      logDensity (BD a b) x
        | a <= 0 || b <= 0 = m_NaN
        | x <= 0 = m_neg_inf
        | x >= 1 = m_neg_inf
        | otherwise = (a-1)*log x + (b-1)*log1p (-x) - logBeta a b

      quantile (BD a b) p
        | p == 0         = 0
        | p == 1         = 1
        | p > 0 && p < 1 = invIncompleteBeta a b p
        | otherwise      =
            error $ "Statistics.Distribution.Gamma.quantile: p must be in [0,1] range. Got: "++show p

    instance D.ContGen BetaDistribution where
      genContVar = D.genContinuous