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_type =
  | RHF | ROHF | CASSCF
  | Natural of string
  | Localized of string

type t = 
  {
    simulation      : Simulation.t;      (* Simulation which produced the MOs *)
    mo_type         : mo_type;           (* Kind of MOs (RHF, CASSCF, Localized...) *)
    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 *)
    one_e_ints      : Mat.t lazy_t;      (* Kinetic energy integrals *)
  }


let size t =
  Mat.dim2 t.mo_coef 

let simulation    t = t.simulation
let mo_type       t = t.mo_type
let ao_basis      t = Si.ao_basis t.simulation
let mo_occupation t = t.mo_occupation
let mo_coef       t = t.mo_coef
let eN_ints       t = Lazy.force t.eN_ints
let ee_ints       t = Lazy.force t.ee_ints
let kin_ints      t = Lazy.force t.kin_ints
let two_e_ints    t = Lazy.force t.ee_ints
let one_e_ints    t = Lazy.force t.one_e_ints

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 1 mo_num in
  let range_ao = list_range 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 1 beta) 
            ) range_mo
        ) (list_range 1 delta)  
    ) range_mo;
  Printf.eprintf "\n%!";
  eri_mo


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


let of_rhf hf =
  let mo_num        = Vec.dim hf.HF.eigenvalues in
  let mo_coef       = hf.HF.eigenvectors in
  let simulation    = hf.HF.simulation in
  let mo_type       = RHF in
  let nocc          = hf.HF.nocc in
  let mo_occupation =
    Array.init mo_num (fun i ->
        if i < nocc then 2. else 0.)
    |> Vec.of_array
  in
  let result = 
    make ~simulation ~mo_type ~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 one-e = %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 two-e = %20.15f\n" !e2;
    Printf.printf "Energy       = %20.15f\n" (Si.nuclear_repulsion simulation +. !e +. !e2)
  in
  result


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