(** Two electron integrals
*)


open Common
open Linear_algebra
open Gaussian
open Operators

open Constants
let cutoff = integrals_cutoff

module Bs   = Basis
module Cs   = Contracted_shell
module Csp  = Contracted_shell_pair
module Cspc = Contracted_shell_pair_couple
module Fis  = Four_idx_storage


module type Two_ei_structure =
  sig
    val name : string
    val class_of_contracted_shell_pair_couple :
      ?operator:Operator.t -> Cspc.t -> float Zmap.t
  end



module Make(T : Two_ei_structure) = struct

  include Four_idx_storage
  type t = Basis.t Four_idx_storage.t

  let class_of_contracted_shell_pair_couple = T.class_of_contracted_shell_pair_couple

  let store_class ?(cutoff=integrals_cutoff) data contracted_shell_pair_couple cls =
    let to_powers x =
      let open Zkey in
      match to_powers x with
      | Three x -> x
      | _ -> assert false
    in

    let shell_p = Cspc.shell_pair_p contracted_shell_pair_couple
    and shell_q = Cspc.shell_pair_q contracted_shell_pair_couple
    in

    let expand i_c powers_i shell = (
       Cs.index shell + i_c + 1,
       to_powers powers_i)
    in

    let l1 =
        Array.mapi (fun i x -> expand i x (Csp.shell_a shell_p))
           (Cs.zkey_array (Csp.shell_a shell_p))
    in
    let l2 =
        Array.mapi (fun i x -> expand i x (Csp.shell_b shell_p))
           (Cs.zkey_array (Csp.shell_b shell_p))
    in
    let l3 =
        Array.mapi (fun i x -> expand i x (Csp.shell_a shell_q))
           (Cs.zkey_array (Csp.shell_a shell_q))
    in
    let l4 =
        Array.mapi (fun i x -> expand i x (Csp.shell_b shell_q))
                (Cs.zkey_array (Csp.shell_b shell_q))
    in

    Array.iter (fun (l_c,xl) ->
       Array.iter (fun (k_c,xk) ->
          Array.iter (fun (j_c,xj) ->
             Array.iter (fun (i_c,xi) ->
                   let key = Zkey.of_powers_twelve  xi xj xk xl in
                    let value = Zmap.find cls key in
                    if abs_float value > cutoff then
                      set_chem data i_c j_c k_c l_c value
                  ) l1
              ) l2
          ) l3
    ) l4



  let of_basis ?operator basis =

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

    let eri_array =
      Fis.create ~size:n `Dense
    (*
      Fis.create ~size:n `Sparse
    *)
    in


    let t0 = Unix.gettimeofday () in

    let shell_pairs =
      Csp.of_contracted_shell_array shell
      |> Array.of_list
    in

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

    let t0 = Unix.gettimeofday () in

    let f shell_p =

      let sp =
        Csp.shell_pairs shell_p
      in

      try
          Array.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 cspc =
                if Array.length sp < Array.length sq then
                  Cspc.make ~cutoff shell_p shell_q
                else
                  Cspc.make ~cutoff shell_q shell_p
              in

              match cspc with
              | Some cspc ->
                let cls =
                  class_of_contracted_shell_pair_couple ?operator cspc
                in
                store_class ~cutoff eri_array cspc cls
              | None -> ()
          ) shell_pairs;
      with Exit -> ()
    in

    let pool =  Domainslib.Task.setup_pool ~num_domains:Qcaml.num_domains () in
    let _ =
      Domainslib.Task.run pool (fun _ ->
      let n = Array.length shell_pairs in
      let i_prev = ref 0 in
      shell_pairs
      |> Array.map (fun sp -> Domainslib.Task.async pool (fun _ -> f sp) )
      |> Array.iteri (fun i task ->
                      let i = ((10 * i+1)/n) in
                      if !i_prev <> i then (
                         i_prev := i; 
                         Printf.printf "%3d %%\n%!" (i*10) );
                      ignore (Domainslib.Task.await pool task)) 
      )
    in
    Domainslib.Task.teardown_pool pool;
    Printf.printf "Computed %s Integrals in %f seconds\n%!"
      T.name (Unix.gettimeofday () -. t0);
    eri_array


end