module Math.SpecialFunction ( incbeta
                            , beta
                            , 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| < \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 forall a. Num a => a -> a -> a
+ a
x forall a. Num a => a -> a -> a
* forall a. (Floating a, Ord a) => a -> a
gamma (a
0.5forall a. Num a => a -> a -> a
*(a
𝜈forall a. Num a => a -> a -> a
+a
1)) forall a. Fractional a => a -> a -> a
/ (forall a. Floating a => a -> a
sqrt(forall a. Floating a => a
piforall a. Num a => a -> a -> a
*a
𝜈) forall a. Num a => a -> a -> a
* forall a. (Floating a, Ord a) => a -> a
gamma(a
𝜈forall a. Fractional a => a -> a -> a
/a
2)) forall a. Num a => a -> a -> a
* forall a. (Eq a, Fractional a) => [a] -> [a] -> a -> a
hypergeometric [a
0.5, a
0.5forall a. Num a => a -> a -> a
*(a
𝜈forall a. Num a => a -> a -> a
+a
1)] [a
1.5] (-a
xforall a b. (Num a, Integral b) => a -> b -> a
^(Int
2::Int)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
xforall a. Fractional a => a -> a -> a
/a
2)forall a. Floating a => a -> a -> a
**a
𝛼forall a. Fractional a => a -> a -> a
/forall a. (Floating a, Ord a) => a -> a
gamma(a
𝛼forall a. Num a => a -> a -> a
+a
1))forall a. Num a => a -> a -> a
*forall a. (Eq a, Fractional a) => [a] -> [a] -> a -> a
hypergeometric [] [a
𝛼forall a. Num a => a -> a -> a
+a
1] (-(a
xforall a b. (Num a, Integral b) => a -> b -> a
^(Int
2::Int))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 = forall a. Floating a => a
piforall a. Fractional a => a -> a -> a
/(a
2forall a. Num a => a -> a -> a
*forall a. (Ord a, Floating a) => a -> a -> a
agm a
1 (forall a. Floating a => a -> a
sqrt (a
1forall a. Num a => a -> a -> a
-a
kforall 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
aforall a. Num a => a -> a -> a
+a
b)forall a. Fractional a => a -> a -> a
/a
2
        b' :: a
b' = forall a. Floating a => a -> a
sqrt(a
aforall a. Num a => a -> a -> a
*a
b)
    in if (a
a'forall a. Num a => a -> a -> a
-a
b')forall a. Fractional a => a -> a -> a
/a
b'forall a. Ord a => a -> a -> Bool
<a
1e-15 Bool -> Bool -> Bool
&& (a
b'forall a. Num a => a -> a -> a
-a
a')forall a. Fractional a => a -> a -> a
/a
b'forall a. Ord a => a -> a -> Bool
<a
1e-15
        then a
a'
        else 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 = forall a. (Floating a, Ord a) => a -> a -> a
incgamma (a
0.5forall a. Num a => a -> a -> a
*a
r) (a
0.5forall a. Num a => a -> a -> a
*a
x) forall a. Fractional a => a -> a -> a
/ forall a. (Floating a, Ord a) => a -> a
gamma (a
0.5forall 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
1forall a. Fractional a => a -> a -> a
/a
a) forall a. Num a => a -> a -> a
* a
x forall a. Floating a => a -> a -> a
** a
a forall a. Num a => a -> a -> a
* forall a. (Eq a, Fractional a) => [a] -> [a] -> a -> a
hypergeometric [a
a] [a
1forall a. Num a => a -> a -> a
+a
a] (-a
x)
-- TODO: writeup?
--
-- 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|<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
zforall a. Floating a => a -> a -> a
**a
aforall a. Fractional a => a -> a -> a
/a
a forall a. Num a => a -> a -> a
* forall a. (Eq a, Fractional a) => [a] -> [a] -> a -> a
hypergeometric [a
a,a
1forall a. Num a => a -> a -> a
-a
b] [a
aforall a. Num a => a -> a -> a
+a
1] a
z

{-# SPECIALIZE regbeta :: Double -> Double -> Double -> Double #-}
{-# SPECIALIZE regbeta :: Float -> Float -> Float -> Float #-}
-- | \(I(z;a,b) = \displaystyle\frac{B(z;a,b)}{B(a,b)}\)
regbeta :: (Floating a, Ord a) => a -> a -> a -> a
regbeta :: forall a. (Floating a, Ord a) => a -> a -> a -> a
regbeta a
z a
a a
b = forall a. (Floating a, Ord a) => a -> a -> a -> a
incbeta a
z a
a a
b forall a. Fractional 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 = forall a. (Floating a, Ord a) => a -> a -> a -> a
regbeta (a
n forall a. Num a => a -> a -> a
* a
x forall a. Fractional a => a -> a -> a
/ (a
m forall a. Num a => a -> a -> a
+ a
n forall a. Num a => a -> a -> a
* a
x)) (a
0.5 forall a. Num a => a -> a -> a
* a
n) (a
0.5 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 = forall a. Floating a => a -> a
exp (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 = forall a. (Floating a, Ord a) => a -> a
gammaln a
x forall a. Num a => a -> a -> a
+ forall a. (Floating a, Ord a) => a -> a
gammaln a
y forall a. Num a => a -> a -> a
- forall a. (Floating a, Ord a) => a -> a
gammaln (a
xforall 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 = forall a. Floating a => a -> a
exp forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 = -forall a. Floating a => a -> a
log a
0
gammaln a
z | a
z forall a. Ord a => a -> a -> Bool
>= a
0.5 = (a
z' forall a. Num a => a -> a -> a
+ a
1forall a. Fractional a => a -> a -> a
/a
2) forall a. Num a => a -> a -> a
* forall a. Floating a => a -> a
log (a
z' forall a. Num a => a -> a -> a
+ a
𝛾 forall a. Num a => a -> a -> a
+ a
1forall a. Fractional a => a -> a -> a
/a
2) forall a. Num a => a -> a -> a
- (a
z' forall a. Num a => a -> a -> a
+ a
1forall a. Fractional a => a -> a -> a
/a
2 forall a. Num a => a -> a -> a
+ a
𝛾) forall a. Num a => a -> a -> a
+ forall a. Floating a => a -> a
log (forall a. Floating a => a -> a
sqrt (a
2forall a. Num a => a -> a -> a
*forall a. Floating a => a
pi) forall a. Num a => a -> a -> a
* (a
c0 forall a. Num a => a -> a -> a
+ forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [a]
series))
    where series :: [a]
series = forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\a
c Int
x -> a
c forall a. Fractional a => a -> a -> a
/ (a
z' forall a. Num a => a -> a -> 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
607forall a. Fractional a => a -> a -> a
/a
128
          z' :: a
z' = a
zforall a. Num a => a -> a -> a
-a
1
gammaln a
z = forall a. Floating a => a -> a
log forall a. Floating a => a
pi forall a. Num a => a -> a -> a
- forall a. Floating a => a -> a
log (forall a. Floating a => a -> a
sin (forall a. Floating a => a
pi forall a. Num a => a -> a -> a
* a
z)) forall a. Num a => a -> a -> a
- forall a. (Floating a, Ord a) => a -> a
gammaln (a
1 forall a. Num a => a -> a -> a
- a
z)