#lang typed/racket/base

(require (only-in racket/math exact-truncate)
         "../../flonum.rkt"
         "../../base.rkt"
         "log-gamma-zeros.rkt"
         "gamma.rkt")

(provide fllog-gamma log-gamma)

(define +fllog-gamma-max.0 2.5563481638716906e+305)

(: fllog-gamma-stirling (Float -> Float))
;; Computes log(Gamma(x)) using 5 terms from Stirling's series
;; For x >= 143, relative error ε <= epsilon.0
(define (fllog-gamma-stirling x)
  (let* ([x  (fl- x 1.0)]
         [log-x  (fllog x)])
    (fl+ (fl+ (fl- (fl* x log-x) x)
              (fl* 0.5 (fl+ (fllog (fl* 2.0 pi)) log-x)))
         (let ([1/x  (fl/ 1.0 x)])
           (fl* 1/x (fl+ (fl* #i-1/360 (fl* 1/x 1/x)) #i1/12))))))

(: fllog-gamma-taylor-0.00 (Float -> Float))
;; Good for -0.25 <= x <= 0.25
(define (fllog-gamma-taylor-0.00 x)
  (fl- (fl* x ((make-flpolyfun
                (-5.7721566490153286060651209008240243104216e-1
                 +8.2246703342411321823620758332301259460947e-1
                 -4.0068563438653142846657938717048333025499e-1
                 +2.7058080842778454787900092413529197569369e-1
                 -2.0738555102867398526627309729140683361142e-1
                 +1.6955717699740818995241965496515342131697e-1
                 -1.4404989676884611811997107854997096565712e-1
                 +1.2550966952474304242233565481358155815737e-1
                 -1.1133426586956469049087252991471245116507e-1
                 +1.000994575127818085337145958900319017006e-1
                 -9.0954017145829042232609298411497266951692e-2
                 +8.3353840546109004024886499837311639246699e-2
                 -7.6932516411352191472827064348181338131867e-2
                 +7.1432946295361336059232753221795380981961e-2
                 -6.6668705882420468032903448567376337505842e-2
                 +6.250095514121304074198328571797729512645e-2
                 -5.882397865868458233895727060550370760077e-2
                 +5.5555767627403611102214247869145663318357e-2
                 -5.2631679379616660733627666155673426387014e-2
                 +5.0000047698101693639805657601934172467295e-2
                 -4.7619070330142227990783957939028779668856e-2
                 +4.5454556293204669442408636529463034227041e-2
                 -4.3478266053040259361351002947335603579408e-2
                 +4.1666669150341210469144983851675330658383e-2
                 -4.000000119214014058609120744254820277464e-2))
               x))
       (fllog (flabs x))))

(: fllog-gamma-taylor-0.25 (Float -> Float))
;; Good for 0.2 <= x <= 0.3
(define (fllog-gamma-taylor-0.25 x)
  ((make-flpolyfun
    (+1.2880225246980774573706104402197172959254e0
     -4.2274535333762654080895301460966835773672e0
     +8.5986645772535553696356595596676120107535e0
     -2.1554623322922820055556327863140664981215e1
     +6.4115922667049516500949635159543127365351e1
     -2.0486979490531611444631855960434965829803e2
     +6.8271182467733512784147775999936646633190e2
     -2.3406019221713673821639922031360341786831e3
     +8.1920211733386271770257584477642417071823e3
     -2.9127126102413786952705664233364266579143e4
     +1.0485761076831147539267450999973685838075e5
     -3.8130037145777091303332298313162618015895e5
     +1.3981013390649692882433574562368969426983e6
     -5.1622203119232470494907885079964686151354e6
     +1.9173961145999447781172798758783596423575e7
     -7.1582788269012640570224029670280449249173e7
     +2.6843545600175936387731944917825625449913e8
     -1.0105805402366187664874395906083166992963e9
     +3.8177487075565563809425864802719934216549e9
     -1.4467258260211284827499035133145761712036e10
     +5.4975581388800576465277096191993634564536e10
     -2.0943078624304805825772615550814087377792e11
     +7.9964482020072760812343178266260768051649e11
     -3.0595106164201741696941758669213590210283e12
     +1.1728124029610666863432083838677809746289e13
     -4.503599627370496015111579019136059236865e13))
   (fl- x 0.25)))

(: fllog-gamma-taylor-0.35 (Float -> Float))
;; Good for 0.3 <= x <= 0.4
(define (fllog-gamma-taylor-0.35 x)
  ((make-flpolyfun
    (+9.3458122714623255657034665561114707129281e-1
     -2.9710708698259454387463064442084878587473e0
     +4.6202295211029058483300084111499864142794e0
     -7.9556316522059548941376774196765046265487e0
     +1.6746588107729381853722502005240049590574e1
     -3.8127454700915928345673649988447464150541e1
     +9.0693847100977686331535034003589466561002e1
     -2.2205503134099860114590905607713594917160e2
     +5.5510435554655568572023613074757278601120e2
     -1.4097672131852369301046404032575154781671e3
     +3.6251013644124743041583769635697422453608e3
     -9.4158380862889149009699386717380613087566e3
     +2.4660521806577226426593494121580423161875e4
     -6.5038734381109325447389888110993862111771e4
     +1.7255174121793485037024340821203990106480e5
     -4.6013797446743480714344276040988335634682e5
     +1.2325124301416449829627359292354595578751e6
     -3.3143191388536834678869738614710180597041e6
     +8.9434008501593831003677839979897453897692e6
     -2.4207701548801412719812272885598523140688e7
     +6.5706618489250456194219732829388693984488e7
     -1.7879351969839140350109553939332746596063e8
     +4.8761869008634574265006998747622410016702e8
     -1.3326225070681491886298992486407039142238e9
     +3.6488473407817484073734090581475464119248e9
     -1.0008266991858446872669247506273021213142e10))
   (fl- x 0.35)))

(: fllog-gamma-taylor-0.50 (Float -> Float))
;; Good for 0.4 <= x <= 0.6
(define (fllog-gamma-taylor-0.50 x)
  ((make-flpolyfun
    (+5.7236494292470008707171367567652935582365e-1
     -1.9635100260214234794409763329987555671932e0
     +2.4674011002723396547086227499690377838284e0
     -2.8047994407057199992660557101933833117850e0
     +4.0587121264167682181850138620293796354053e0
     -6.4289520818888935432544660160336118419539e0
     +1.0682102150836715967002438262804665542969e1
     -1.8294336889643457001236326975846312638455e1
     +3.2004965728809475817695591977463297330129e1
     -5.6891809859347556840835862786418062545350e1
     +1.0240174503557579012999003159550263543971e2
     -1.8618287309751204945015123384833490545011e2
     +3.4133397703631637148191021683379116271523e2
     -6.3015424192538580035392648407595334063812e2
     +1.1702859591569047686584101960326737266275e3
     -2.1845334856492714760341472992072204510539e3
     +4.0960000951793966250258746295276420361119e3
     -7.7101177067724468917494684155339764589405e3
     +1.4563555593150464825167749579161451619264e4
     -2.7594105286901080206051448205759553700168e4
     +5.2428800015036983408359217419948124894894e4
     -9.9864380961928103573100568175792144311320e4
     +1.9065018182425722465630287142143639284758e5
     -3.6472208696039589096044455278104007553545e5    
     +6.9905066666914181140109626025108513265177e5
     -1.3421772800015840896262969798376701341757e6))
   (fl- x 0.5)))

(: fllog-gamma-taylor-0.75 (Float -> Float))
;; Good for 0.55 <= x <= 0.95
(define (fllog-gamma-taylor-0.75 x)
  ((make-flpolyfun
    (+2.0328095143129537148143297186242969975967e-1
     -1.0858608797864721696268867628171806931701e0
     +1.2709398238358032491988314402085391245601e0
     -8.8377220272293993857211781840640151306496e-1
     +8.2347135561877499001058663292694680113391e-1
     -8.5667171512847893782435290872592064449950e-1
     +9.4271297621469404667828882013212841811724e-1
     -1.0731997029951139942576497722938390391533e0
     +1.2500532365986323043130984663624093306882e0
     -1.4805455721621498022975132817814440767646e0
     +1.7761481181337004352823540578403095184479e0
     -2.1526459337642405867437897637062028968286e0
     +2.6308757825693465467919143116597833317597e0
     -3.2379295134270085772495537411513722635691e0
     +4.0088272799485265938930157299134890960034e0
     -4.9887426871564629090301417504909640897591e0
     +6.2359175733403762705424672939794981911256e0
     -7.8254593925079467335525526631300376066016e0
     +9.8542790701861880654077061516254611016851e0
     -1.2447508711571266543755503178318247050213e1
     +1.5766843449126578573345890445253001140810e1
     -2.0021388006808847242857973248859918171351e1
     +2.5481766281397976151808667749688257328831e1
     -3.2498484383736992094096004929199618085860e1
     +4.1525841075370756346859758598306419016717e1
     -5.3153076531697230407881186981792953854319e1))
   (fl- x 0.75)))

(: fllog-gamma-taylor-1.00 (Float -> Float))
;; Good for 0.75 <= x <= 1.25
(define (fllog-gamma-taylor-1.00 x)
  (fl* (fl- x 1.0)
       ((make-flpolyfun
         (-5.7721566490153286060651209008240243104216e-1
          +8.2246703342411321823620758332301259460947e-1
          -4.0068563438653142846657938717048333025499e-1
          +2.7058080842778454787900092413529197569369e-1
          -2.0738555102867398526627309729140683361142e-1
          +1.6955717699740818995241965496515342131697e-1
          -1.4404989676884611811997107854997096565712e-1
          +1.2550966952474304242233565481358155815737e-1
          -1.1133426586956469049087252991471245116507e-1
          +1.0009945751278180853371459589003190170060e-1
          -9.0954017145829042232609298411497266951692e-2
          +8.3353840546109004024886499837311639246699e-2
          -7.6932516411352191472827064348181338131867e-2
          +7.1432946295361336059232753221795380981961e-2
          -6.6668705882420468032903448567376337505842e-2
          +6.2500955141213040741983285717977295126450e-2
          -5.8823978658684582338957270605503707600770e-2
          +5.5555767627403611102214247869145663318357e-2
          -5.2631679379616660733627666155673426387014e-2
          +5.0000047698101693639805657601934172467295e-2
          -4.7619070330142227990783957939028779668856e-2
          +4.5454556293204669442408636529463034227041e-2
          -4.3478266053040259361351002947335603579408e-2
          +4.1666669150341210469144983851675330658383e-2
          -4.0000001192140140586091207442548202774640e-2))
        (fl- x 1.0))))

(: fllog-gamma-taylor-1.50 (Float -> Float))
;; Good for 1.15 <= x <= 1.9
(define (fllog-gamma-taylor-1.50 x)
  ((make-flpolyfun
    (-1.2078223763524522234551844578164721225185e-1
     +3.6489973978576520559023667001244432806837e-2
     +4.6740110027233965470862274996903778382841e-1
     -1.3813277403905333259938904352671664511829e-1
     +5.8712126416768218185013862029379635405284e-2
     -2.8952081888893543254466016033611841953893e-2
     +1.5435484170049300335771596137998876302295e-2
     -8.6226039291712869506126901320269241689428e-3
     +4.9657288094758176955919774632973301288730e-3
     -2.9209704586679519469738975291736564606011e-3
     +1.7450355757901299900315955026354397136770e-3
     -1.0549156938676319694156665167236319325019e-3
     +6.4370298303814857688350045782938190040240e-4
     -3.9577153964650777263792210718679196822994e-4
     +2.4487119048294412448174695944091318291752e-4
     -1.5231593814270081396587388711772058669604e-4
     +9.5179396625025874629527642036111920748339e-5
     -5.9713623362337703709651623517764060417344e-5
     +3.7594909269612194023605896063708381243259e-5
     -2.3743185469209342942601658963326064219265e-5
     +1.5036983408359217419948124894894133810908e-5
     -9.5471511921481872234111919303675575211366e-6
     +6.0754064744846896032545746657610345645862e-6
     -3.8741518300097701723444233615366600085264e-6
     +2.4751447344295935844184659851074357373265e-6
     -1.5840896262969798376701341756813850557999e-6))
   (fl- x 1.5)))

(: fllog-gamma-taylor-2.00 (Float -> Float))
;; Good for 1.6 <= x <= 2.55
(define (fllog-gamma-taylor-2.00 x)
  (fl* (fl- x 2.0)
       ((make-flpolyfun
         (+4.2278433509846713939348790991759756895784e-1
          +3.2246703342411321823620758332301259460947e-1
          -6.7352301053198095133246053837149996921661e-2
          +2.0580808427784547879000924135291975693686e-2
          -7.3855510286739852662730972914068336114162e-3
          +2.8905103307415232857529882984867546503010e-3
          -1.1927539117032609771139356928281085142662e-3
          +5.0966952474304242233565481358155815736809e-4
          -2.2315475845357937976141880360134005395629e-4
          +9.9457512781808533714595890031901700599964e-5
          -4.4926236738133141700207502406357860783221e-5
          +2.0507212775670691553166503978305913365992e-5
          -9.4394882752683959039874251044150549438348e-6
          +4.3748667899074878041817932239524105325324e-6
          -2.0392157538013662367819007096708391751769e-6
          +9.5514121304074198328571797729512645039176e-7
          -4.4924691987645660432942903311936547590378e-7
          +2.1207184805554665869231359010776280096232e-7
          -1.0043224823968099608720830500533438202816e-7
          +4.7698101693639805657601934172467295268577e-8
          -2.2711094608943164910319981160621236400060e-8
          +1.0838659214896954091074917579681586085324e-8
          -5.1834750419700466551212486470576691703107e-9
          +2.4836745438024783171850086639917157755945e-9
          -1.1921401405860912074425482027746404023922e-9))
        (fl- x 2.0))))

(: fllog-gamma-taylor-3.00 (Float -> Float))
;; Good for 2.2 <= x <= 3.85
(define (fllog-gamma-taylor-3.00 x)
  ((make-flpolyfun
    (+6.9314718055994530941723212145817656807550e-1
     +9.2278433509846713939348790991759756895784e-1
     +1.9746703342411321823620758332301259460947e-1
     -2.5685634386531428466579387170483330254994e-2
     +4.9558084277845478790009241352919756936856e-3
     -1.1355510286739852662730972914068336114162e-3
     +2.8634366407485661908632163182008798363431e-4
     -7.6682483131832405685364264256679942837650e-5
     +2.1388274743042422335654813581558157368093e-5
     -6.1408695646904908725299147124511650673996e-6
     +1.8012627818085337145958900319017005999641e-6
     -5.3703219267859624566204786090331532867543e-7
     +1.6216069233735821983317064497258003265830e-7
     -4.9464236806857442448963565953516482296279e-8
     +1.5212772050344947038936081095267675389550e-8
     -4.7105454680329034485673763375058418435462e-9
     +1.4668966344919832857179772951264503917565e-9
     -4.5900627351542785884079782524782884495948e-10
     +1.4422218749110313675803455220724540676060e-10
     -4.5478091654680297734620794808066238685290e-11
     +1.4385873327305657601934172467295268576888e-11
     -4.5632659372125293676002082402840191072311e-12
     +1.4510739168404547112812160452224489605314e-12
     -4.6245280564448120820516879679960509329496e-13
     +1.4767816120227551834199732504910892780276e-13
     -4.7245078278707442548202774640402392180051e-14))
   (fl- x 3.0)))

(: fllog-gamma-taylor-4.00 (Float -> Float))
;; Good for 2.95 <= x <= 5.2
(define (fllog-gamma-taylor-4.00 x)
  ((make-flpolyfun
    (+1.7917594692280550008124773583807022727230e0
     +1.2561176684318004727268212432509309022912e0
     +1.4191147786855766268065202776745703905392e-1
     -1.3339955374185749454233708158137651242649e-2
     +1.8693886746981281259145043822055559405991e-3
     -3.1250576118427333211671869058378834392645e-4
     +5.7719978661047748487327576035908742665064e-5
     -1.1361430156458442657080248318343016846479e-5
     +2.3363009585583497857386422662098872873405e-6
     -4.9584029521372863551771580419982578412385e-7
     +1.0775400096550504349223035942629881510428e-7
     -2.3847713635254224115484323789557212087699e-8
     +5.3543237407814910272762308544872747942770e-9
     -1.2161233925261412779191308402571721339304e-9
     +2.7883218400430620027947022687788898118330e-10
     -6.4430842949148520985319622895686038139193e-11
     +1.4985814153309870848584571810776705153549e-11
     -3.5048396836871796660863826389900181909766e-12
     +8.2358795148107109192169064520359062021452e-13
     -1.9432337901017173479568063057928609389652e-13
     +4.6013373343451035322886311091085999887070e-14
     -1.0929443891193949628371365298531818267889e-14
     +2.6032144200298055629043122724648881160952e-15
     -6.2156719159210783121433122139932582601310e-16
     +1.4873780760261127809520410055697145662790e-16
     -3.5662792412112287754095936864694572751781e-17))
   (fl- x 4.0)))

(: fllog-gamma-integer (Float -> Float))
(define (fllog-gamma-integer x)
  (cond [(x . fl<= . 0.0)  +inf.0]
        [else  (fllog-factorial (fl- x 1.0))]))

(: fllog-gamma-large-negative (Float -> Float))
(define (fllog-gamma-large-negative x)
  (cond [(x . fl< . -170.0)  (fl- (fl- (fllog pi) (fllog (flabs (fl* x (flsinpix x)))))
                                  (fllog-gamma (- x)))]
        [else  (fllog (flabs (flgamma x)))]))

(: fllog-gamma-small-positive (Float -> Float))
;; Good for 0 <= x <= 5
(define (fllog-gamma-small-positive x)
  (cond [(x . fl< . 0.85)
         (cond [(x . fl< . 0.2)   (fllog-gamma-taylor-0.00 x)]
               [(x . fl< . 0.3)   (fllog-gamma-taylor-0.25 x)]
               [(x . fl< . 0.4)   (fllog-gamma-taylor-0.35 x)]
               [(x . fl< . 0.6)   (fllog-gamma-taylor-0.50 x)]
               [else              (fllog-gamma-taylor-0.75 x)])]
        [else
         (cond [(x . fl< . 1.20)  (fllog-gamma-taylor-1.00 x)]
               [(x . fl< . 1.75)  (fllog-gamma-taylor-1.50 x)]
               [(x . fl< . 2.35)  (fllog-gamma-taylor-2.00 x)]
               [(x . fl< . 3.4)   (fllog-gamma-taylor-3.00 x)]
               [else              (fllog-gamma-taylor-4.00 x)])]))

(: fllog-gamma-small-negative (Float -> Float))
(define (fllog-gamma-small-negative x)
  (cond [(x . fl> . -0.25)  (fllog-gamma-taylor-0.00 x)]
        [else  (fllog (flabs (flgamma x)))]))

(: fllog-gamma-large-positive (Float -> Float))
(define (fllog-gamma-large-positive x)
  (cond [(x . fl< . 150.0)  (fllog (flgamma x))]
        [else  (fllog-gamma-stirling x)]))

(define: fllog-gamma-hash : (HashTable Float Float) (make-weak-hash))

(: fllog-gamma (Float -> Float))
(define (fllog-gamma x)
  (cond [(integer? x)  (fllog-gamma-integer x)]
        [(x . fl> . +fllog-gamma-max.0)  +inf.0]
        [(x . fl= . -inf.0)  +inf.0]
        [(eqv? x +nan.0)     +nan.0]
        [else
         (hash-ref!
          fllog-gamma-hash x
          (λ ()
            (cond [(x . fl< . 0.0)
                   (define y (fllog-gamma-special-negative x))
                   (cond [(not (fl= y 0.0))  y]
                         [(x . fl< . -5.5)  (fllog-gamma-large-negative x)]
                         [else  (fllog-gamma-small-negative x)])]
                  [(x . fl< . 4.5)  (fllog-gamma-small-positive x)]
                  [else  (fllog-gamma-large-positive x)])))]))

(: log-gamma (case-> (One -> Zero)
                     (Flonum -> Flonum)
                     (Real -> (U Zero Flonum))))
(define (log-gamma x)
  (cond [(flonum? x)  (fllog-gamma x)]
        [(single-flonum? x)  (fllog-gamma (fl x))]
        [(integer? x)
         (cond [(x . <= . 0)
                (raise-argument-error 'log-gamma "Real, not Zero or Negative-Integer" x)]
               [(eqv? x 1)  0]
               [else  (fllog-factorial (fl (- x 1)))])]
        [else  (fllog-gamma (fl x))]))