(** Electron-electron repulsion integrals *)

open Util
open Constants
open Bigarray

type t = (float, float32_elt, fortran_layout) Bigarray.Genarray.t

module Am  = AngularMomentum
module Bs  = Basis
module Cs  = ContractedShell
module Csp = ContractedShellPair

let cutoff = integrals_cutoff

(** (00|00)^m : Fundamental electron repulsion integral
    $ \int \int \phi_p(r1) 1/r_{12} \phi_q(r2) dr_1 dr_2 $

    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
  let f = two_over_sq_pi *. (sqrt exp_pq) in
  let result = boys_function ~maxm t in
  let rec aux accu k = function
    | 0 -> result.(k) <- result.(k) *. accu
    | l -> 
      begin
        result.(k) <- result.(k) *. accu;
        let new_accu = -. accu *. exp_pq in
        aux new_accu (k+1) (l-1)
      end
  in
  aux f 0 maxm;
  result


(** Compute all the integrals of a contracted class *)
let contracted_class_shell_pairs ?schwartz_p ?schwartz_q shell_p shell_q : float Zmap.t =
  TwoElectronRR.contracted_class_shell_pairs ~zero_m ?schwartz_p ?schwartz_q shell_p shell_q

let contracted_class_shell_pairs_vec ?schwartz_p ?schwartz_q shell_p shell_q : float Zmap.t =
  TwoElectronRRVectorized.contracted_class_shell_pairs ~zero_m ?schwartz_p ?schwartz_q shell_p shell_q


let cutoff2 = cutoff *. cutoff
(*
type n_cls = { n : int ; cls : Zkey.t array }
*)
exception NullIntegral

(*
(** Unique index for integral <ij|kl> *)
let index i j k l =
  let f i k = 
    let (p,r) = 
      if i <= k then (i,k) else (k,i)
    in p+ (r*r-r)/2
  in
  let p = f i k and q = f j l in
  f p q
*)


let of_basis basis = 
  let to_powers x = 
    let open Zkey in
    match to_powers x with
    | Three x -> x
    | _ -> assert false
  in

  let n = Bs.size basis
  and shell = Bs.contracted_shells basis
  in


  (* Pre-compute all shell pairs *)
  let shell_pairs =
    Csp.of_contracted_shell_array shell
  in

  (* Pre-compute diagonal integrals for Schwartz *)
  let t0 = Unix.gettimeofday () in

  let schwartz = 
    Array.map (fun pair_array -> Array.map (function 
        | None -> (Zmap.create 0, 0.)
        | Some pair ->
          let cls = 
            contracted_class_shell_pairs pair pair
          in
          (cls, Zmap.fold (fun key value accu -> max (abs_float value) accu) cls 0. )
      ) pair_array ) shell_pairs
  in

  let icount = ref 0 in
  for i=0 to (Array.length shell) - 1 do
    print_int (Cs.index shell.(i)) ; print_newline ();
    for j=0 to i do
      let schwartz_p, schwartz_p_max = schwartz.(i).(j) in
        if (schwartz_p_max >= cutoff) then
          icount := !icount + 1;
    done;
  done;
  Printf.printf "%d shell pairs computed in %f seconds\n" !icount (Unix.gettimeofday () -. t0);

  (* Group shell pairs by common pairs of atoms *)
  

  (* 4D data initialization *)
  let eri_array = 
    Genarray.create Float32 fortran_layout [| n ; n ; n ; n|]
  in
  Genarray.fill eri_array 0.;

  (* Compute ERIs *)

  let t0 = Unix.gettimeofday () in
  let inn = ref 0 and out = ref 0 in

  for i=0 to (Array.length shell) - 1 do
    print_int (Cs.index shell.(i)) ; print_newline ();
    for j=0 to i do
      let schwartz_p, schwartz_p_max = schwartz.(i).(j) in
      try
        if (schwartz_p_max < cutoff) then raise NullIntegral;

        let shell_p = 
          match shell_pairs.(i).(j) with
          | None -> raise NullIntegral
          | Some x -> x
        in

        let sp =
          Csp.shell_pairs shell_p
        in

        for k=0 to i do
          for l=0 to k do
            let schwartz_q, schwartz_q_max = schwartz.(k).(l) in
            try
            if schwartz_p_max *. schwartz_q_max < cutoff2 then
              raise NullIntegral;

            let shell_q =
              match shell_pairs.(k).(l) with
              | None -> raise NullIntegral
              | Some x -> x
            in

            let sq =
              Csp.shell_pairs shell_q
            in

            let swap = 
                Array.length sp > Array.length sq 
            in

            (* Compute all the integrals of the class *)
            let cls =
              if swap then
                if (Array.length sp) + (Array.length sq)  < 4 then 
                  contracted_class_shell_pairs ~schwartz_p:schwartz_q ~schwartz_q:schwartz_p shell_q shell_p 
                else
                  contracted_class_shell_pairs_vec ~schwartz_p:schwartz_q ~schwartz_q:schwartz_p shell_q shell_p 
              else
                if (Array.length sp) + (Array.length sq)  < 4 then 
                  contracted_class_shell_pairs ~schwartz_p ~schwartz_q shell_p shell_q 
                else
                  contracted_class_shell_pairs_vec ~schwartz_p ~schwartz_q shell_p shell_q 
            in


            (* Write the data in the output file *)
            Array.iteri (fun i_c powers_i ->
                let i_c = Cs.index shell.(i) + i_c + 1 in
                let xi = to_powers powers_i in
                Array.iteri (fun j_c powers_j ->
                    let j_c = Cs.index shell.(j) + j_c + 1 in
                    let xj = to_powers powers_j in
                    Array.iteri (fun k_c powers_k ->
                        let k_c = Cs.index shell.(k) + k_c + 1 in 
                        let xk = to_powers powers_k in
                        Array.iteri (fun l_c powers_l ->
                            let l_c = Cs.index shell.(l) + l_c + 1 in
                            let xl = to_powers powers_l in
                            let key = 
                              if swap then
                                Zkey.of_powers_twelve  xk xl xi xj 
                              else
                                Zkey.of_powers_twelve  xi xj xk xl 
                            in
                            let value = 
                              Zmap.find cls key 
                            in
                              eri_array.{i_c,k_c,j_c,l_c} <- value;
                              eri_array.{j_c,k_c,i_c,l_c} <- value;
                              eri_array.{i_c,l_c,j_c,k_c} <- value;
                              eri_array.{j_c,l_c,i_c,k_c} <- value;
                              eri_array.{k_c,i_c,l_c,j_c} <- value;
                              eri_array.{k_c,j_c,l_c,i_c} <- value;
                              eri_array.{l_c,i_c,k_c,j_c} <- value;
                              eri_array.{l_c,j_c,k_c,i_c} <- value;
                            if (abs_float value > cutoff) then
                              (inn := !inn + 1;
                              )
                            else
                              out := !out + 1;
                          ) Am.(zkey_array (Singlet (Cs.ang_mom shell.(l))))
                      ) Am.(zkey_array (Singlet (Cs.ang_mom shell.(k))))
                  ) Am.(zkey_array (Singlet (Cs.ang_mom shell.(j))))
              ) Am.(zkey_array (Singlet (Cs.ang_mom shell.(i))))
            with NullIntegral ->  ()
          done;
        done;
      with NullIntegral -> ()
    done;
  done;
  Printf.printf "In: %d  Out:%d\n" !inn !out ;
  Printf.printf "Computed ERIs in %f seconds\n%!" (Unix.gettimeofday () -. t0);
  eri_array


(** Write all integrals to a file with the <ij|kl> convention *)
let to_file ~filename eri_array =
  let oc = open_out filename in
  (* Print ERIs *)
  for l_c=1 to (Genarray.nth_dim eri_array 3) do
    for k_c=1 to l_c do
      for j_c=1 to l_c do
        for i_c=1 to k_c do
          let value = eri_array.{i_c,j_c,k_c,l_c} in
          if (abs_float value > cutoff) then
            Printf.fprintf oc " %5d %5d %5d %5d%20.15f\n" i_c j_c k_c l_c value;
        done;
      done;
    done;
  done;
  close_out oc