(** 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 = FourIdxStorage.t

let get_chem = FourIdxStorage.get_chem
let get_phys = FourIdxStorage.get_phys

let set_chem = FourIdxStorage.set_chem
let set_phys = FourIdxStorage.set_phys


module Am  = AngularMomentum
module As  = AtomicShell
module Asp = AtomicShellPair
module Bs  = Basis
module Cs  = ContractedShell
module Csp = ContractedShellPair


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





let class_of_contracted_shell_pairs shell_p shell_q = 
  if Array.length (Csp.shell_pairs shell_p) + (Array.length (Csp.shell_pairs shell_q))  < 4 then 
    TwoElectronRR.contracted_class_shell_pairs ~zero_m shell_p shell_q
  else
    TwoElectronRRVectorized.contracted_class_shell_pairs ~zero_m shell_p shell_q





let filter_contracted_shell_pairs ?(cutoff=integrals_cutoff) shell_pairs = 
  let t0 = Unix.gettimeofday () in

  let schwartz = 
    List.map (fun pair -> 
        let cls = 
          class_of_contracted_shell_pairs pair pair
        in
        (pair, 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)
    |> List.map fst
  in

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

let store_class ?(cutoff=integrals_cutoff) data cls shell_p shell_q = 
    let to_powers x = 
      let open Zkey in
      match to_powers x with
      | Three x -> x
      | _ -> assert false
    in
    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 = Zkey.of_powers_twelve  xi xj xk xl in
                    let value = Zmap.find cls key in
                    set_chem data i_c j_c k_c l_c value;
                  ) (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))






let of_basis basis = 

  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

  (* Pre-compute diagonal integrals for Schwartz *)
  let schwartz = 
    filter_contracted_shell_pairs shell_pairs
  in

  (* 4D data initialization *)
  let eri_array = 
    FourIdxStorage.create ~size:n `Dense
    (*
    FourIdxStorage.create ~size:n `Sparse
    *)
  in

  (* Compute ERIs *)

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

  let ishell = ref 0 in
  List.iter (fun shell_p ->
      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
      in

      try
        List.iter (fun shell_q ->
            let () = 
              if  Cs.index (Csp.shell_a shell_q) >
                  Cs.index (Csp.shell_a shell_p) then
                raise Exit
            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 f p q = 
              let cls = class_of_contracted_shell_pairs p q in
              store_class ~cutoff:integrals_cutoff eri_array cls p q
            in
            if swap then
              f shell_q shell_p
            else
              f shell_p shell_q
          ) 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 = FourIdxStorage.to_file