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

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

  (* Orthogonalization matrix *)
  let m_X = 
    match Lazy.force simulation.overlap_ortho with
    | Lowdin    x -> x
    | Canonical x -> x
    | Svd       x -> x
  in

  (* Core Hamiltonian *)
  let m_Hc =
    let m_T = Lazy.force simulation.kin_ints
    and m_V = Lazy.force simulation.eN_ints
    in Mat.add m_T m_V
  in

  (* Overlap matrix *)
  let m_S =
    Lazy.force  simulation.overlap
  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 =
        Fock.make  ~density:m_P  simulation
      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 =
        simulation.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

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

      if not converged then
          loop (nSCF+1) (energy :: iterations) m_C
      else
          let iterations = 
            List.rev iterations
            |> Array.of_list
          in
          { HartreeFock_type.
            guess ;
            eigenvectors = m_C ;
            eigenvalues ;
            energy = iterations.(Array.length iterations - 1) ;
            iterations ;
          }
  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