module Math.SpecialFunction ( incbeta
                            , beta
                            , regbeta
                            , bessel1
                            , gamma
                            , gammaln
                            , agm
                            , completeElliptic
                            , fcdf
                            , chisqcdf
                            , tcdf
                            ) where

import           Math.Hypergeometric

{-# SPECIALIZE tcdf :: Double -> Double -> Double #-}
{-# SPECIALIZE tcdf :: Float -> Float -> Float #-}
-- | Converges if and only if \(|x| \leq \sqrt{\nu} \)
--
-- @since 0.1.2.0
tcdf :: (Floating a, Ord a)
     => a -- ^ \(\nu\) (degrees of freedom)
     -> a -- ^ \(x\)
     -> a
tcdf :: forall a. (Floating a, Ord a) => a -> a -> a
tcdf a
𝜈 a
x = a
0.5 a -> a -> a
forall a. Num a => a -> a -> a
+ a
x a -> a -> a
forall a. Num a => a -> a -> a
* a -> a
forall a. (Floating a, Ord a) => a -> a
gamma (a
0.5a -> a -> a
forall a. Num a => a -> a -> a
*(a
𝜈a -> a -> a
forall a. Num a => a -> a -> a
+a
1)) a -> a -> a
forall a. Fractional a => a -> a -> a
/ (a -> a
forall a. Floating a => a -> a
sqrt(a
forall a. Floating a => a
pia -> 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, Ord a) => a -> a
gamma(a
𝜈a -> 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. (Ord a, Fractional a) => [a] -> [a] -> a -> a
hypergeometric [a
0.5, a
0.5a -> a -> a
forall a. Num a => a -> a -> a
*(a
𝜈a -> a -> a
forall a. Num a => a -> a -> a
+a
1)] [a
1.5] (-(a
xa -> Int -> a
forall a b. (Num a, Integral b) => a -> b -> a
^(Int
2::Int))a -> a -> a
forall a. Fractional a => a -> a -> a
/a
𝜈)

{-# SPECIALIZE bessel1 :: Double -> Double -> Double #-}
{-# SPECIALIZE bessel1 :: Float -> Float -> Float #-}
-- | Bessel functions of the first kind, \( J_\alpha(x)\).
--
-- @since 0.1.3.0
bessel1 :: (Floating a, Ord a)
        => a -- ^ \(\alpha\)
        -> a -- ^ \(x\)
        -> a
bessel1 :: forall a. (Floating a, Ord a) => a -> a -> a
bessel1 a
𝛼 a
x = ((a
xa -> a -> a
forall a. Fractional a => a -> a -> a
/a
2)a -> a -> a
forall a. Floating a => a -> a -> a
**a
𝛼a -> a -> a
forall a. Fractional a => a -> a -> a
/a -> a
forall a. (Floating a, Ord a) => a -> a
gamma(a
𝛼a -> a -> a
forall a. Num a => a -> a -> a
+a
1))a -> a -> a
forall a. Num a => a -> a -> a
*[a] -> [a] -> a -> a
forall a. (Ord a, Fractional a) => [a] -> [a] -> a -> a
hypergeometric [] [a
𝛼a -> a -> a
forall a. Num a => a -> a -> a
+a
1] (-(a
xa -> Int -> a
forall a b. (Num a, Integral b) => a -> b -> a
^(Int
2::Int))a -> a -> a
forall a. Fractional a => a -> a -> a
/a
4)

{-# SPECIALIZE completeElliptic :: Double -> Double #-}
{-# SPECIALIZE completeElliptic :: Float -> Float #-}
-- | [Complete elliptic integral of the first kind](https://mathworld.wolfram.com/CompleteEllipticIntegraloftheFirstKind.html)
--
-- @since 0.1.4.0
completeElliptic :: (Ord a, Floating a) => a -> a
completeElliptic :: forall a. (Ord a, Floating a) => a -> a
completeElliptic a
k = a
forall a. Floating a => a
pia -> a -> a
forall a. Fractional a => a -> a -> a
/(a
2a -> a -> a
forall a. Num a => a -> a -> a
*a -> a -> a
forall a. (Ord a, Floating a) => a -> a -> a
agm a
1 (a -> a
forall a. Floating a => a -> a
sqrt (a
1a -> a -> a
forall a. Num a => a -> a -> a
-a
ka -> Int -> a
forall a b. (Num a, Integral b) => a -> b -> a
^(Int
2::Int))))

{-# SPECIALIZE agm :: Double -> Double -> Double #-}
{-# SPECIALIZE agm :: Float -> Float -> Float #-}
-- | [Arithmetic-geometric mean](https://mathworld.wolfram.com/Arithmetic-GeometricMean.html)
--
-- @since 0.1.4.0
agm :: (Ord a, Floating a) => a -> a -> a
agm :: forall a. (Ord a, Floating a) => a -> a -> a
agm a
a a
b =
    let a' :: a
a' = (a
aa -> a -> a
forall a. Num a => a -> a -> a
+a
b)a -> a -> a
forall a. Fractional a => a -> a -> a
/a
2
        b' :: a
b' = a -> a
forall a. Floating a => a -> a
sqrt(a
aa -> a -> a
forall a. Num a => a -> a -> a
*a
b)
    in if (a
a'a -> a -> a
forall a. Num a => a -> a -> a
-a
b')a -> a -> a
forall a. Fractional a => a -> a -> a
/a
b'a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<a
1e-15 Bool -> Bool -> Bool
&& (a
b'a -> a -> a
forall a. Num a => a -> a -> a
-a
a')a -> a -> a
forall a. Fractional a => a -> a -> a
/a
b'a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<a
1e-15
        then a
a'
        else a -> a -> a
forall a. (Ord a, Floating a) => a -> a -> a
agm a
a' a
b'

-- | @since 0.1.2.0
{-# SPECIALIZE chisqcdf :: Double -> Double -> Double #-}
{-# SPECIALIZE chisqcdf :: Float -> Float -> Float #-}
chisqcdf :: (Floating a, Ord a)
         => a -- ^ \(r\) (degrees of freedom)
         -> a -- ^ \(\chi^2\)
         -> a
chisqcdf :: forall a. (Floating a, Ord a) => a -> a -> a
chisqcdf a
r a
x = a -> a -> a
forall a. (Floating a, Ord a) => a -> a -> a
incgamma (a
0.5a -> a -> a
forall a. Num a => a -> a -> a
*a
r) (a
0.5a -> a -> a
forall a. Num a => a -> a -> a
*a
x) a -> a -> a
forall a. Fractional a => a -> a -> a
/ a -> a
forall a. (Floating a, Ord a) => a -> a
gamma (a
0.5a -> a -> a
forall a. Num a => a -> a -> a
*a
r)

{-# SPECIALIZE incgamma :: Double -> Double -> Double #-}
{-# SPECIALIZE incgamma :: Float -> Float -> Float #-}
-- | \(a^{-1}x^a{}_1F_1(a;1+a;-x) \)
incgamma :: (Floating a, Ord a) => a -> a -> a
incgamma :: forall a. (Floating a, Ord a) => a -> a -> a
incgamma a
a a
x = (a
1a -> a -> a
forall a. Fractional a => a -> a -> a
/a
a) a -> a -> a
forall a. Num a => a -> a -> a
* a
x a -> a -> a
forall a. Floating a => a -> a -> a
** a
a a -> a -> a
forall a. Num a => a -> a -> a
* [a] -> [a] -> a -> a
forall a. (Ord a, Fractional a) => [a] -> [a] -> a -> a
hypergeometric [a
a] [a
1a -> a -> a
forall a. Num a => a -> a -> a
+a
a] (-a
x)
-- chisqcdf 10 28 works better this way than w/ e^-x ... x

-- (1 2 H. _1.1) 1 hangs indefinitely

{-# SPECIALIZE incbeta :: Double -> Double -> Double -> Double #-}
{-# SPECIALIZE incbeta :: Float -> Float -> Float -> Float #-}
-- | Incomplete beta function, \(|z|\leq 1\)
--
-- 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, Ord a)
        => a -- ^ \(z\)
        -> a -- ^ \(a\)
        -> a -- ^ \(b\)
        -> a
incbeta :: forall a. (Floating a, Ord 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. (Ord 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 regbeta :: Double -> Double -> Double -> Double #-}
{-# SPECIALIZE regbeta :: Float -> Float -> Float -> Float #-}
-- | [Regularized beta function](https://mathworld.wolfram.com/RegularizedBetaFunction.html), \(|z|\leq 1\)
--
-- \(I(z;a,b) = \displaystyle\frac{B(z;a,b)}{B(a,b)}\)
regbeta :: (Floating a, Ord a) 
        => a -- ^ \(z\)
        -> a -- ^ \(a\)
        -> a -- ^ \(b\)
        -> a
regbeta :: forall a. (Floating a, Ord a) => a -> a -> a -> a
regbeta a
z a
a a
b = a -> a -> a -> a
forall a. (Floating a, Ord a) => a -> a -> a -> a
incbeta a
z a
a a
b a -> a -> a
forall a. Fractional a => a -> a -> a
/ a -> a -> a
forall a. (Floating a, Ord a) => a -> a -> a
beta a
a a
b

{-# SPECIALIZE fcdf :: Double -> Double -> Double -> Double #-}
{-# SPECIALIZE fcdf :: Float -> Float -> Float -> Float #-}
-- | @since 0.1.2.0
fcdf :: (Floating a, Ord a)
     => a -- ^ \(n\)
     -> a -- ^ \(m\)
     -> a -- ^ \(x\)
     -> a
fcdf :: forall a. (Floating a, Ord a) => a -> a -> a -> a
fcdf a
n a
m a
x = a -> a -> a -> a
forall a. (Floating a, Ord a) => a -> a -> a -> a
regbeta (a
n a -> a -> a
forall a. Num a => a -> a -> a
* a
x a -> a -> a
forall a. Fractional a => a -> a -> a
/ (a
m a -> a -> a
forall a. Num a => a -> a -> a
+ a
n a -> a -> a
forall a. Num a => a -> a -> a
* a
x)) (a
0.5 a -> a -> a
forall a. Num a => a -> a -> a
* a
n) (a
0.5 a -> a -> a
forall a. Num a => a -> a -> a
* a
m) -- we can use hypergeo because nx/(m+nx) < 1

{-# SPECIALIZE beta :: Double -> Double -> Double #-}
{-# SPECIALIZE beta :: Float -> Float -> Float #-}
-- | \(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 #-}
{-# SPECIALIZE betaln :: Float -> Float -> Float #-}
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

{-# SPECIALIZE gammaln :: Double -> Double #-}
{-# SPECIALIZE gammaln :: Float -> Float #-}
-- | \(\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 a. Num 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)