(** Electron-electron repulsion integrals *)

open Util
open Constants

type t = FourIdxStorage.t


module Am  = AngularMomentum
module As  = AtomicShell
module Asp = AtomicShellPair
module Bs  = Basis
module Cs  = ContractedShell
module Csp = ContractedShellPair
module Cspc = ContractedShellPairCouple
module Fis = FourIdxStorage


let get_chem = Fis.get_chem
let get_phys = Fis.get_phys

let set_chem = Fis.set_chem
let set_phys = Fis.set_phys

let to_file = Fis.to_file

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





let class_of_contracted_shell_pair_couple shell_pair_couple = 
  let shell_p = Cspc.shell_pair_p shell_pair_couple
  and shell_q = Cspc.shell_pair_q shell_pair_couple
  in
  if Array.length (Csp.shell_pairs shell_p) + (Array.length (Csp.shell_pairs shell_q))  < 4 then 
    TwoElectronRR.contracted_class_shell_pair_couple ~zero_m shell_pair_couple
  else
    TwoElectronRRVectorized.contracted_class_shell_pairs ~zero_m shell_p shell_q





let filter_contracted_shell_pairs ?(cutoff=integrals_cutoff) shell_pairs = 
  List.map (fun pair -> 
        match Cspc.make ~cutoff pair pair with
        | Some cspc -> 
          let cls = class_of_contracted_shell_pair_couple cspc in
          (pair, Zmap.fold (fun key value accu -> max (abs_float value) accu) cls 0. )
          (* TODO \sum_k |coef_k * integral_k| *)
        | None -> (pair, -1.) 
    ) shell_pairs
  |> List.filter (fun (_, schwartz_p_max) -> schwartz_p_max >= cutoff)
  |> List.map fst


(*  TODO
let filter_contracted_shell_pair_couples ?(cutoff=integrals_cutoff) shell_pair_couples = 
  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

*)


let store_class ?(cutoff=integrals_cutoff) push_socket 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 msg = ref [] 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
                  msg := (i_c,j_c,k_c,l_c,value) :: !msg;
                ) (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));
    Zmq.Socket.send_all push_socket ["0" ; Bytes.to_string (Marshal.to_bytes !msg []) ]






let of_basis 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
    |> filter_contracted_shell_pairs ~cutoff
  in

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


  let t0 = Unix.gettimeofday () in
  let ishell = ref 0 in



  let zmq_addr = Printf.sprintf "ipc://%d" (Unix.getpid ()) in
  let () = 
    match Unix.fork () with
    | 0 ->    begin
                let zmq = ref None in

                Parmap.pariter ~chunksize:1 ~ncores:4
                  ~init:(fun _ ->
                    let zmq_context = 
                      Zmq.Context.create ()
                    in
                    let push_socket = 
                      Zmq.Socket.create zmq_context Zmq.Socket.push
                    in
                    Zmq.Socket.connect push_socket zmq_addr;
                    zmq := Some (zmq_context, push_socket)
                  )

                  (fun shell_p ->
                    let push_socket = 
                      match !zmq with
                      | Some (_, push_socket) -> push_socket
                      | None -> failwith "ZMQ"
                    in
                    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 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 cspc in
                                        store_class ~cutoff push_socket cspc cls
                          | None -> ()
                        ) shell_pairs
                    with Exit -> ()
                  ) (Parmap.L shell_pairs)
                  ~finalize:(fun _ ->
                    let zmq_context, push_socket = 
                      match !zmq with
                      | Some (zmq_context, push_socket) -> zmq_context, push_socket
                      | None -> failwith "ZMQ"
                    in
                    Zmq.Socket.close push_socket;
                    Zmq.Context.terminate zmq_context
                );
                let zmq_context = 
                  Zmq.Context.create ()
                in
                let push_socket = Zmq.Socket.create zmq_context Zmq.Socket.push in
                Zmq.Socket.connect push_socket zmq_addr;
                Zmq.Socket.send_all push_socket [ "1" ; ""];
                Zmq.Socket.close push_socket;
                Zmq.Context.terminate zmq_context;
                ignore @@ exit 0
              end

    | pid ->  begin
                let zmq_context = 
                  Zmq.Context.create ()
                in
                let pull_socket = 
                  Zmq.Socket.create zmq_context Zmq.Socket.pull
                in
                Zmq.Socket.bind pull_socket zmq_addr;

                try
                  while true do
                    match Zmq.Socket.recv_all pull_socket with
                    | "0" :: rest :: [] ->
                        List.iter (fun (i,j,k,l,value) ->
                          set_chem eri_array i j k l value) (Marshal.from_bytes (Bytes.of_string rest) 0) 
                    | "1" :: _ -> raise Exit
                    | _ -> invalid_arg "ERI"
                  done
                with Exit -> ();

                Zmq.Socket.close pull_socket;
                Zmq.Context.terminate zmq_context;
                ignore (Unix.wait ())
              end
  in
  Printf.printf "Computed ERIs in %f seconds\n%!" (Unix.gettimeofday () -. t0);
  eri_array