    {-# LANGUAGE OverloadedStrings #-}
    {-# LANGUAGE DeriveDataTypeable, DeriveGeneric #-}
    -- |
    -- Module    : Statistics.Distribution.Geometric
    -- Copyright : (c) 2009 Bryan O'Sullivan
    -- License   : BSD3
    --
    -- Maintainer  : bos@serpentine.com
    -- Stability   : experimental
    -- Portability : portable
    --
    -- The Geometric distribution. There are two variants of
    -- distribution. First is the probability distribution of the number
    -- of Bernoulli trials needed to get one success, supported on the set
    -- [1,2..] ('GeometricDistribution'). Sometimes it's referred to as
    -- the /shifted/ geometric distribution to distinguish from another
    -- one.
    --
    -- Second variant is probability distribution of the number of
    -- failures before first success, defined over the set [0,1..]
    -- ('GeometricDistribution0').
    module Statistics.Distribution.Geometric
        (
          GeometricDistribution
        , GeometricDistribution0
        -- * Constructors
        , geometric
        , geometricE
        , geometric0
        , geometric0E
        -- ** Accessors
        , gdSuccess
        , gdSuccess0
        ) where

    import Control.Applicative
    import Control.Monad       (liftM)
    import Data.Aeson          (FromJSON(..), ToJSON, Value(..), (.:))
    import Data.Binary         (Binary(..))
    import Data.Data           (Data, Typeable)
    import GHC.Generics        (Generic)
    import Numeric.MathFunctions.Constants (m_pos_inf, m_neg_inf)
    import qualified System.Random.MWC.Distributions as MWC

    import qualified Statistics.Distribution as D
    import Statistics.Internal

    ----------------------------------------------------------------

    -- | Distribution over [1..]
    newtype GeometricDistribution = GD {
          gdSuccess :: Double
        } deriving (Eq, Typeable, Data, Generic)

    instance Show GeometricDistribution where
      showsPrec i (GD x) = defaultShow1 "geometric" x i
    instance Read GeometricDistribution where
      readPrec = defaultReadPrecM1 "geometric" geometricE

    instance ToJSON GeometricDistribution
    instance FromJSON GeometricDistribution where
      parseJSON (Object v) = do
        x <- v .: "gdSuccess"
        maybe (fail $ errMsg x) return  $ geometricE x
      parseJSON _ = empty

    instance Binary GeometricDistribution where
      put (GD x) = put x
      get = do
        x <- get
        maybe (fail $ errMsg x) return  $ geometricE x

    instance D.Distribution GeometricDistribution where
        cumulative = cumulative

    instance D.DiscreteDistr GeometricDistribution where
        probability (GD s) n
          | n < 1     = 0
          | otherwise = s * (1-s) ** (fromIntegral n - 1)
        logProbability (GD s) n
           | n < 1     = m_neg_inf
           | otherwise = log s + log (1-s) * (fromIntegral n - 1)

    instance D.Mean GeometricDistribution where
        mean (GD s) = 1 / s

    instance D.Variance GeometricDistribution where
        variance (GD s) = (1 - s) / (s * s)

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

    instance D.MaybeVariance GeometricDistribution where
        maybeStdDev   = Just . D.stdDev
        maybeVariance = Just . D.variance

    instance D.Entropy GeometricDistribution where
      entropy (GD s)
        | s == 0 = m_pos_inf
        | s == 1 = 0
        | otherwise = negate $ (s * log s + (1-s) * log (1-s)) / s

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

    instance D.DiscreteGen GeometricDistribution where
      genDiscreteVar (GD s) g = MWC.geometric1 s g

    instance D.ContGen GeometricDistribution where
      genContVar d g = fromIntegral `liftM` D.genDiscreteVar d g

    cumulative :: GeometricDistribution -> Double -> Double
    cumulative (GD s) x
      | x < 1        = 0
      | isInfinite x = 1
      | isNaN      x = error "Statistics.Distribution.Geometric.cumulative: NaN input"
      | otherwise    = 1 - (1-s) ^ (floor x :: Int)

    -- | Create geometric distribution.
    geometric :: Double                -- ^ Success rate
              -> GeometricDistribution
    geometric x = maybe (error $ errMsg x) id $ geometricE x

    -- | Create geometric distribution.
    geometricE :: Double                -- ^ Success rate
               -> Maybe GeometricDistribution
    geometricE x
      | x >= 0 && x <= 1 = Just (GD x)
      | otherwise        = Nothing

    errMsg :: Double -> String
    errMsg x = "Statistics.Distribution.Geometric.geometric: probability must be in [0,1] range. Got " ++ show x

    ----------------------------------------------------------------

    -- | Distribution over [0..]
    newtype GeometricDistribution0 = GD0 {
          gdSuccess0 :: Double
        } deriving (Eq, Typeable, Data, Generic)

    instance Show GeometricDistribution0 where
      showsPrec i (GD0 x) = defaultShow1 "geometric0" x i
    instance Read GeometricDistribution0 where
      readPrec = defaultReadPrecM1 "geometric0" geometric0E

    instance ToJSON GeometricDistribution0
    instance FromJSON GeometricDistribution0 where
      parseJSON (Object v) = do
        x <- v .: "gdSuccess0"
        maybe (fail $ errMsg x) return  $ geometric0E x
      parseJSON _ = empty

    instance Binary GeometricDistribution0 where
      put (GD0 x) = put x
      get = do
        x <- get
        maybe (fail $ errMsg x) return  $ geometric0E x

    instance D.Distribution GeometricDistribution0 where
        cumulative (GD0 s) x = cumulative (GD s) (x + 1)

    instance D.DiscreteDistr GeometricDistribution0 where
        probability    (GD0 s) n = D.probability    (GD s) (n + 1)
        logProbability (GD0 s) n = D.logProbability (GD s) (n + 1)

    instance D.Mean GeometricDistribution0 where
        mean (GD0 s) = 1 / s - 1

    instance D.Variance GeometricDistribution0 where
        variance (GD0 s) = D.variance (GD s)

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

    instance D.MaybeVariance GeometricDistribution0 where
        maybeStdDev   = Just . D.stdDev
        maybeVariance = Just . D.variance

    instance D.Entropy GeometricDistribution0 where
      entropy (GD0 s) = D.entropy (GD s)

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

    instance D.DiscreteGen GeometricDistribution0 where
      genDiscreteVar (GD0 s) g = MWC.geometric0 s g

    instance D.ContGen GeometricDistribution0 where
      genContVar d g = fromIntegral `liftM` D.genDiscreteVar d g

    -- | Create geometric distribution.
    geometric0 :: Double                -- ^ Success rate
               -> GeometricDistribution0
    geometric0 x = maybe (error $ errMsg0 x) id $ geometric0E x

    -- | Create geometric distribution.
    geometric0E :: Double                -- ^ Success rate
                -> Maybe GeometricDistribution0
    geometric0E x
      | x >= 0 && x <= 1 = Just (GD0 x)
      | otherwise        = Nothing

    errMsg0 :: Double -> String
    errMsg0 x = "Statistics.Distribution.Geometric.geometric0: probability must be in [0,1] range. Got " ++ show x