(** Two electron integrals
*)

open Constants
let cutoff = integrals_cutoff

module Bs  = Basis
module Cs  = ContractedShell
module Csp = ContractedShellPair
module Cspc = ContractedShellPairCouple
module Fis = FourIdxStorage

module type TwoEI_structure =
  sig
    val name : string
    val class_of_contracted_shell_pair_couple :
      basis:Basis.t -> ContractedShellPairCouple.t -> float Zmap.t
  end


module Make(T : TwoEI_structure) = struct

  include FourIdxStorage

  let class_of_contracted_shell_pair_couple = T.class_of_contracted_shell_pair_couple 

  let filter_contracted_shell_pairs ?(cutoff=integrals_cutoff) ~basis shell_pairs =
    List.rev_map (fun pair ->
	match Cspc.make ~cutoff pair pair with
	| Some cspc ->
	  let cls = class_of_contracted_shell_pair_couple ~basis 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.rev_map fst


(*  TODO                                                                                      
    let filter_contracted_shell_pair_couples
    ?(cutoff=integrals_cutoff) shell_pair_couples = 
    List.rev_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.rev_map fst

*)


  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

    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_serial 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 ~basis ~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

    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 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 ~basis cspc
                in
                store_class ~cutoff eri_array cspc cls
              | None -> ()
            ) shell_pairs
        with Exit -> ()
      ) shell_pairs ;
    Printf.printf "Computed ERIs in %f seconds\n%!" (Unix.gettimeofday () -. t0);
    eri_array







  (* Parallel functions *)



  let of_basis_parallel basis = 

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

    let store_class_parallel
        ?(cutoff=integrals_cutoff) 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 result = 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
                      result := (i_c, j_c, k_c, l_c, value) :: !result
                    ) (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));
      !result
    in



    let t0 = Unix.gettimeofday () in

    let shell_pairs =
      Csp.of_contracted_shell_array shell
      |> filter_contracted_shell_pairs ~basis ~cutoff
    in

    if Parallel.master then
      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 max_int in

    let input_stream = Stream.of_list shell_pairs in

    let f shell_p =
      let () = 
        if Parallel.rank < 2 && 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

      let result = ref [] in
      let () =
        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 ~basis cspc
                in
                result := (store_class_parallel ~cutoff cspc cls) :: !result;
              | None -> ()
            ) shell_pairs;
        with Exit -> ()
      in
      List.concat !result |> Array.of_list
    in

    let eri_array = 
      if Parallel.master then
        Fis.create ~size:n `Dense
      else
        Fis.create ~size:n `Dense
    in

    Farm.run ~ordered:true ~f input_stream
    |> Stream.iter (fun l ->
        Array.iter (fun (i_c,j_c,k_c,l_c,value) ->
            set_chem eri_array i_c j_c k_c l_c value) l);

    if Parallel.master then
      Printf.printf
        "Computed %s Integrals in parallel in %f seconds\n%!" T.name (Unix.gettimeofday () -. t0);
    Fis.broadcast eri_array



  let of_basis = 
    match Parallel.size with
    | 1 -> of_basis_serial
    | _ -> of_basis_parallel


end