module Math.SpecialFunction ( incbeta
                            , beta
                            , gamma
                            , gammaln
                            ) where

import           Math.Hypergeometric

-- prop_betamatch :: Double -> Double -> Bool
-- prop_betamatch x y = x <= 0 || y <= 0 || abs (beta x y - incbeta 1 x y) < 1e-15

{-# SPECIALIZE incbeta :: Double -> Double -> Double -> Double #-}
-- | Incomplete beta function.
--
-- Calculated with \(B(z;a,b)=\displaystyle\frac{z^a}{a}{}_2F_1(a, 1-b; a+1; z)\)
--
-- @since 0.1.1.0
incbeta :: (Floating a, Eq a)
        => a -- ^ \(z\)
        -> a -- ^ \(a\)
        -> a -- ^ \(b\)
        -> a
incbeta z a b = z**a/a * hypergeometric [a,1-b] [a+1] z

{-# SPECIALIZE beta :: Double -> Double -> Double #-}
-- | \(B(x, y) = \displaystyle\frac{\Gamma(x)\Gamma(y)}{\Gamma(x+y)}\)
--
-- This uses 'gammaln' under the hood to extend its domain somewhat.
--
-- @since 0.1.1.0
beta :: (Floating a, Ord a) => a -> a -> a
beta x y = exp (betaln x y)

{-# SPECIALIZE betaln :: Double -> Double -> Double #-}
betaln :: (Floating a, Ord a) => a -> a -> a
betaln x y = gammaln x + gammaln y - gammaln (x+y)

-- | \(\Gamma(z)\)
--
-- @since 0.1.1.0
gamma :: (Floating a, Ord a) => a -> a
gamma = exp . gammaln

-- gamma from beta:
-- Γ(z)Γ(1-z) = 𝜋/sin(𝜋z)
--
-- THENCE, B(z,1-z)=Γ(z)Γ(1-z)/Γ(1)=...

{-# SPECIALIZE gammaln :: Double -> Double #-}
-- | \(\text{log} (\Gamma(z))\)
--
-- Lanczos approximation.
-- This is exactly the approach described in Press, William H. et al. /Numerical Recipes/, 3rd ed., extended to work on negative real numbers.
--
-- @since 0.1.1.0
gammaln :: (Floating a, Ord a)
        => a -- ^ \( z \)
        -> a
gammaln 0 = -log 0
gammaln z | z >= 0.5 = (z' + 1/2) * log (z' + 𝛾 + 1/2) - (z' + 1/2 + 𝛾) + log (sqrt (2*pi) * (c0 + sum series))
    where series = zipWith (\c x -> c / (z' + fromIntegral x)) coeff [(1::Int)..]
          c0 = 0.999999999999997092
          -- constants from Numerical Recipes
          coeff = [ 57.1562356658629235
                  , -59.5979603554754912
                  , 14.1360979747417471
                  , -0.491913816097620199
                  , 0.339946499848118887e-4
                  , 0.465236289270485756e-4
                  , -0.983744753048795646e-4
                  , 0.158088703224912494e-3
                  , -0.210264441724104883e-3
                  , 0.217439618115212643e-3
                  , -0.164318106536763890e-3
                  , 0.844182239838527433e-4
                  , -0.261908384015814087e-4
                  , 0.368991826595316234e-5
                  ]
          𝛾 = 607/128
          z' = z-1
gammaln z = log pi - log (sin (pi * z)) - gammaln (1 - z)