open Lacaml.D
open Util
open Constants

(** One-electron orthogonal basis set, corresponding to Molecular Orbitals. *)

module HF = HartreeFock_type
module Si = Simulation

type mo_class =
| Core of int
| Inactive of int
| Active of int
| Virtual of int
| Deleted of int

type mo_type =
| RHF | ROHF | CASSCF
| Natural of string
| Localized of string

type t = 
{
  ao_basis        : AOBasis.t;         (* Atomic basis set on which the MOs are built. *)
  mo_type         : mo_type;           (* Kind of MOs (RHF, CASSCF, Localized...       *)
  mo_class        : mo_class array;     (* CI-Class of the MOs *)
  mo_occupation   : Vec.t;             (* Occupation numbers *)
  mo_coef         : Mat.t;             (* Matrix of the MO coefficients in the AO basis *)
  eN_ints         : NucInt.t lazy_t;   (* Electron-nucleus potential integrals *)
  ee_ints         : ERI.t lazy_t;      (* Electron-electron potential integrals *)
  kin_ints        : KinInt.t lazy_t;   (* Kinetic energy integrals *)
}



let mo_matrix_of_ao_matrix ~mo_coef ao_matrix = 
  xt_o_x ~x:mo_coef ~o:ao_matrix


let ao_matrix_of_mo_matrix ~mo_coef ~ao_overlap mo_matrix =
  let sc = gemm ao_overlap mo_coef  in
  gemm sc @@
  gemm ~transb:`T mo_matrix sc


let four_index_transform ~mo_coef eri_ao =

  let ao_num = Mat.dim1 mo_coef in
  let mo_num = Mat.dim2 mo_coef in
  let eri_mo = ERI.create ~size:mo_num `Dense in

  let mo_num_2 = mo_num * mo_num in
  let ao_num_2 = ao_num * ao_num in
  let ao_mo_num = ao_num * mo_num in

  let range_mo = list_range ~start:1 mo_num in
  let range_ao = list_range ~start:1 ao_num in

  let u = Mat.create mo_num_2 mo_num
  and o = Mat.create ao_num ao_num_2 
  and p = Mat.create ao_num_2 mo_num
  and q = Mat.create ao_mo_num mo_num
  in
  Printf.eprintf "Transforming %d integrals : %!" mo_num;
  List.iter (fun delta ->
      Printf.eprintf "%d %!" delta;
      Mat.fill u 0.;

      List.iter (fun l ->
          if abs_float mo_coef.{l,delta} > epsilon then
          begin
            let jk = ref 0 in
            List.iter (fun k ->
                List.iter (fun j ->
                    incr jk;
                    ERI.get_chem_all_i eri_ao ~j ~k ~l
                    |> Array.iteri (fun i x -> o.{i+1,!jk} <- x)
                ) range_ao
            ) range_ao;
            (* o_i_jk *)

            let p =
              gemm ~transa:`T ~c:p o mo_coef
              (* p_jk_alpha  = \sum_i o_i_jk c_i_alpha *)
            in
            let p' = 
              Bigarray.reshape_2 (Bigarray.genarray_of_array2 p) ao_num ao_mo_num
              (* p_j_kalpha *)
            in

            let q =
              gemm ~transa:`T ~c:q p' mo_coef
              (* q_kalpha_beta  = \sum_j p_j_kalpha c_j_beta *)
            in
            let q' =
              Bigarray.reshape_2 (Bigarray.genarray_of_array2 q) ao_num mo_num_2
              (* q_k_alphabeta  = \sum_j p_j_kalpha c_j_beta *)
            in

            ignore @@
              gemm ~transa:`T ~beta:1. ~alpha:mo_coef.{l,delta} ~c:u q' mo_coef ;
              (* u_alphabeta_gamma = \sum_k q_k_alphabeta c_k_gamma *)
          end
      ) range_ao;
      let u =
        Bigarray.reshape 
          (Bigarray.genarray_of_array2 u)
          [| mo_num ; mo_num ; mo_num |]
        |> Bigarray.array3_of_genarray 
      in
      List.iter (fun gamma ->
          List.iter (fun beta ->
              List.iter (fun alpha ->
                 let x = u.{alpha,beta,gamma} in
                 if x <> 0. then
                    ERI.set_chem eri_mo alpha beta gamma delta x
              ) (list_range ~start:1 beta) 
          ) range_mo
      ) (list_range ~start:1 delta)  
  ) range_mo;
  Printf.eprintf "\n%!";

  eri_mo


let make ~ao_basis ~mo_type ~mo_class ~mo_occupation ~mo_coef ()  =
  let eN_ints = lazy (
    Lazy.force ao_basis.AOBasis.eN_ints 
    |> NucInt.matrix
    |> mo_matrix_of_ao_matrix ~mo_coef   
    |> NucInt.of_matrix
  )
  and kin_ints = lazy (
    Lazy.force ao_basis.AOBasis.kin_ints 
    |> KinInt.matrix
    |> mo_matrix_of_ao_matrix ~mo_coef  
    |> KinInt.of_matrix
  )
  and ee_ints = lazy (
    Lazy.force ao_basis.AOBasis.ee_ints 
    |> four_index_transform ~mo_coef   
  )
  in
  { ao_basis ; mo_type ; mo_class ; mo_occupation ; mo_coef   ;
    eN_ints ; ee_ints ; kin_ints }


let of_rhf ~frozen_core hf =
  let simulation = hf.HF.simulation in
  let nocc = hf.HF.nocc in
  let ncore =
    if frozen_core then Nuclei.small_core simulation.Si.nuclei
    else 0
  in
  let mo_num = Vec.dim hf.HF.eigenvalues in

  let ao_basis = simulation.Si.ao_basis in
  let mo_type  = RHF in
  let mo_class =
    Array.init mo_num (fun i ->
      if (i < ncore) then Core i
      else 
        if (i < nocc ) then Inactive i
        else Virtual i)
  in
  let mo_occupation =
      Array.init mo_num (fun i ->
        if i < nocc then 2. else 0.)
      |> Vec.of_array
  in
  let mo_coef = hf.HF.eigenvectors in
  let result = 
    make ~ao_basis ~mo_type ~mo_class ~mo_occupation ~mo_coef ()
  in

  let () =
    let e = ref 0. in
    let t = KinInt.matrix (Lazy.force result.kin_ints) in
    let v = NucInt.matrix (Lazy.force result.eN_ints) in
    let g = Lazy.force result.ee_ints in
    for i = 1 to 5 do
      e := !e +. 2. *. (t.{i,i} +. v.{i,i})
    done;
    Printf.printf "Energy Mono = %20.15f\n" !e;
    let e2 = ref 0. in
    for i = 1 to 5 do
      for j = i+1 to 5 do
        e2 := !e2 +. 2. *. (ERI.get_phys g i j i j -.
                            ERI.get_phys g i j j i)
      done;
    done;
    for i = 1 to 5 do
      for j = 1 to 5 do
        e2 := !e2 +. ERI.get_phys g i j i j 
      done;
    done;
    Printf.printf "Energy bi = %20.15f\n" !e2;
    Printf.printf "Energy = %20.15f\n" (simulation.Si.nuclear_repulsion +. !e +. !e2)
  in
  result


let of_hartree_fock ~frozen_core = function
| HF.RHF hf -> of_rhf ~frozen_core hf
| _ -> assert false