open Util

type t = {
  expo      : float array;
  coef      : float array;
  center    : Coordinate.t;
  totAngMom : Angular_momentum.t;
  size      : int;
  norm_coef : (int array -> float) array;
}

let size a = a.size
let expo a i = a.expo.(i)
let coef a i = a.coef.(i)
let center a = a.center
let totAngMom a = a.totAngMom
let norm_coef a i = a.norm_coef.(i)


let to_string s =
  let coord =
    Coordinate.to_Bohr s.center
  in
  let open Printf in
  ( sprintf "%2s %10.6f %10.6f %10.6f " (Angular_momentum.to_string s.totAngMom)
      (Coordinate.x coord) (Coordinate.y coord) (Coordinate.z coord) ) ^
  (Array.map2 (fun e c -> sprintf "%16.8e  %16.8e" e c) s.expo s.coef
   |> Array.to_list |> String.concat (sprintf "\n%36s" " ") ) 

(** Normalization coefficient of contracted function i, which depends on the
    exponent and the angular momentum. Two conventions can be chosen : a single
    normalisation factor for all functions of the class, or a coefficient which
    depends on the powers of x,y and z.
*)
let compute_norm_coef s =
  let atot =
    Angular_momentum.to_int s.totAngMom
  in
  Array.mapi (fun i alpha ->
      let c =
        ((alpha +. alpha) *. pi_inv)**(1.5) *. (pow (4. *. alpha) atot)
      in
      let result a = 
        let dfa = Array.map (fun j ->
            fact (j+j) /. ( float_of_int (1 lsl j) *. fact j)
          ) a
        in sqrt (c /.  (dfa.(0) *.dfa.(1) *. dfa.(2)))
      in
      result
    ) s.expo


let create ~expo ~coef ~center ~totAngMom =
  assert (Array.length expo = Array.length coef);
  assert (Array.length expo > 0);
  let tmp =
    { expo ; coef ; center ; totAngMom ; size=Array.length expo ; norm_coef = [||]}
  in
  { tmp with norm_coef = compute_norm_coef tmp }