open Util
open Constants

(** Computes all the overlap integrals of the contracted shell pair *)
let contracted_class shell_a shell_b : float Zmap.t = 

  let shell_p =
     Shell_pair.create_array shell_a shell_b
  in

  (* Pre-computation of integral class indices *)
  let class_indices = 
    Angular_momentum.zkey_array (Angular_momentum.Doublet
         Contracted_shell.(totAngMom shell_a, totAngMom shell_b))
  in

  let contracted_class = 
    Array.make (Array.length class_indices) 0.
  in

  (* Compute all integrals in the shell for each pair of significant shell pairs *)

  for ab=0 to (Array.length shell_p - 1)
  do
      let coef_prod =
         shell_p.(ab).Shell_pair.coef 
      in
      (** Screening on thr product of coefficients *)
      if (abs_float coef_prod) > 1.e-4*.cutoff then
        begin
            let center_ab = 
              shell_p.(ab).Shell_pair.center_ab 
            in
            let center_a = 
              shell_p.(ab).Shell_pair.center_a
            in
            let expo_inv = 
              shell_p.(ab).Shell_pair.expo_inv
            in
            let norm_coef_scale =
	      shell_p.(ab).Shell_pair.norm_coef_scale
	    in

            Array.iteri (fun i key ->
                      let (angMomA,angMomB) =
                        let a = Zkey.to_int_array Zkey.Kind_6 key in
                        ( [| a.(0) ;  a.(1) ;  a.(2) |], 
                          [| a.(3) ;  a.(4) ;  a.(5) |] )
                      in
                      let f k = 
                        Overlap_primitives.hvrr (angMomA.(k), angMomB.(k)) 
                        expo_inv
                        (Coordinate.coord center_ab k,
                         Coordinate.coord center_a  k) 
                      in
		      let norm =  norm_coef_scale.(i) in
                      let coef_prod = coef_prod *. norm in
                      let integral = chop norm (fun () -> (f 0)*.(f 1)*.(f 2)) in 
                      contracted_class.(i) <- contracted_class.(i) +. coef_prod *. integral
                ) class_indices
        end
  done;
  let result = 
    Zmap.create (Array.length contracted_class)
  in
  Array.iteri (fun i key -> Zmap.add result key contracted_class.(i)) class_indices;
  result





(** Write all overlap integrals to a file *)
let to_file ~filename basis =
  let to_int_tuple x = 
    let open Zkey in
    match to_int_tuple Kind_3 x with
    | Three x -> x
    | _ -> assert false
  in

  let oc = open_out filename in
  for i=0 to (Array.length basis) - 1 do
    print_int basis.(i).Contracted_shell.indice ; print_newline ();
    for j=0 to i do
      (* Compute all the integrals of the class *)
      let cls =
        contracted_class basis.(i) basis.(j) 
      in

      (* Write the data in the output file *)
      Array.iteri (fun i_c powers_i ->
          let i_c = basis.(i).Contracted_shell.indice + i_c + 1 in
          let xi = to_int_tuple powers_i in
          Array.iteri (fun j_c powers_j ->
              let j_c = basis.(j).Contracted_shell.indice + j_c + 1 in
              let xj = to_int_tuple powers_j in
              let key = 
                Zkey.of_int_tuple (Zkey.Six (xi,xj))
              in
              let value = 
                try
                  Zmap.find cls key 
                with Not_found -> failwith "Bug in overlap integrals"
              in
              if (abs_float value > cutoff) then
                Printf.fprintf oc "%4d %4d %20.12e\n"
                  i_c j_c value
            ) basis.(j).Contracted_shell.powers
        ) basis.(i).Contracted_shell.powers;
    done;
  done;
  close_out oc