```{-# LANGUAGE FlexibleInstances, TemplateHaskell #-}
module Math.Gamma
( Gamma(..)
, Factorial(..)
, IncGamma(..)
) where

import Math.Gamma.Lanczos
import Math.Gamma.Incomplete
import Math.Factorial

import Data.Complex
import Data.List (findIndex)
import GHC.Float (float2Double, double2Float)
import qualified Data.Vector.Unboxed as V
import Math.ContinuedFraction
import Math.Sequence.Converge

-- |Gamma function.  Minimal definition is ether 'gamma' or 'lnGamma'.
class (Floating a, Factorial a) => Gamma a where
-- |The gamma function:  gamma z == integral from 0 to infinity of
-- @\t -> t**(z-1) * exp (negate t)@
gamma :: a -> a
gamma 0 = 0/0
gamma z
| z == abs z    = exp (lnGamma z)
| otherwise     = pi / (sin (pi * z) * exp (lnGamma (1-z)))

-- |Natural log of the gamma function
lnGamma :: a -> a
lnGamma z = log (gamma z)

-- |Natural log of the factorial function
lnFactorial :: Integral b => b -> a
lnFactorial n = lnGamma (fromIntegral n+1)

floatGammaInfCutoff :: Double
floatGammaInfCutoff = \$( do
let Just cutoff = findIndex isInfinite (scanl (*) (1::Float) [1..])
litE (IntegerL (1 + toInteger cutoff))
)

instance Gamma Float where
gamma = double2Float . gam . float2Double
where
gam x
| x >= floatGammaInfCutoff  = 1/0
| otherwise = case properFraction x of
(n,0) | n < 1     -> 0/0
| otherwise -> factorial (n-1 :: Integer)
_     | x < (-20) -> let s = pi / sin (pi * x)
in signum s * exp (log (abs s) - lnGamma (1-x))
| otherwise -> reflect (gammaLanczos g cs) x

g = pi
cs = [ 1.0000000249904433
, 9.100643759042066
,-4.3325519094475
, 0.12502459858901147
, 1.1378929685052916e-4
,-9.555011214455924e-5
]

lnGamma = double2Float . reflectLn (lnGammaLanczos g cs) . float2Double
where
g = pi
cs = [ 1.0000000249904433
, 9.100643759042066
,-4.3325519094475
, 0.12502459858901147
, 1.1378929685052916e-4
,-9.555011214455924e-5
]

lnFactorial n
| n' < 0                = error "lnFactorial n: n < 0"
| n' < toInteger nFacs  = facs V.! fromIntegral n
| otherwise             = lnGamma (fromIntegral n+1)
where
n' = toInteger n
nFacs       = 2000 -- limited only by time and space
facs        = V.map lnGamma (V.enumFromN 1 nFacs)

doubleGammaInfCutoff :: Double
doubleGammaInfCutoff = \$( do
let Just cutoff = findIndex isInfinite (scanl (*) (1::Double) [1..])
litE (IntegerL (1 + toInteger cutoff))
)

instance Gamma Double where
gamma x
| x >= doubleGammaInfCutoff  = 1/0
| otherwise = case properFraction x of
(n,0) | n < 1     -> 0/0
| otherwise -> factorial (n-1 :: Integer)
_     | x < (-50) -> let s = pi / sin (pi * x)
in signum s * exp (log (abs s) - lnGammaLanczos g cs (1-x))
| otherwise -> reflect (gammaLanczos g cs) x
where
g = 2*pi
cs = [   0.9999999999999858
, 311.6011750541472
,-498.6511904603639
, 244.08472899976877
, -38.670364643074194
,   1.3350900101370549
,  -1.8977221899565682e-3
,   8.475264614349149e-7
,   2.59715567376858e-7
,  -2.7166437850607517e-7
,   6.151114806136299e-8
]

lnGamma = reflectLn (lnGammaLanczos g cs)
where
g = exp pi / pi
cs = [    1.0000000000000002
, 1002.5049417114732
,-1999.6140446432912
, 1352.1626218340114
, -360.6486475548049
,   33.344988357090685
,    -0.6637188712004668
,     5.16644552377916e-4
,     1.684651140163429e-7
,    -1.8148207145896904e-7
,     6.171532716135051e-8
,    -9.014004881476154e-9
]

lnFactorial n
| n' < 0                = error "lnFactorial n: n < 0"
| n' < toInteger nFacs  = facs V.! fromIntegral n
| otherwise             = lnGamma (fromIntegral n+1)
where
n' = toInteger n
nFacs       = 2000 -- limited only by time and space
facs        = V.map lnGamma (V.enumFromN 1 nFacs)

complexDoubleToFloat :: Complex Double -> Complex Float
complexDoubleToFloat (a :+ b) = double2Float a :+ double2Float b
complexFloatToDouble :: Complex Float -> Complex Double
complexFloatToDouble (a :+ b) = float2Double a :+ float2Double b

instance Gamma (Complex Float) where
gamma = complexDoubleToFloat . gamma . complexFloatToDouble
lnGamma = complexDoubleToFloat . reflectLnC (lnGammaLanczos g cs) . complexFloatToDouble
where
g = pi
cs = [ 1.0000000249904433
, 9.100643759042066
,-4.3325519094475
, 0.12502459858901147
, 1.1378929685052916e-4
,-9.555011214455924e-5
]

lnFactorial n
| n' < 0                = error "lnFactorial n: n < 0"
| n' < toInteger nFacs  = facs V.! fromIntegral n
| otherwise             = lnGamma (fromIntegral n+1)
where
n' = toInteger n
nFacs       = 2000 -- limited only by time and space
facs        = V.map lnGamma (V.enumFromN 1 nFacs)

instance Gamma (Complex Double) where
gamma = reflectC (gammaLanczos g cs)
where
g = 2*pi
cs = [   1.0000000000000002
, 311.60117505414695
,-498.65119046033163
, 244.08472899875767
, -38.67036462939322
,   1.3350899103585203
,  -1.8972831806242229e-3
,  -3.935368195357295e-7
,   2.592464641764731e-6
,  -3.2263565156368265e-6
,   2.5666169886566876e-6
,  -1.3737776806198937e-6
,   4.4551204024819644e-7
,  -6.576826592057796e-8
]

lnGamma = reflectLnC (lnGammaLanczos g cs)
where
g = exp pi / pi
cs = [    1.0000000000000002
, 1002.5049417114732
,-1999.6140446432912
, 1352.1626218340114
, -360.6486475548049
,   33.344988357090685
,    -0.6637188712004668
,     5.16644552377916e-4
,     1.684651140163429e-7
,    -1.8148207145896904e-7
,     6.171532716135051e-8
,    -9.014004881476154e-9
]

lnFactorial n
| n' < 0                = error "lnFactorial n: n < 0"
| n' < toInteger nFacs  = facs V.! fromIntegral n
| otherwise             = lnGamma (fromIntegral n+1)
where
n' = toInteger n
nFacs       = 2000 -- limited only by time and space
facs        = V.map lnGamma (V.enumFromN 1 nFacs)

-- |Incomplete gamma functions.
class Gamma a => IncGamma a where
-- |Lower gamma function: lowerGamma s x == integral from 0 to x of
-- @\t -> t**(s-1) * exp (negate t)@
lowerGamma :: a -> a -> a
-- |Natural log of lower gamma function
lnLowerGamma :: a -> a -> a
-- |Regularized lower incomplete gamma function: lowerGamma s x / gamma s
p :: a -> a -> a

-- |Upper gamma function: lowerGamma s x == integral from x to infinity of
-- @\t -> t**(s-1) * exp (negate t)@
upperGamma :: a -> a -> a
-- |Natural log of upper gamma function
lnUpperGamma :: a -> a -> a
-- |Regularized upper incomplete gamma function: upperGamma s x / gamma s
q :: a -> a -> a

-- |This instance uses the Double instance.
instance IncGamma Float where
lowerGamma   s x = double2Float \$ (lowerGamma   :: Double -> Double -> Double) (float2Double s) (float2Double x)
lnLowerGamma s x = double2Float \$ (lnLowerGamma :: Double -> Double -> Double) (float2Double s) (float2Double x)
p s x = double2Float \$ (p :: Double -> Double -> Double) (float2Double s) (float2Double x)

upperGamma   s x = double2Float \$ (upperGamma   :: Double -> Double -> Double) (float2Double s) (float2Double x)
lnUpperGamma s x = double2Float \$ (lnUpperGamma :: Double -> Double -> Double) (float2Double s) (float2Double x)
q s x = double2Float \$ (q :: Double -> Double -> Double) (float2Double s) (float2Double x)

-- |I have not yet come up with a good strategy for evaluating these
-- functions for negative @x@.  They can be rather numerically unstable.
instance IncGamma Double where
lowerGamma s x
| x < 0     = error "lowerGamma: x < 0 is not currently supported."
| x == 0    = 0
| x >= s+1  = gamma s - upperGamma s x
| otherwise = lowerGammaHypGeom s x

upperGamma s x
| x < 0     = error "upperGamma: x < 0 is not currently supported."
| x == 0    = gamma s
| x < s+1   = q s x * gamma s
| otherwise = converge . concat
\$ modifiedLentz 1e-30 (upperGammaCF s x)

lnLowerGamma s x
| x < 0     = error "lnLowerGamma: x < 0 is not currently supported."
| x == 0    = log 0
| x >= s+1  = log (p s x) + lnGamma s
| otherwise = lnLowerGammaHypGeom s x

lnUpperGamma s x
| x < 0     = error "lnUpperGamma: x < 0 is not currently supported."
| x == 0    = lnGamma s
| x < s+1   = log (q s x) + lnGamma s
| otherwise =
converge (lnUpperGammaConvergents s x)

p s x
| x < 0     = error "p: x < 0 is not currently supported."
| x == 0    = 0
| x >= s+1  = 1 - q s x
| otherwise = pHypGeom s x

q s x
| x < 0     = error "q: x < 0 is not currently supported."
| x == 0    = 1
| x < s+1   = 1 - p s x
| otherwise =
converge . concat
\$ modifiedLentz 1e-30 (qCF s x)

```