type t = 
  {
    size : int;
    contracted_shells : ContractedShell.t array;
  }

module Cs = ContractedShell
module Gb = GeneralBasis
module Ps = PrimitiveShell


(** Returns an array of the basis set per atom *)
let of_nuclei_and_general_basis n b =
  let result = 
    Array.map (fun (e, center) ->
        List.assoc e b
        |> Array.map (fun (totAngMom, shell) ->
            let lc = 
              Array.map (fun Gb.{exponent ; coefficient} ->
                coefficient, Ps.make totAngMom center exponent) shell
            in
            Cs.make lc)
      ) n
    |> Array.to_list
    |> Array.concat
  in
  Array.iteri (fun i x ->
      if (i > 0) then
        result.(i) <- Cs.with_index x (
            Cs.index result.(i-1)  +
            Cs.size_of_shell result.(i-1) ) 
    ) result ;

  let size =
    let n = ref 0 in
    for i=0 to (Array.length result) - 1 do
      n := !n + Cs.size_of_shell result.(i)
    done; !n
  in
  { contracted_shells = result ; size }



let to_string b =
  let b = b.contracted_shells in
  let line ="
-----------------------------------------------------------------------
" in 
  "
                          Atomic Basis set
                          ----------------

-----------------------------------------------------------------------
 #   Angular   Coordinates (Bohr)       Exponents       Coefficients
    Momentum  X        Y        Z
-----------------------------------------------------------------------
"
  ^
  ( Array.map (fun p -> Format.(fprintf str_formatter "%a" Cs.pp p;
    flush_str_formatter ())) b
    |> Array.to_list
    |> String.concat line
  )
  ^ line



let of_nuclei_and_basis_filename ~nuclei filename = 
  let general_basis = 
    GeneralBasis.read filename
  in
  of_nuclei_and_general_basis nuclei general_basis