(** Electron-nucleus repulsion integrals *)

open Util
open Constants
open Bigarray

(** (0|0)^m : Fundamental electron-nucleus repulsion integral
    $ \int \phi_p(r1) 1/r_{C} dr_1 $

    maxm        : Maximum total angular momentum
    expo_pq_inv : $1./p + 1./q$ where $p$ and $q$ are the exponents of
                $\phi_p$ and $\phi_q$
    norm_pq_sq  : square of the distance between the centers of $\phi_p$
                and $\phi_q$
*)

let zero_m ~maxm ~expo_pq_inv ~norm_pq_sq =
  let exp_pq = 1. /. expo_pq_inv in
  let t = norm_pq_sq *. exp_pq in
  boys_function ~maxm t
  |> Array.mapi (fun m fm ->
      two_over_sq_pi *. fm *.
      (pow exp_pq m) *. (sqrt exp_pq)
    )


(** Compute all the integrals of a contracted class *)
let contracted_class_shell_pair shell_p geometry: float Zmap.t =
  OneElectronRR.contracted_class_shell_pair ~zero_m shell_p geometry


let cutoff2 = cutoff *. cutoff
exception NullIntegral

let to_file ~filename basis geometry =
  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

  (* Pre-compute all shell pairs *)
  let shell_pairs =
    Array.mapi (fun i shell_a -> Array.map (fun shell_b -> 
        Shell_pair.create_array shell_a shell_b) (Array.sub basis 0 (i+1)) ) basis
  in
  Printf.printf "%d shells\n" (Array.length basis);

  let eni_array = 
    let n = ref 0 in
    for i=0 to (Array.length basis) - 1 do
      n := !n + (Array.length (basis.(i).Contracted_shell.powers))
    done;
    let n = !n in
    Array2.create Float64 c_layout n  n
  in
  Array2.fill eni_array 0.;

  (* Compute Integrals *)
  let t0 = Unix.gettimeofday () in

  let inn = ref 0 and out = ref 0 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
      let
        shell_p = shell_pairs.(i).(j)
      in

      (* Compute all the integrals of the class *)
      let cls =
        contracted_class_shell_pair shell_p geometry
      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 = 
                Zmap.find cls key 
              in
              if (abs_float value > cutoff) then
                (inn := !inn + 1;
                 eni_array.{(i_c-1),(j_c-1)} <- value;
                )
              else
                out := !out + 1;
            ) basis.(j).Contracted_shell.powers
        ) basis.(i).Contracted_shell.powers;
    done;
  done;

  Printf.printf "Computed %d non-zero ENIs in %f seconds\n" !inn (Unix.gettimeofday () -. t0);

  (* Print ENIs *)
  for i_c=1 to (Array2.dim1 eni_array) do
    for j_c=1 to (Array2.dim2 eni_array) do
      let value = eni_array.{(i_c-1),(j_c-1)} in
      if (value <> 0.) then
        Printf.fprintf oc " %5d %5d %20.15f\n" i_c j_c value;
    done;
  done;
  Printf.printf "In: %d  Out:%d\n" !inn !out ;
  close_out oc