(** Two electron integral functor for operators that are separable among %{ $(x,y,z)$ %}.
    It is parameterized by the [zero_m] function.
*)

open Constants
let cutoff = integrals_cutoff

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


include FourIdxStorage

type f12_factor =
{
  expo_s  : float ;
  coef_g  : float array;
  expo_sg : float array;
  expo_sg_inv : float array;
}

let make_gaussian_corr_factor expo_s coef_g expo_sg = 
  let expo_sg_inv = 
    Array.map (fun x -> 1. /. (x *. expo_s *. expo_s)) expo_sg 
  in
  { 
    expo_s ; coef_g ; expo_sg ; expo_sg_inv 
  }


(* -1/expo_s *. exp (-expo_s r) *)
let gaussian_geminal expo_s =
  let coef_g =
    [| 0.3144 ; 0.3037 ; 0.1681 ; 0.09811 ; 0.06024 ; 0.03726 |]
    |> Array.map (fun x -> -. x /. expo_s)
  and expo_sg =
    [| 0.2209 ; 1.004 ;  3.622 ; 12.16 ;   45.87 ;  254.4 |]
  in
  make_gaussian_corr_factor expo_s coef_g expo_sg



(* exp (-expo_s r) *)
let simple_gaussian_geminal expo_s =
  let coef_g =
    [| 0.3144 ; 0.3037 ; 0.1681 ; 0.09811 ; 0.06024 ; 0.03726 |]
  and expo_sg =
    [| 0.2209 ; 1.004 ;  3.622 ; 12.16 ;   45.87 ;  254.4 |]
  in
  make_gaussian_corr_factor expo_s coef_g expo_sg



(** r12 * exp ( -expo_s * r) *)
let gaussian_geminal_times_r12 expo_s =
  let coef_g =
    [| 0.2454 ; 0.2938 ; 0.1815 ; 0.11281 ; 0.07502 ; 0.05280 |]
  and expo_sg =
    [| 0.1824 ; 0.7118;  2.252 ; 6.474 ;   19.66 ;  77.92 |]
  in make_gaussian_corr_factor expo_s coef_g expo_sg


(* exp (-expo_s r) *)
let simple_gaussian_geminal' expo_s =
  let coef_g =
    [|
    -3.4793465193721626604883567779324948787689208984375 ;
    -0.00571703486454788484955047422886309504974633455276489257812 ;
    4.14878218728681513738365538301877677440643310546875 ;
    0.202874298181392742623785352407139725983142852783203125 ;
    0.0819187742387294803858566183407674543559551239013671875 ;
    0.04225945671351955673644695821167260874062776565551757812 ;
    |]
  and expo_sg =
    [|
    0.63172472556807146570889699432882480323314666748046875; 
    26.3759196683467962429858744144439697265625; 
    0.63172102793029016876147352377302013337612152099609375; 
    7.08429025944207335641067402320913970470428466796875; 
    42.4442841447001910637482069432735443115234375;
    391.44036073596890901171718724071979522705078125 ;
    |]
  in make_gaussian_corr_factor expo_s coef_g expo_sg




let one_over_r =
  
  let coef_g = [|
    841.88478132     ;  70.590185207     ;  18.3616020768    ;  7.2608642093    ;
    3.57483416444    ;  2.01376031082    ;  1.24216542801    ;  0.81754348620   ;
    0.564546514023   ;  0.404228610699   ;  0.297458536575   ;  0.223321219537  ;
    0.169933732064   ;  0.130190978230   ;  0.099652303426   ;  0.075428246546  ;
    0.0555635614051  ;  0.0386791283055  ;  0.0237550435652  ;  0.010006278387  ;
    |]
  and expo_sg_inv =
    [| 84135.654509    ;  2971.58727634    ;  474.716025959    ;  130.676724560    ;
        47.3938388887   ;  20.2078651631    ;  9.5411021938     ;  4.8109546955     ;
        2.52795733067   ;  1.35894103210    ;  0.73586710268    ;  0.39557629706    ;
        0.20785895177   ;  0.104809693858   ;  0.049485682527   ;  0.021099788990   ;
        0.007652472186  ;  0.0021065225215  ;  0.0003365204879  ;  0.0000118855674  |]
  in make_gaussian_corr_factor 1.0 coef_g expo_sg_inv






module Zero_m = struct
  let name = "F12"
end 

let class_of_contracted_shell_pair_couple f12 shell_pair_couple = 
    F12RR.contracted_class_shell_pair_couple 
      f12.expo_sg_inv f12.coef_g shell_pair_couple


let filter_contracted_shell_pairs f12 ?(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 f12 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 basis ?(cutoff=integrals_cutoff) data f12 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 lambda_inv = -. 1. /. f12.expo_s 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
                  (*
                  lambda_inv *. (s.{i_c,j_c} *. s.{k_c,l_c} -. value)
                  lambda_inv *. value
                  *)
                  value
                  |> set_chem data i_c j_c k_c l_c
                ) (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 f12 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 f12 ~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 f12 cspc
              in
              store_class basis ~cutoff eri_array f12 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 f12 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 f12 ~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 (List.rev 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
    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 f12 cspc
            in
            result := (store_class_parallel ~cutoff cspc cls) :: !result;
          | None -> ()
        ) shell_pairs;
      raise Exit
    with Exit -> List.concat !result |> Array.of_list
  in

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

  (*
  let lambda_inv = -. 1. /. f12.expo_s in
  *)
  Farm.run ~ordered:false ~f input_stream
  |> Stream.iter (fun l ->
      Array.iter (fun (i_c,j_c,k_c,l_c,value) ->
      (*
        lambda_inv *. (s.{i_c,j_c} *. s.{k_c,l_c} -. value)
        lambda_inv *. value
        *)
        value
        |> set_chem eri_array i_c j_c k_c l_c) l);

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



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