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 𝜈 x = 0.5 + x * gamma (0.5*(𝜈+1)) / (sqrt(pi*𝜈) * gamma(𝜈/2)) * hypergeometric [0.5, 0.5*(𝜈+1)] [1.5] (-x^(2::Int)/𝜈)

{-# 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 𝛼 x = ((x/2)**𝛼/gamma(𝛼+1))*hypergeometric [] [𝛼+1] (-(x^(2::Int))/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 k = pi/(2*agm 1 (sqrt (1-k^(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 a b =
    let a' = (a+b)/2
        b' = sqrt(a*b)
    in if (a'-b')/b'<1e-15 && (b'-a')/b'<1e-15
        then a'
        else agm 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 r x = incgamma (0.5*r) (0.5*x) / gamma (0.5*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 a x = (1/a) * x ** a * hypergeometric [a] [1+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 z a b = z**a/a * hypergeometric [a,1-b] [a+1] 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 z a b = incbeta z a b / beta 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 n m x = regbeta (n * x / (m + n * x)) (0.5 * n) (0.5 * 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 x y = exp (betaln x y)

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

{-# 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 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)