(** %{ $ \langle ij | H F | kl \rangle $ %} integrals. *)

open Lacaml.D

module Fis = FourIdxStorage

type t = (float, Bigarray.float64_elt, Bigarray.fortran_layout) Bigarray.Genarray.t 
       * (float, Bigarray.float64_elt, Bigarray.fortran_layout) Bigarray.Genarray.t 
       * (float, Bigarray.float64_elt, Bigarray.fortran_layout) Bigarray.Genarray.t 
       * (float, Bigarray.float64_elt, Bigarray.fortran_layout) Bigarray.Genarray.t 


let make ~simulation ~mo_basis ~aux_basis_filename () =
  
  let f12 = Util.of_some @@ Simulation.f12 simulation in                                      
  let mo_num = MOBasis.size mo_basis in

  (* Add auxiliary basis set *)                                                               
  let aux_basis = 
    let s =
      let charge        = Charge.to_int @@ Simulation.charge simulation 
      and multiplicity  = Electrons.multiplicity @@ Simulation.electrons simulation
      and nuclei        = Simulation.nuclei simulation
      in
      let general_basis = 
	Basis.general_basis @@ Simulation.basis simulation
      in
      GeneralBasis.combine [
	general_basis ; GeneralBasis.read aux_basis_filename
      ]
      |> Basis.of_nuclei_and_general_basis nuclei
      |> Simulation.make ~f12 ~charge ~multiplicity ~nuclei 
    in
    MOBasis.of_mo_basis s mo_basis
  in

  let aux_num = MOBasis.size aux_basis in

  (* Fire calculation of F12 and ERI *)
  let f12 = 
    MOBasis.f12_ints aux_basis
  in
  let eri =
    MOBasis.two_e_ints aux_basis
  in

  (* Compute the <ij|QHF|kl> integrals *)
  if Parallel.master then Printf.eprintf "Computing HF12 integrals\n%!";

  let result_s, result_o, resultm_s, resultm_o =
    Bigarray.Genarray.create Float64 Bigarray.fortran_layout [| mo_num ; mo_num ; mo_num ; mo_num |] , 
    Bigarray.Genarray.create Float64 Bigarray.fortran_layout [| mo_num ; mo_num ; mo_num ; mo_num |] , 
    Bigarray.Genarray.create Float64 Bigarray.fortran_layout [| mo_num ; mo_num ; mo_num ; mo_num ; mo_num |] , 
    Bigarray.Genarray.create Float64 Bigarray.fortran_layout [| mo_num ; mo_num ; mo_num ; mo_num ; mo_num |] 
  in


  let h_s  = Bigarray.Array3.create Float64 Bigarray.fortran_layout mo_num aux_num aux_num in
  let f_s  = Bigarray.Array3.create Float64 Bigarray.fortran_layout aux_num aux_num mo_num in
  let h_o  = Bigarray.Array3.create Float64 Bigarray.fortran_layout mo_num aux_num aux_num in
  let f_o  = Bigarray.Array3.create Float64 Bigarray.fortran_layout aux_num aux_num mo_num in
  let hf_s = Mat.create mo_num mo_num in
  let hf_o = Mat.create mo_num mo_num in

  let hfm_s  = Bigarray.Array3.create Float64 Bigarray.fortran_layout mo_num mo_num mo_num in
  let hfm_o  = Bigarray.Array3.create Float64 Bigarray.fortran_layout mo_num mo_num mo_num in

  for a=1 to mo_num do
    for b=1 to mo_num do
     for i=1 to mo_num do
        h_s.{i, a, b} <- 0. ;
        h_o.{i, a, b} <- 0. 
     done
    done
  done;

  for k=1 to mo_num do
    for b=1 to mo_num do
      for a=1 to mo_num do
        f_s.{a, b, k} <- 0. ;
        f_o.{a, b, k} <- 0.
     done
    done
  done;

  let task (j,l) = 
    
      let h i a b = 
         h_s.{i, a, b} <- ERI.get_phys eri i j a b -. ERI.get_phys eri i j b a ;
         h_o.{i, a, b} <- ERI.get_phys eri i j a b
      and f a b k =
         f_s.{a, b, k} <- 0.25 *. (F12.get_phys f12 a b k l -. F12.get_phys f12 a b l k) ;
         f_o.{a, b, k} <- 0.375 *. F12.get_phys f12 a b k l +. 0.125 *. F12.get_phys f12 b a k l
      in

      for a=mo_num+1 to aux_num do
        for b=mo_num+1 to aux_num do
         for i=1 to mo_num do
            h i a b
         done
        done
      done;

      for k=1 to mo_num do
        for b=mo_num+1 to aux_num do
          for a=mo_num+1 to aux_num do
            f a b k
         done
        done
      done;

(*
      let h i a b = 
         h_s.{i, a, b} <- 0. ;
         h_o.{i, a, b} <- 0.
      and f a b k =
         f_s.{a, b, k} <- 0. ;
         f_o.{a, b, k} <- 0.
      in
*)

      for m=1 to mo_num do
        for a=mo_num+1 to aux_num do
          for i=1 to mo_num do
              h i a m ;
              h i m a
          done
        done
      done;

      for k=1 to mo_num do
        for m=1 to mo_num do
          for a=mo_num+1 to aux_num do
            f a m k ;
            f m a k
          done
        done
      done;

      let h_o' =
        Bigarray.(reshape (genarray_of_array3 h_o)) [| mo_num ; aux_num*aux_num |]
        |> Bigarray.array2_of_genarray
      in
      let f_o' =
        Bigarray.(reshape (genarray_of_array3 f_o)) [| aux_num*aux_num ; mo_num |] 
        |> Bigarray.array2_of_genarray
      in
      let h_s' =
        Bigarray.(reshape (genarray_of_array3 h_s)) [| mo_num ; aux_num*aux_num |]
        |> Bigarray.array2_of_genarray
      in
      let f_s' =
        Bigarray.(reshape (genarray_of_array3 f_s)) [| aux_num*aux_num ; mo_num |] 
        |> Bigarray.array2_of_genarray
      in
      let hf_s = gemm ~c:hf_s h_s' f_s' in
      let hf_o = gemm ~c:hf_o h_o' f_o' in

      let () = 
        for m=1 to mo_num do
          let h_o' =
            Mat.init_cols mo_num aux_num (fun i a -> h_o.{i,m,a})
          in
          let f_o' =
            Mat.init_cols aux_num mo_num (fun a k -> f_o.{m,a,k})
          in
          let h_s' =
            Mat.init_cols mo_num aux_num (fun i a -> h_s.{i,m,a})
          in
          let f_s' =
            Mat.init_cols aux_num mo_num (fun a k -> f_s.{m,a,k})
          in
          let r_s, r_o =
            gemm h_s' f_s' , 
            gemm h_o' f_o'
          in
          for k = 1 to mo_num do
            for i = 1 to mo_num do
              hfm_s.{m,i,k} <- r_s.{i,k};
              hfm_o.{m,i,k} <- r_o.{i,k}
            done
          done
        done
      in
      hf_s, hf_o, hfm_s, hfm_o, j, l
  in

  let tasks = 
    let rec next accu = function
      | _, 0 -> accu
      | 0, l -> next accu  (mo_num, l-1)
      | j, l -> next ((j,l) :: accu)  ((j-1), l)
    in
    next [] (mo_num, mo_num)
    |> Stream.of_list
  in


  Farm.run ~f:task ~ordered:true tasks
  |> Stream.iter (fun (hf_s, hf_o, hfm_s, hfm_o, j, l) ->
    for k=1 to mo_num do
      for i=1 to mo_num do
        result_s.{i,j,k,l} <- hf_s.{i,k} ;
        result_o.{i,j,k,l} <- hf_o.{i,k} ;
        for m=1 to mo_num do
          resultm_s.{m,i,j,k,l} <- hfm_s.{m,i,k} ;
          resultm_o.{m,i,j,k,l} <- hfm_o.{m,i,k}
        done
      done
    done );



  (*
  for l=1 to mo_num do
    for k=1 to mo_num do
      for j=1 to mo_num do
        for i=1 to mo_num do
          Printf.printf "%d %d %d %d %e\n" i j k l result.{i,j,k,l} 
        done
      done
    done
  done;
  Printf.printf "%!";
  *)

  Parallel.broadcast (lazy (result_s, result_o, resultm_s, resultm_o) )