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

open Lacaml.D

module Fis = FourIdxStorage


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

type t = {
  simulation : Simulation.t ;
  aux_basis  : MOBasis.t ;
  hf12     : q ; (* hf12.{i,j,k,l} = \sum_{ab} <ij|H|ab> <ab|F12|kl>
                                   + \sum_{am} <ij|H|am> <am|F12|kl>
                                   + \sum_{bm} <ij|H|mb> <mb|F12|kl>
                 *)
  hf12_anti: q ; (* hf12_anti.{i,j,k,l} = 
                      \sum_{ab} (<ij|H|ab> - <ij|H|ba>) (<ab|F12|kl> - <ab|F12|lk>)
                    + \sum_{am} (<ij|H|am> - <ij|H|ma>) (<am|F12|kl> - <am|F12|lk>)
                    + \sum_{bm} (<ij|H|mb> - <ij|H|bm>) (<mb|F12|kl> - <mb|F12|lk>)
                    *)
  hf12_single     : q ; (* hf12.{m,i,j,k,l} = \sum_{a} <ij|H|am> <am|F12|kl> *)
  hf12_single_anti: q ; (* hf12_anti.{m,i,j,k,l} = 
                           \sum_{ab} (<ij|H|am> - <ij|H|ma>) (<am|F12|kl> - <am|F12|lk>) *)
}


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 simulation =
    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

  let aux_basis = 
    MOBasis.of_mo_basis simulation 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 hf12, hf12_anti, hf12_single, hf12_single_anti =
    let create4 n =
      Bigarray.Genarray.create Float64 Bigarray.fortran_layout [| n ; n ; n ; n |] 
    in
    let create5 n =
      Bigarray.Genarray.create Float64 Bigarray.fortran_layout [| n ; n ; n ; n ; n |] 
    in
    create4 mo_num ,
    create4 mo_num ,
    create5 mo_num ,
    create5 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
        hf12.{i,j,k,l}      <- hf_o.{i,k} ;
        hf12_anti.{i,j,k,l} <- hf_s.{i,k} ;
        for m=1 to mo_num do
          hf12_single.{m,i,j,k,l}      <- hfm_o.{m,i,k} ;
          hf12_single_anti.{m,i,j,k,l} <- hfm_s.{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 "%!";
  *)

  let result = 
    { simulation ; aux_basis ;
      hf12 ; hf12_anti ;
      hf12_single ; hf12_single_anti
    } 
  in
  Parallel.broadcast (lazy result)