open Util
open Lacaml.D
open Simulation

let make ?guess:(guess=`Hcore) ?max_scf:(max_scf=64)
  ?threshold_SCF:(threshold_SCF=1.e-6) simulation =

  (* Number of occupied MOs *)
  let nocc =
    simulation.electrons.Electrons.n_alpha
  in

  let nuclear_repulsion = 
    simulation.nuclear_repulsion
  in

  (* Initial guess *)
  let guess = 
    Guess.make ~guess simulation 
  in

  (* Orthogonalization matrix *)
  let m_X = 
   Lazy.force simulation.overlap_ortho 
  in


  (* Overlap matrix *)
  let m_S =
    Lazy.force  simulation.overlap
  in

  let m_T = Lazy.force simulation.kin_ints
  and m_V = Lazy.force simulation.eN_ints
  in

  (* SCF iterations *)
  let rec loop nSCF iterations m_C =

      (* Density matrix over nocc occupied MOs *)
      let m_P =
        gemm  ~alpha:2.  ~transb:`T  ~k:nocc  m_C  m_C
      in

      (* Fock matrix in AO basis *)
      let m_F, m_Hc, m_J, m_K =
        let x = 
          Fock.make  ~density:m_P  simulation
        in
        x.Fock.fock, x.Fock.core, x.Fock.coulomb, x.Fock.exchange
      in

      (* Fock matrix in MO basis *)
      let m_Fmo =
        xt_o_x  m_F  m_X
      in

      (* MOs in old MO basis *)
      let m_C', eigenvalues =
        diagonalize_symm  m_Fmo
      in

      (* MOs in AO basis *)
      let m_C =
        gemm  m_X  m_C'
      in

      (* Hartree-Fock energy *)
      let energy =
        nuclear_repulsion +. 0.5 *.
          Mat.gemm_trace  m_P  (Mat.add  m_Hc  m_F)
      in

      (* Convergence criterion *)
      let commutator =
        let fps =
          gemm  m_F  (gemm  m_P  m_S)
        and spf =
          gemm  m_S  (gemm  m_P  m_F)
        in
        Mat.sub  fps  spf
        |> Mat.as_vec
        |> amax
      in

      let converged =
        nSCF = max_scf  || (abs_float commutator) < threshold_SCF
      in

      let gap =
        eigenvalues.{nocc+1} -. eigenvalues.{nocc};
      in

      Printf.printf "%d  %16.10f  %11.4e  %10.4f\n%!" nSCF energy commutator gap;

      if not converged then
          loop (nSCF+1) ( (energy, commutator, gap) :: iterations) m_C
      else
          let iterations = 
            List.rev ( (energy, commutator, gap) :: iterations )
            |> Array.of_list
          in
          { HartreeFock_type.
            nocc;
            guess ;
            eigenvectors = m_C ;
            eigenvalues ;
            energy ;
            nuclear_repulsion;
            iterations ;
            kin_energy = Mat.gemm_trace m_P m_T;
            eN_energy = Mat.gemm_trace m_P m_V;
            coulomb_energy  = 0.5 *. Mat.gemm_trace m_P m_J;
            exchange_energy = 0.5 *. Mat.gemm_trace m_P m_K;
          }
  in


  (* Guess coefficients *)
  let m_H = 
    match guess with
    | Guess.Hcore m_H -> m_H 
  in
  let m_Hmo =
    xt_o_x  m_H  m_X
  in
  let m_C', _ =
    diagonalize_symm  m_Hmo
  in
  let m_C =
    gemm  m_X  m_C'
  in

  loop 1 [] m_C