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

module Cs = Contracted_shell
module Gb = General_basis

(** 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 expo = Array.map (fun Gb.{exponent ; coefficient} ->
                exponent) shell
            and coef = Array.map (fun Gb.{exponent ; coefficient} ->
                coefficient) shell
            in
            Cs.make ~expo ~coef ~totAngMom ~center ~index:0)
      ) n
    |> Array.to_list
    |> Array.concat
  in
  Array.iteri (fun i x ->
      if (i > 0) then
        result.(i) <- Cs.with_index x (
            result.(i-1).index  +
            (Array.length result.(i-1).powers )) 
    ) result ;

  let size =
    let n = ref 0 in
    for i=0 to (Array.length result) - 1 do
      n := !n + (Array.length (result.(i).powers ))
    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 i -> 
        Cs.to_string i) b
    |> Array.to_list
    |> String.concat line
  )
  ^ line

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