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$ %} *)
  norm_coef : float array;       (** Normalization coefficients of primitive functions {% $1/\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. *)
  prim : PrimitiveShell.t array;
}

module Am = AngularMomentum
module Co = Coordinate
module Ps = PrimitiveShell


let make ?(index=0) lc = 
  assert (Array.length lc > 0);

  let coef = Array.map fst lc
  and prim = Array.map snd lc
  in

  let center = Ps.center prim.(0) in
  let rec unique_center = function
  | 0 -> true
  | i -> if Ps.center prim.(i) = center then unique_center (i-1) else false
  in
  if not (unique_center (Array.length prim - 1)) then
    invalid_arg "ContractedShell.make Coordinate.t differ";
    
  let totAngMom = Ps.totAngMom prim.(0) in
  let rec unique_angmom = function
  | 0 -> true
  | i -> if Ps.totAngMom prim.(i) = totAngMom then unique_angmom (i-1) else false
  in
  if not (unique_angmom (Array.length prim - 1)) then
    invalid_arg "ContractedShell.make: AngularMomentum.t differ";
    
  let expo = Array.map Ps.expo prim in

  let norm_coef =
    Array.map Ps.norm_coef prim
  in
  let norm_coef_scale = Ps.norm_coef_scale prim.(0)
  in
  { index ; expo ; coef ; center ; totAngMom ; norm_coef ;
    norm_coef_scale ; prim }


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" " ") ) 



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

let prim x = x.prim


(** {2 Printers} *)

let pp ppf x =
  let open Format in
  fprintf ppf "@[<2>{@ ";
  fprintf ppf "@[<2>expo =@ %a ;@]@ " pp_float_array_size x.expo;
  fprintf ppf "@[<2>coef =@ %a ;@]@ " pp_float_array_size x.coef;
  fprintf ppf "@[<2>center =@ %a ;@]@ " Co.pp_angstrom x.center;
  fprintf ppf "@[<2>totAngMom =@ %a ;@]@ " Am.pp_string x.totAngMom;
  fprintf ppf "@[<2>norm_coef =@ %a ;@]@ " pp_float_array_size x.norm_coef;
  fprintf ppf "@[<2>norm_coef_scale =@ %a ;@]@ " pp_float_array_size x.norm_coef_scale;
  fprintf ppf "@[<2>index =@ %d ;@]@ " x.index;
  fprintf ppf "}@,@]"