(** Electron-electron repulsion integrals *)

open Util
open Constants
open Bigarray

let max_ao = 1 lsl 14

type index_pair = { first : int ; second : int }

type t = 
  | Dense  of (float, float64_elt, fortran_layout) Bigarray.Genarray.t
  | Sparse of (int, float) Hashtbl.t


let key_of_indices ~r1 ~r2 =
  let { first=i ; second=k } = r1 and { first=j ; second=l } = r2 in
  let i,k = if i<=k then i,k else k,i
  and j,l = if j<=l then j,l else l,j in
  let i,k,j,l = if k<=l then i,k,j,l else j,l,i,k in
  ((((((i lsl 15) lor k) lsl 15) lor j) lsl 15) lor l)


let get ~r1 ~r2 = function
  | Dense  t -> let { first=i ; second=k } = r1 and { first=j ; second=l } = r2 in t.{i,j,k,l}
  | Sparse t -> let key = key_of_indices ~r1 ~r2 in
                try Hashtbl.find t key
                with Not_found -> 0.


let set ~r1 ~r2 ~value = function
  | Dense  t -> let { first=i ; second=k } = r1 and { first=j ; second=l } = r2 in
                t.{i,j,k,l} <- value;
                t.{k,j,i,l} <- value;
                t.{i,l,k,j} <- value;
                t.{k,l,i,j} <- value;
                t.{j,i,l,k} <- value;
                t.{j,k,l,i} <- value;
                t.{l,i,j,k} <- value;
                t.{l,k,j,i} <- value;

  | Sparse t -> let key = key_of_indices ~r1 ~r2 in
                Hashtbl.replace t key value
    

let increment ~r1 ~r2 ~value = function
  | Dense  t -> let { first=i ; second=k } = r1 and { first=j ; second=l } = r2 in
                t.{i,j,k,l} <- t.{i,j,k,l} +. value;
                t.{k,j,i,l} <- t.{k,j,i,l} +. value;
                t.{i,l,k,j} <- t.{i,l,k,j} +. value;
                t.{k,l,i,j} <- t.{k,l,i,j} +. value;
                t.{j,i,l,k} <- t.{j,i,l,k} +. value;
                t.{j,k,l,i} <- t.{j,k,l,i} +. value;
                t.{l,i,j,k} <- t.{l,i,j,k} +. value;
                t.{l,k,j,i} <- t.{l,k,j,i} +. value;
  | Sparse t -> let key = key_of_indices ~r1 ~r2 in
                let old_value =
                  try Hashtbl.find t key
                  with Not_found -> 0.
                in
                Hashtbl.replace t key (old_value +. value)
    

let create = function
  | `Dense n -> 
    let eri_array = 
      Genarray.create Float64 fortran_layout [| n ; n ; n ; n|]
    in
    Genarray.fill eri_array 0.;
    Dense eri_array

  | `Sparse n -> 
    let eri_array = Hashtbl.create (n*n+13) in
    Hashtbl.add eri_array (-1) (float_of_int n);
    Sparse eri_array

let size = function
  | Dense t -> Genarray.nth_dim t 3
  | Sparse t -> Hashtbl.find t (-1) |> int_of_float

(** TODO : remove epsilons *)


let get_chem t i j k l = get ~r1:{ first=i ; second=j } ~r2:{ first=k ; second=l } t
let get_phys t i j k l = get ~r1:{ first=i ; second=k } ~r2:{ first=j ; second=l } t

let set_chem t i j k l value = set ~r1:{ first=i ; second=j } ~r2:{ first=k ; second=l } ~value t
let set_phys t i j k l value = set ~r1:{ first=i ; second=k } ~r2:{ first=j ; second=l } ~value t 


module Am  = AngularMomentum
module As  = AtomicShell
module Asp = AtomicShellPair
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 contracted_class_atomic_shell_pairs ?schwartz_p ?schwartz_q shell_p shell_q : float Zmap.t =
  TwoElectronRR.contracted_class_atomic_shell_pairs ~zero_m ?schwartz_p ?schwartz_q shell_p shell_q

(*
let contracted_class_atomic_shell_pairs_vec ?schwartz_p ?schwartz_q shell_p shell_q : float Zmap.t =
  TwoElectronRRVectorized.contracted_class_atomic_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
  (*TODO
    and atomic_shells = Bs.atomic_shells basis
  *)
  in


  (* Pre-compute all shell pairs *)
  let shell_pairs =
    Csp.of_contracted_shell_array shell
  in
  (*TODO
    let atomic_shell_pairs =
    Asp.of_atomic_shell_array ~cutoff atomic_shells
    in
  *)

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

  let schwartz = 
    List.map (fun pair -> 
        let cls = 
          contracted_class_shell_pairs pair pair
          (*TODO
            contracted_class_atomic_shell_pairs pair pair
          *)
        in
        (pair, cls, Zmap.fold (fun key value accu -> max (abs_float value) accu) cls 0. )
      ) shell_pairs
    |> List.filter (fun (_, _, schwartz_p_max) -> schwartz_p_max >= cutoff)
  in

  Printf.printf "%d shell pairs computed in %f seconds\n"
    (List.length schwartz) (Unix.gettimeofday () -. t0);


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


  (* 4D data initialization *)
  let eri_array = 
    (*
    create (`Dense n)
    *)
    create (`Sparse n)
  in

  (* Compute ERIs *)

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

  (*TODO
    for i=0 to (Array.length atomic_shells) - 1 do
  *)
  let ishell = ref 0 in
  List.iter (fun (shell_p, schwartz_p, schwartz_p_max) ->
      let () = 
        if (Cs.index (Csp.shell_a shell_p) > !ishell) then
          (ishell := Cs.index (Csp.shell_a shell_p) ; print_int !ishell ; print_newline ())
      in

      let sp =
        Csp.shell_pairs shell_p
        (*TODO
          Asp.atomic_shell_pairs shell_p
        *)
      in

      try
        List.iter (fun (shell_q, schwartz_q, schwartz_q_max) ->
            let () = 
              if  Cs.index (Csp.shell_a shell_q) >
                  Cs.index (Csp.shell_a shell_p) then
                raise Exit
            in

            try
              if schwartz_p_max *. schwartz_q_max < cutoff2 then
                raise NullIntegral;

              let sq =
                Csp.shell_pairs shell_q
                (*TODO
                  Asp.atomic_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
                  (*TODO
                    contracted_class_atomic_shell_pairs ~schwartz_p:schwartz_q ~schwartz_q:schwartz_p shell_q shell_p 
                  *)
                  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 (Csp.shell_a shell_p) + i_c + 1 in
                  let xi = to_powers powers_i in
                  Array.iteri (fun j_c powers_j ->
                      let j_c = Cs.index (Csp.shell_b shell_p) + j_c + 1 in
                      let xj = to_powers powers_j in
                      Array.iteri (fun k_c powers_k ->
                          let k_c = Cs.index (Csp.shell_a shell_q) + k_c + 1 in 
                          let xk = to_powers powers_k in
                          Array.iteri (fun l_c powers_l ->
                              let l_c = Cs.index (Csp.shell_b shell_q) + 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
                              set_chem eri_array i_c j_c k_c l_c value;
                              if (abs_float value > cutoff) then
                                (inn := !inn + 1;
                                )
                              else
                                out := !out + 1;
                            ) (Cs.zkey_array (Csp.shell_b shell_q))
                        ) (Cs.zkey_array (Csp.shell_a shell_q))
                    ) (Cs.zkey_array (Csp.shell_b shell_p))
                ) (Cs.zkey_array (Csp.shell_a shell_p))
            with NullIntegral ->  ()
          ) schwartz
      with Exit -> ()
    ) schwartz;
  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 size eri_array 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 = get_phys 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