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

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


(** Returns an array of the basis set per atom *)
let of_nuclei_and_general_basis nucl bas =
  let index_ = ref 0 in
  let contracted_shells = 
    Array.map (fun (e, center) ->
        List.assoc e bas
        |> Array.map (fun (totAngMom, shell) ->
            let lc = 
              Array.map (fun Gb.{exponent ; coefficient} ->
                coefficient, Ps.make totAngMom center exponent) shell
            in
            let result = Cs.make ~index:!index_ lc in
            index_ := !index_ + Cs.size_of_shell result;
            result
          )
    ) nucl
    |> Array.to_list
    |> Array.concat
  in
  let contracted_atomic_shells = lazy(
    let uniq_center_angmom =
      Array.map (fun x -> Cs.center x, Cs.totAngMom x) contracted_shells
      |> Array.to_list
      |> List.sort_uniq compare
    in
    let csl = 
      Array.to_list contracted_shells
    in
    List.map (fun (center, totAngMom) ->
      let a = 
        List.filter (fun x -> Cs.center x = center && Cs.totAngMom x = totAngMom) csl
        |> Array.of_list
      in
      Ca.make ~index:(Cs.index a.(0)) a
    ) uniq_center_angmom
    |> List.sort (fun x y -> compare (Ca.index x) (Ca.index y))
    |> Array.of_list
  ) in
  { contracted_shells ; contracted_atomic_shells ; size = !index_ }


let size x = x.size

let contracted_atomic_shells x = Lazy.force x.contracted_atomic_shells

let contracted_shells x = x.contracted_shells



let to_string b =
  let b = contracted_atomic_shells b 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" Ca.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