(** Functions from libm *)
external erf_float : float -> float = "erf_float_bytecode" "erf_float" [@@unboxed]  [@@noalloc]
external erfc_float : float -> float = "erfc_float_bytecode" "erfc_float" [@@unboxed]  [@@noalloc]
external gamma_float : float -> float = "gamma_float_bytecode" "gamma_float" [@@unboxed]  [@@noalloc]

open Constants

let factmax = 150





(* Incomplete gamma function : Int_0^x   exp(-t) t^(a-1) dt  
  p:  1 / Gamma(a) * Int_0^x   exp(-t) t^(a-1) dt  
  q:  1 / Gamma(a) * Int_x^inf exp(-t) t^(a-1) dt  

  reference - Haruhiko Okumura: C-gengo niyoru saishin algorithm jiten
  (New Algorithm handbook in C language) (Gijyutsu hyouron
  sha, Tokyo, 1991) p.227 [in Japanese] *)

let incomplete_gamma ~alpha x = 
  let a = alpha in
  let a_inv = 1./. a in
  let gf = gamma_float alpha in
  let loggamma_a = log gf in
  let rec p_gamma x =
    if x >= 1. +. a then 1. -. q_gamma x 
    else if x = 0. then 0.
    else
      let rec pg_loop prev res term k =
        if k > 1000. then failwith "p_gamma did not converge."
        else if prev = res then res
        else
          let term = term *. x /. (a +. k) in
          pg_loop res (res +. term) term (k +. 1.)
      in
      let r0 =  exp (a *. log x -. x -. loggamma_a) *. a_inv in
      pg_loop min_float r0 r0 1.

  and q_gamma x =
    if x < 1. +. a then 1. -. p_gamma x 
    else
      let rec qg_loop prev res la lb w k =
        if k > 1000. then failwith "q_gamma did not converge."
        else if prev = res then res
        else
          let k_inv = 1. /. k in
          let kma = (k -. 1. -. a) *. k_inv in
          let la, lb =
            lb, kma *. (lb -. la) +. (k +. x) *. lb *. k_inv
          in
          let w = w *. kma  in
          let prev, res = res, res +. w /. (la *. lb) in
          qg_loop prev res la lb w (k +. 1.)
      in
      let w = exp (a *. log x -. x -. loggamma_a) in
      let lb = (1. +. x -. a) in
      qg_loop min_float (w /. lb) 1. lb w 2.0
  in
  gf *. p_gamma x 





let fact_memo =
  let rec aux accu_l accu = function
    | 0 ->  aux [1.] 1. 1
    | i when (i = factmax) ->
      let x = (float_of_int factmax) *. accu in
      List.rev (x::accu_l)
    | i ->  let x = (float_of_int i) *. accu in
      aux (x::accu_l) x (i+1)
  in
  aux [] 0. 0
  |> Array.of_list



(** Factorial function. 
    @raise Invalid_argument for negative or arguments >100. *)
let fact = function
  | i when (i < 0) ->
    raise (Invalid_argument "Argument of factorial should be non-negative")
  | i when (i > 150) ->
    raise (Invalid_argument "Result of factorial is infinite")
  | i -> fact_memo.(i)



(** Integer powers of floats *)
let rec pow a = function
  | 0 -> 1.
  | 1 -> a
  | 2 -> a *. a
  | 3 -> a *. a *. a
  | -1 -> 1. /. a
  | n when (n<0) -> pow (1./.a) (-n)
  | n ->
    let b = pow a (n / 2) in
    b *. b *. (if n mod 2 = 0 then 1. else a)
;;


(** In chop f g, evaluate g only if f is non zero, and return f *. (g ()) *)
let chop f g =
  if (abs_float f) < cutoff then 0.
  else f *. (g ())



(** Generalized Boys function. 
    maxm : Maximum total angular momentum
*)
let boys_function ~maxm t =
  match maxm with
  | 0 ->
    begin
      if t = 0. then [| 1. |] else
        let sq_t = sqrt t in
        [| (sq_pi_over_two /. sq_t) *.  erf_float sq_t |]
    end
  | _ ->
    begin
      let result =
        Array.init (maxm+1) (fun m -> 1. /. float_of_int (2*m+1))
      in
      if t <> 0. then
        begin
          let fmax = 
            let t_inv = sqrt (1. /. t) in
            let n = float_of_int maxm in
            let dm = 0.5 +. n in
            let f  = (pow t_inv (maxm+maxm+1) ) in
            (incomplete_gamma dm t) *. 0.5 *. f
          in
          let emt = exp (-. t) in
          result.(maxm) <- fmax;
          for n=maxm-1 downto 0 do
            result.(n) <- ( (t+.t) *. result.(n+1) +. emt) *. result.(n)
          done
        end;
      result
    end