open Lacaml.D
open Util

(** 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 = 
  gemm ~transa:`T mo_coef   @@
  gemm ao_matrix mo_coef  


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

          let jk = ref 0 in
          List.iter (fun k ->
              List.iter (fun j ->
                  incr jk;
                  ERI.get_phys_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. ~c:u q' mo_coef 
            (* u_alphabeta_gamma = \sum_k q_k_alphabeta c_k_gamma *)

      ) 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_phys eri_mo alpha beta gamma delta x
              ) range_mo
          ) range_mo
      ) range_mo
  ) 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
  make ~ao_basis ~mo_type ~mo_class ~mo_occupation ~mo_coef ()


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