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 :: forall a. (Floating a, Eq a) => a -> a -> a -> a
incbeta a
z a
a a
b = a
za -> a -> a
forall a. Floating a => a -> a -> a
**a
aa -> a -> a
forall a. Fractional a => a -> a -> a
/a
a a -> a -> a
forall a. Num a => a -> a -> a
* [a] -> [a] -> a -> a
forall a. (Eq a, Fractional a) => [a] -> [a] -> a -> a
hypergeometric [a
a,a
1a -> a -> a
forall a. Num a => a -> a -> a
-a
b] [a
aa -> a -> a
forall a. Num a => a -> a -> a
+a
1] a
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 :: forall a. (Floating a, Ord a) => a -> a -> a
beta a
x a
y = a -> a
forall a. Floating a => a -> a
exp (a -> a -> a
forall a. (Floating a, Ord a) => a -> a -> a
betaln a
x a
y)

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

-- | \(\Gamma(z)\)
--
-- @since 0.1.1.0
gamma :: (Floating a, Ord a) => a -> a
gamma :: forall a. (Floating a, Ord a) => a -> a
gamma = a -> a
forall a. Floating a => a -> a
exp (a -> a) -> (a -> a) -> a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> a
forall a. (Floating a, Ord a) => a -> a
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 :: forall a. (Floating a, Ord a) => a -> a
gammaln a
0 = -a -> a
forall a. Floating a => a -> a
log a
0
gammaln a
z | a
z a -> a -> Bool
forall a. Ord a => a -> a -> Bool
>= a
0.5 = (a
z' a -> a -> a
forall a. Num a => a -> a -> a
+ a
1a -> a -> a
forall a. Fractional a => a -> a -> a
/a
2) a -> a -> a
forall a. Num a => a -> a -> a
* a -> a
forall a. Floating a => a -> a
log (a
z' a -> a -> a
forall a. Num a => a -> a -> a
+ a
𝛾 a -> a -> a
forall a. Num a => a -> a -> a
+ a
1a -> a -> a
forall a. Fractional a => a -> a -> a
/a
2) a -> a -> a
forall a. Num a => a -> a -> a
- (a
z' a -> a -> a
forall a. Num a => a -> a -> a
+ a
1a -> a -> a
forall a. Fractional a => a -> a -> a
/a
2 a -> a -> a
forall a. Num a => a -> a -> a
+ a
𝛾) a -> a -> a
forall a. Num a => a -> a -> a
+ a -> a
forall a. Floating a => a -> a
log (a -> a
forall a. Floating a => a -> a
sqrt (a
2a -> a -> a
forall a. Num a => a -> a -> a
*a
forall a. Floating a => a
pi) a -> a -> a
forall a. Num a => a -> a -> a
* (a
c0 a -> a -> a
forall a. Num a => a -> a -> a
+ [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [a]
series))
    where series :: [a]
series = (a -> Int -> a) -> [a] -> [Int] -> [a]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\a
c Int
x -> a
c a -> a -> a
forall a. Fractional a => a -> a -> a
/ (a
z' a -> a -> a
forall a. Num a => a -> a -> a
+ Int -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
x)) [a]
coeff [(Int
1::Int)..]
          c0 :: a
c0 = a
0.999999999999997092
          -- constants from Numerical Recipes
          coeff :: [a]
coeff = [ a
57.1562356658629235
                  , -a
59.5979603554754912
                  , a
14.1360979747417471
                  , -a
0.491913816097620199
                  , a
0.339946499848118887e-4
                  , a
0.465236289270485756e-4
                  , -a
0.983744753048795646e-4
                  , a
0.158088703224912494e-3
                  , -a
0.210264441724104883e-3
                  , a
0.217439618115212643e-3
                  , -a
0.164318106536763890e-3
                  , a
0.844182239838527433e-4
                  , -a
0.261908384015814087e-4
                  , a
0.368991826595316234e-5
                  ]
          𝛾 :: a
𝛾 = a
607a -> a -> a
forall a. Fractional a => a -> a -> a
/a
128
          z' :: a
z' = a
za -> a -> a
forall a. Num a => a -> a -> a
-a
1
gammaln a
z = a -> a
forall a. Floating a => a -> a
log a
forall a. Floating a => a
pi a -> a -> a
forall a. Num a => a -> a -> a
- a -> a
forall a. Floating a => a -> a
log (a -> a
forall a. Floating a => a -> a
sin (a
forall a. Floating a => a
pi a -> a -> a
forall a. Num a => a -> a -> a
* a
z)) a -> a -> a
forall a. Num a => a -> a -> a
- a -> a
forall a. (Floating a, Ord a) => a -> a
gammaln (a
1 a -> a -> a
forall a. Num a => a -> a -> a
- a
z)