#lang typed/racket/base

(require "../../flonum.rkt"
         "../../base.rkt"
         "../polynomial/chebyshev.rkt"
         "../number-theory/bernoulli.rkt"
         "gamma.rkt"
         "hurwitz-zeta.rkt"
         "tan-diff.rkt")

(provide flpsi0 flpsi psi0 psi)

(: flpsi0-huge (Flonum -> Flonum))
;; Compute using first two terms from asymptotic expansion
;; Error <= 1 ulp for x >= 1e7
(define (flpsi0-huge x)
  (fl- (fllog x) (fl/ 0.5 x)))

(: flpsi0-large (Flonum -> Flonum))
;; Compute using more terms from asymptotic expansion
;; Error <= 1 ulp for x >= 44
(define (flpsi0-large x)
  (define 1/x (/ 1.0 x))
  (define 1/x^2 (* 1/x 1/x))
  (+ (fllog x)
     (let* ([y  (* 1/x^2 #i-1/252)]
            [y  (* 1/x^2 (+ #i1/120 y))]
            [y  (* 1/x (+ #i-1/12 y))])
       (* 1/x (+ -0.5 y)))))

(: flpsi0-15-44 (Flonum -> Flonum))
;; Compute using Chebyshev fit of transformed function
;; Error <= 1 ulp for 15 <= x <= 44
(define (flpsi0-15-44 x)
  (define y
    ((inline-chebyshev-flpoly-fun
      15.0 44.0
      (-2.651473128433205899236218386166733001125e6
       -1.762146136898477454946670140604454842406e6
       -5.717880156460911821842439315667552305722e5
       -8.999795051532912396661010381758994001778e4
       -5.525632813934147220013901752731996650924e3
       +3.730714289863418037571565563246348880127e-7
       -9.689086103062769553596723282091729538813e-8
       +2.512217593011105663102858605710866764550e-8
       -6.502769240064928944990050767511278012820e-9
       +1.679934847962426417907078746025458161173e-9
       -4.316007423661544654578361918157615904084e-10
       +1.042642491290735190419839693626327816406e-10))
     x))
  (+ (/ y (flexpt x 5.0)) (fllog x) (/ 0.5 x)))

(: flpsi0-6-15 (Flonum -> Flonum))
;; Error <= 1 ulp for 6 <= x <= 15
(define (flpsi0-6-15 x)
  (define y
    ((inline-chebyshev-flpoly-fun
      6.0 15.0
      (-1.925244301146350628816428684716550045429e4
       -1.198353263149330621624048804034028776287e4
       -3.477937539308286982692458873651726693812e3
       -4.803046790639916261784265178273773686249e2
       -2.562890803820053728218998699641891321552e1
       +3.740788808963323665370593041281819399850e-7
       -7.713638977946100306711746793694236435924e-8
       +1.565178892640717100414675722794707618947e-8
       -3.117542465640122863625790111806482043942e-9
       +6.073503129025465167965511518460566324779e-10
       -1.150949496047179890084998241564162255119e-10
       +2.102858491281106124889246275868833089257e-11
       -3.647260410862808863678434035974178985902e-12
       +5.827390822374468084948661977003242859188e-13
       -7.999114123557451949902041477681255993810e-14))
     x))
  (+ (/ y (flexpt x 5.0)) (fllog x)))

(: flpsi0-1.9-3 (Flonum -> Flonum))
(define (flpsi0-1.9-3 x)
  (define y
    ((inline-chebyshev-flpoly-fun
      1.9 3.0
      (-1.068595135558250365970154135318779762977
       0.007555110290613685894455357538792217876164
       -8.211196131506054788097165956947618267216e-4
       8.819898077656088066728485086945404403996e-5
       -9.377207860942770361265708932757532015772e-6
       9.882656181596999489012188167296378462293e-7
       -1.033878691914051272793824247210082411098e-7
       1.075053486743124253187897895133480499813e-8
       -1.112459077942683078328275732383680238168e-9
       1.146881360570377686259482136117612783879e-10
       -1.179172668817595317199544685547649548159e-11
       1.210211518974043571256683371091400286333e-12
       -1.240872724505738918165638170689453141739e-13
       1.272005760582855661944481497028186095588e-14
       -1.304284165442490577500377306573196471693e-15
       1.324746805869742319628788142120853886066e-16))
     x))
  (+ (/ y (flexpt x 1.0)) (fllog x)))

(: flpsi0-pos-zero (Flonum -> Flonum))
;; Error <= 2 ulp for 1.01 <= x <= 1.9
(define (flpsi0-pos-zero x)
  ((make-flpolyfun
    (1.0585432012233855e-08
     +9.6767223576078877096707761813142672147231142e-1
     -4.42763160500353127553797599467585510491966309e-1
     +2.58499753782129775648776735361540543429029315e-1
     -1.63942699544933748044376456834206467101233261e-1
     +1.07824045952487191135128855124325739325915008e-1
     -7.21995575193409626696893908469264780134510431e-2
     +4.8804285262128765080401907661996397609430272e-2
     -3.31611242500726233186197345853966159593666507e-2
     +2.25976465448926011125449117808694365399392534e-2
     -1.54247646368189203846188963269378209631264895e-2
     +1.0538790670881555307791836736009529596710663e-2
     -7.20453368572977191622981741327584775896463279e-3
     +4.92678087965504785805744172756649359819423221e-3
     -3.36980127720035923906475716305572919765314786e-3
     +2.30512605073145773629820036982065731913189791e-3
     -1.57693657080739955525227889461909068304999766e-3
     +1.07882505658776450345367614926996141127935916e-3
     -7.3807083404537248486922908431997361289422551e-4
     +5.04953190252646215739884819303074593735197609e-4
     -3.45467970810508971431806463930371137534420811e-4
     +2.36355976724137594774864968876273185881939724e-4
     -1.61706193084464479583646631218793971111171482e-4
     +1.10633707815531817492134312083342362402565131e-4
     -7.56917816597543344357441178188820822372883337e-5
     +5.1785747875351335968463825889133484772977912e-5
     -3.54300637882398506423667399737040919259208752e-5
     +2.42400610389497340053120082460952804952267374e-5
     -1.65842386724073585234799128562703946147149189e-5
     +1.13463820371245229094354235118485703293845643e-5
     -7.76281586742261875732585414254027665196056168e-6
     ))
   (- x 1.46163215590743)))

(: flpsi0-taylor-1 (Flonum -> Flonum))
;; Error <= 1 ulp for 0.75 <= x <= 1.23
(define (flpsi0-taylor-1 x)
  ((make-flpolyfun
    (-5.77215664901532860606512090082402431042159336e-1
     +1.6449340668482264364724151666460251892189499
     -1.20205690315959428539973816151144999076498629
     +1.08232323371113819151600369654116790277475095
     -1.03692775514336992633136548645703416805708092
     +1.01734306198444913971451792979092052790181749
     -1.00834927738192282683979754984979675959986356
     +1.00407735619794433937868523850865246525896079
     -1.00200839282608221441785276923241206048560585
     +1.00099457512781808533714595890031901700601953
     -1.00049418860411946455870228252646993646860644
     +1.00024608655330804829863799804773967096041609
     -1.00012271334757848914675183652635739571427511
     +1.0000612481350587048292585451051353337474817
     -1.00003058823630702049355172851064506258762795
     +1.00001528225940865187173257148763672202323739
     -1.00000763719763789976227360029356302921308825
     +1.0000038172932649998398564616446219397304547
     -1.00000190821271655393892565695779510135325857
     +1.00000095396203387279611315203868344934594379
     -1.00000047693298678780646311671960437304596645
     +1.00000023845050272773299000364818675299493504
     -1.00000011921992596531107306778871888232638726
     +1.0000000596081890512594796124402079358012275
     -1.00000002980350351465228018606370506936601184
     +1.00000001490155482836504123465850663069862886
     -1.00000000745071178983542949198100417060411946))
   (fl- x 1.0)))

(define pi^2/6 1.644934066848226436472415166646025189221)
(define -zeta3.0 -1.202056903159594285399738161511449990768)
(define pi^4/90 1.082323233711138191516003696541167902776)

(: flpsi0-tiny (Flonum -> Flonum))
(define (flpsi0-tiny x)
  (+ (/ -1.0 x)
     (- gamma.0)
     (* x (+ pi^2/6 (* x (+ -zeta3.0 (* x pi^4/90)))))))

(: flpsi0-negative (Flonum -> Flonum))
(define (flpsi0-negative x)
  (cond [(and (x . > . -inf.0) (not (integer? x)))
         (- (flpsi0 (- 1.0 x)) (* pi (flcotpix x)))]
        [else  +nan.0]))

(: flpsi0 (Flonum -> Flonum))
(define (flpsi0 x)
  (cond [(x . < . 3.0)
         (cond [(x . < . 0.75)
                (cond [(x . < . -0.0007)  (flpsi0-negative x)]
                      [(x . < . 0.0007)  (flpsi0-tiny x)]
                      [else  (+ (/ -1.0 x) (flpsi0-pos-zero (+ x 1.0)))])]
               [(x . < . 1.125)  (flpsi0-taylor-1 x)]
               [(x . < . 1.9)  (flpsi0-pos-zero x)]
               [else  (flpsi0-1.9-3 x)])]
        [(x . < . 6.0)
         (cond [(x . < . 4.0)
                (+ (/ -1.0 (+ x 2.0)) (/ -1.0 (+ x 1.0)) (/ -1.0 x) (flpsi0-6-15 (+ x 3.0)))]
               [(x . < . 5.0)
                (+ (/ -1.0 (+ x 1.0)) (/ -1.0 x) (flpsi0-6-15 (+ x 2.0)))]
               [else  (+ (/ -1.0 x) (flpsi0-6-15 (+ x 1.0)))])]
        [(x . < . 15.0)  (flpsi0-6-15 x)]
        [(x . < . 44.0)  (flpsi0-15-44 x)]
        [(x . < . 1e7)   (flpsi0-large x)]
        [(x . < . +inf.0)  (flpsi0-huge x)]
        [(x . = . +inf.0)  +inf.0]
        [else  +nan.0]))

(define pi.128 267257146016241686964920093290467695825/85070591730234615865843651857942052864)

(: flexppi (Flonum -> Flonum))
(define flexppi (make-flexpt pi.128))

(: flpsi (Integer Flonum -> Flonum))
(define (flpsi m x)
  (cond [(m . < . 0)  +nan.0]
        [(m . = . 0)  (flpsi0 x)]
        [(x . fl> . 0.0)
         (define sgn (if (even? m) -1.0 1.0))
         (define m+1 (fl+ (fl m) 1.0))
         (fl* (fl* sgn (flgamma m+1)) (flhurwitz-zeta m+1 x))]
        [(integer? x)
         (cond [(even? m)  (if (fl= 0.0 x) (/ -1.0 x) +nan.0)]
               [else  +inf.0])]
        [(and (x . fl> . -inf.0) (x . fl< . 0.0))
         (define sgn (if (even? m) 1.0 -1.0))
         (define t
           (fl* (flexppi (fl+ (fl m) 1.0))
                (flcot-diff/y m (flcotpix x))))
         (fl- (fl* sgn (flpsi m (fl- 1.0 x))) t)]
        [else  +nan.0]))

(: psi0 (Real -> Flonum))
(define (psi0 x)
  (cond [(flonum? x)  (flpsi0 x)]
        [(single-flonum? x)  (flpsi0 (fl x))]
        [(and (integer? x) (x . <= . 0))
         (raise-argument-error 'psi0 "Real, not Zero or Negative-Integer" x)]
        [else  (flpsi0 (fl x))]))

(: psi (Integer Real -> Flonum))
(define (psi m x)
  (cond [(m . < . 0)  (raise-argument-error 'psi "Natural" 0 m x)]
        [(flonum? x)  (flpsi m x)]
        [(single-flonum? x)  (flpsi m (fl x))]
        [(and (integer? x) (x . <= . 0))
         (raise-argument-error 'psi "Real, not Zero or Negative-Integer" 1 m x)]
        [else  (flpsi m (fl x))]))