(** Electron-electron repulsion integrals *)

open Util
open Constants
open Bigarray

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

(* Input type *)
type input_data = 
  {
    i               : int;
    j               : int;
    shell_pairs     : ContractedShellPair.t array array ;
    schwartz        : (float Zmap.t * float) array array; 
    cutoff          : float
  }

(* Output type *)

type output_integral =  (* <ij|kl> *)
  {
    i1 : int ;       (* Function i for electron 1 *)
    j2 : int ;       (* Function j for electron 2 *)
    k1 : int ;       (* Function k for electron 1 *)
    l2 : int ;       (* Function l for electron 2 *)
    swap : bool ;    (* If true, Compute (kl|ij) instead of (ij|kl) *)
    cls : float Zmap.t;
  }

type output_data = output_integral list


(** (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 when shell pairs are not yet available *)
let contracted_class  shell_a shell_b shell_c shell_d : float Zmap.t =
  TwoElectronRR.contracted_class ~zero_m shell_a shell_b shell_c shell_d 


(** Compute all the integrals of a contracted class *)
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


(** 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


(** Creates the input data structure from the basis set. *)
let make_input basis = 

  let shell = Basis.contracted_shells basis in

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

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

  let schwartz = 
    Array.map (fun pair_array ->
      Array.map (fun pair ->
        let cls = contracted_class_shell_pairs pair pair in
        let m = Zmap.fold (fun _ value accu ->
          max (abs_float value) accu) cls 0. 
        in (cls, m)
      ) pair_array
    ) shell_pairs
  in

  (* Count number of significant shell pairs. *)
  let icount = ref 0 in
  for i=0 to (Array.length shell) - 1 do
    print_int (Contracted_shell.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);

  List.init (Array.length shell) (fun i ->
    List.init (i+1) (fun j ->
        { i ; j ; shell_pairs ; schwartz ; cutoff } ) )
  |> List.fold_left List.rev_append []



exception NullIntegral

let slave_job { i ; j ; shell_pairs ; schwartz ; cutoff} =

  let shell_p = shell_pairs.(i).(j) in
  let schwartz_p, schwartz_p_max = schwartz.(i).(j) in
  if schwartz_p_max < cutoff then [] else
  let schwartz_cutoff = cutoff *. cutoff in
  let sp = shell_p.ContractedShellPair.shell_pairs in

  let f k l  =
    let schwartz_q, schwartz_q_max = schwartz.(k).(l) in
    if schwartz_p_max *. schwartz_q_max < schwartz_cutoff then
      raise NullIntegral;

    let shell_q = shell_pairs.(k).(l) in
    let sq = shell_q.ContractedShellPair.shell_pairs 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 = 2 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 = 2 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
      { i1=i ; j2=k ; k1=j ; l2=l ; swap ; cls }
  in
  let rec loop accu k l = 
    match k, l with
    | -1, -1 -> accu
    | k,  -1 -> loop accu (k-1) (k-1)
    | k,   l ->
      let new_accu =
        let _, schwartz_q_max = schwartz.(k).(l) in
        if schwartz_p_max *. schwartz_q_max > schwartz_cutoff then
            f k l :: accu
        else accu
      in
      loop new_accu k (l-1)
  in
  loop [] i i


let of_basis basis = 

  let shell = Basis.contracted_shells basis in

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

  (* Compute ERIs *)

  let to_int_tuple x = 
    let open Zkey in
    match to_int_tuple Kind_3 x with
    | Three x -> x
    | _ -> assert false
  in

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

  make_input basis 
  |> List.map slave_job
  |> List.iter (fun  output_ij ->
      List.iter (fun  { i1=i ; j2=k ; k1=j ; l2=l ; swap ; cls } ->

        Array.iteri (fun i_c powers_i ->
            let i_c = (Contracted_shell.index shell.(i)) + i_c + 1 in
            let xi = to_int_tuple powers_i in
            Array.iteri (fun j_c powers_j ->
                let j_c = (Contracted_shell.index shell.(j)) + j_c + 1 in
                let xj = to_int_tuple powers_j in
                Array.iteri (fun k_c powers_k ->
                    let k_c = (Contracted_shell.index shell.(k)) + k_c + 1 in
                    let xk = to_int_tuple powers_k in
                    Array.iteri (fun l_c powers_l ->
                        let l_c = (Contracted_shell.index shell.(l)) + l_c + 1 in
                        let xl = to_int_tuple powers_l in
                        let key =
                          if swap then
                            Zkey.of_int_tuple (Zkey.Twelve (xk,xl,xi,xj))
                          else
                            Zkey.of_int_tuple (Zkey.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;
                      ) (Contracted_shell.powers shell.(l))
                  ) (Contracted_shell.powers shell.(k))
              ) (Contracted_shell.powers shell.(j))
          ) (Contracted_shell.powers shell.(i));
    ) output_ij 
  ); 
  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