{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE DeriveGeneric      #-}

{- |
   Module      :  ELynx.Distribution.BirthDeathNearlyCritical
   Description :  Birth and death distribution
   Copyright   :  (c) Dominik Schrempf 2018
   License     :  GPL-3.0-or-later

   Maintainer  :  dominik.schrempf@gmail.com
   Stability   :  unstable
   Portability :  portable

Creation date: Tue Feb 13 13:16:18 2018.

See Gernhard, T. (2008). The conditioned reconstructed process. Journal of
Theoretical Biology, 253(4), 769–778. http://doi.org/10.1016/j.jtbi.2008.04.005.

Distribution of the values of the point process such that it corresponds to
reconstructed trees under the birth and death process; nearly critical birth and
death process with lambda~mu.

Basically, this is a Taylor expansion of Eq. (2) and Eq. (3).

-}

module ELynx.Distribution.BirthDeathNearlyCritical
  ( BirthDeathNearlyCriticalDistribution(..)
  , cumulative
  , density
  , quantile
  )
where

import           Data.Data                      ( Data
                                                , Typeable
                                                )
import           GHC.Generics                   ( Generic )
import qualified Statistics.Distribution       as D

import           ELynx.Distribution.Types

-- | Distribution of the values of the point process such that it corresponds to
-- a reconstructed tree of the birth and death process.
data BirthDeathNearlyCriticalDistribution = BDNCD
  { bdncdTOr :: Time    -- ^ Time to origin of the tree.
  , bdncdLa  :: Rate    -- ^ Birth and death rate.
  , bdncdMu  :: Rate    -- ^ Birth and death rate.
  } deriving (Eq, Typeable, Data, Generic)

instance D.Distribution BirthDeathNearlyCriticalDistribution where
  cumulative = cumulative

-- | Cumulative distribution function section 2.1.2, second formula.
cumulative :: BirthDeathNearlyCriticalDistribution -> Time -> Double
cumulative (BDNCD t l m) s | s <= 0    = 0
                           | s > t     = 1
                           | otherwise = o0 + o1
 where
  o0 = s * (1.0 + t * l) / t / (1.0 + s * l)
  o1 = (-s * s + s * t) * (m - l) / (2.0 * t * (1.0 + s * l) ** 2)

instance D.ContDistr BirthDeathNearlyCriticalDistribution where
  density  = density
  quantile = quantile

-- | Density function section 2.1.2, first formula.
density :: BirthDeathNearlyCriticalDistribution -> Time -> Double
density (BDNCD t l m) s | s < 0     = 0
                        | s > t     = 0
                        | otherwise = o0 + o1
 where
  o0 = (1.0 + t * l) / (t * (1.0 + s * l) ** 2)
  o1 = (-2.0 * s + t - s * t * l) * (m - l) / (2.0 * t * (1.0 + s * l) ** 3)

-- | Inverted cumulative probability distribution 'cumulative'. See also
-- 'D.ContDistr'.
quantile :: BirthDeathNearlyCriticalDistribution -> Double -> Time
quantile (BDNCD t l m) p
  | p >= 0 && p <= 1
  = res
  | otherwise
  = error
    $  "PointProcess.quantile: p must be in [0,1] range. Got: "
    ++ show p
    ++ "."
 where
  den = l * (-3.0 + 2.0 * t * (-1.0 + p) * l) + m
  t1  = (2.0 + t * (l - 4.0 * p * l + m)) / den
  t2Nom =
    4.0
      + t
      * ( l
        * (4.0 + t * l + 8.0 * p * (1.0 + t * l))
        + 2.0
        * (2.0 + t * l - 4.0 * p * (1.0 + t * l))
        * m
        + t
        * m
        * m
        )
  t2  = t2Nom / (den ** 2)
  res = 0.5 * (t1 + sqrt t2)

instance D.ContGen BirthDeathNearlyCriticalDistribution where
  genContVar = D.genContinuous