open Util
open Constants
open Coordinate

type t = {
  expo      : float array;       (** Array of exponents {% $\alpha_i$ %} *)
  coef      : float array;       (** Array of contraction coefficients {% $d_i$ %} *)
  center    : Coordinate.t;      (** Coordinate of the center {% $\mathbf{A} = (X_A,Y_A,Z_A)$ %} *)
  totAngMom : AngularMomentum.t; (** Total angular momentum : {% $l = n_x + n_y + n_z$ %} *)
  size      : int;               (** Number of contracted functions, {% $m$ %} in the formula *)
  norm_coef : float array;       (** Normalization coefficients of primitive functions {% $\mathcal{N}_i$ %} *)
  norm_coef_scale : float array; (** Scaling factors {% $f_i$ %}, given in the same order as [AngularMomentum.zkey_array totAngMom]. *)
  index    : int;                (** Index in the basis set, represented as an array of contracted shells. *)
}

module Am = AngularMomentum


let compute_norm_coef expo totAngMom =
  let atot =
    Am.to_int totAngMom
  in
  let factor int_array = 
       let dfa = Array.map (fun j ->
           ( float_of_int (1 lsl j) *. fact j) /. fact (j+j) 
         ) int_array
       in
       sqrt (dfa.(0) *.dfa.(1) *. dfa.(2))
  in
  let expo = 
    if atot mod 2 = 0 then
      Array.map (fun alpha ->
        let alpha_2 = alpha +. alpha in
        (alpha_2 *. pi_inv)**(0.75) *. (pow (alpha_2 +. alpha_2) (atot/2))
      ) expo
    else
      Array.map (fun alpha ->
        let alpha_2 = alpha +. alpha in
        (alpha_2 *. pi_inv)**(0.75) *. sqrt (pow (alpha_2 +. alpha_2) atot)
      ) expo
  in
  Array.map (fun x -> let f a = x *. (factor a) in f) expo


let make ~index ~expo ~coef ~center ~totAngMom =
  assert (Array.length expo = Array.length coef);
  assert (Array.length expo > 0);
  let norm_coef_func =
    compute_norm_coef expo totAngMom
  in
  let powers =
    Am.zkey_array (Am.Singlet totAngMom) 
  in
  let norm_coef =
    Array.map (fun f -> f [| Am.to_int totAngMom ; 0 ; 0 |])  norm_coef_func
  in
  let norm_coef_scale = 
    Array.map (fun a ->
      (norm_coef_func.(0) (Zkey.to_int_array a)) /. norm_coef.(0)
    ) powers
  in
  { index ; expo ; coef ; center ; totAngMom ; size=Array.length expo ; norm_coef ;
    norm_coef_scale }


let with_index a i =
  { a with index = i }


let to_string s =
  let coord = s.center in
  let open Printf in
  (match s.totAngMom with
   | Am.S -> sprintf "%3d    " (s.index+1)
   | _                  -> sprintf "%3d-%-3d" (s.index+1) (s.index+(Array.length s.norm_coef_scale))
  ) ^
  ( sprintf "%1s %8.3f %8.3f %8.3f " (Am.to_string s.totAngMom)
      (get X coord) (get Y coord) (get 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.
    Returns, for each contracted function, an array of functions taking as
    argument the [|x;y;z|] powers.
*)
let compute_norm_coef expo totAngMom =
  let atot =
    Am.to_int totAngMom
  in
  let factor int_array = 
       let dfa = Array.map (fun j ->
           (float_of_int (1 lsl j) *. fact j) /. fact (j+j) 
         ) int_array
       in
       sqrt (dfa.(0) *.dfa.(1) *. dfa.(2))
  in
  let expo = 
    if atot mod 2 = 0 then
      Array.map (fun alpha ->
        let alpha_2 = alpha +. alpha in
        (alpha_2 *. pi_inv)**(0.75) *. (pow (alpha_2 +. alpha_2) (atot/2))
      ) expo
    else
      Array.map (fun alpha ->
        let alpha_2 = alpha +. alpha in
        (alpha_2 *. pi_inv)**(0.75) *. sqrt (pow (alpha_2 +. alpha_2) atot)
      ) expo
  in
  Array.map (fun x -> let f a = x *. factor a in f) expo


let expo x = x.expo

let coef x = x.coef

let center x = x.center

let totAngMom x = x.totAngMom

let size x = Array.length x.coef

let norm_coef x = x.norm_coef

let norm_coef_scale x = x.norm_coef_scale

let index x = x.index

let size_of_shell x = Array.length x.norm_coef_scale