open Util
open Constants
open Lacaml.D

module Si = Simulation
module El = Electrons
module Ao = AOBasis
module Ov = Overlap

let make ~guess ~max_scf ~level_shift ~threshold_SCF simulation =
  (* Number of occupied MOs *)
  let nocc =
    El.n_alfa @@ Si.electrons simulation
  in

  let nuclear_repulsion = 
    Si.nuclear_repulsion simulation
  in
  
  let ao_basis = 
    Si.ao_basis simulation
  in

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

  (* Orthogonalization matrix *)
  let m_X = 
    Ao.ortho ao_basis
  in


  (* Overlap matrix *)
  let m_S =
    Ao.overlap ao_basis
    |> Ov.matrix
  in

  let m_T = Ao.kin_ints ao_basis |> KinInt.matrix
  and m_V = Ao.eN_ints  ao_basis |> NucInt.matrix
  in

  (* Level shift in MO basis *)
  let m_LSmo =
    Array.init (Mat.dim2 m_X) (fun i ->
      if i > nocc then level_shift else 0.)
    |> Vec.of_array
    |> Mat.of_diag
  in

  (* SCF iterations *)
  let rec loop nSCF iterations energy_prev m_C m_P_prev fock_prev threshold diis =

      (* 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 fock = 
        match fock_prev, threshold > 100. *. threshold_SCF with
        | Some fock_prev, true -> 
            let threshold = 1.e-8 in
            Fock.make_rhf  ~density:(Mat.sub m_P m_P_prev) ~threshold ao_basis
            |> Fock.add fock_prev
        | _ -> Fock.make_rhf  ~density:m_P ao_basis
      in

      let m_F, m_Hc, m_J, m_K =
        let x = fock in
        Fock.(fock x, core x, coulomb x, exchange x)
      in

      (* Add level shift in AO basis *)
      let m_F =
        let m_SC =
          gemm  m_S  m_C
        in
        gemm  m_SC (gemm  m_LSmo  m_SC  ~transb:`T)
        |> Mat.add m_F
      in


      (* Fock matrix in orthogonal basis *)
      let m_F_ortho =
        xt_o_x m_F m_X
      in

      let error_fock = 
        let fps =
          gemm  m_F  (gemm  m_P  m_S)
        and spf =
          gemm  m_S  (gemm  m_P  m_F)
        in
        xt_o_x  (Mat.sub  fps  spf) m_X
      in

      let diis = 
        DIIS.append ~p:(Mat.as_vec m_F_ortho) ~e:(Mat.as_vec error_fock) diis
      in

      let m_F_diis =
        let x = 
          Bigarray.genarray_of_array1 (DIIS.next diis)
        in
        Bigarray.reshape_2 x (Mat.dim1 m_F_ortho) (Mat.dim2 m_F_ortho)
      in


      (* MOs in orthogonal MO basis *)
      let m_C', eigenvalues =
        diagonalize_symm  m_F_diis
      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 error =
        error_fock 
        |> Mat.as_vec 
        |> amax
        |> abs_float
      in

      let converged =
        nSCF = max_scf  || error < threshold_SCF
      in

      let gap =
        if nocc < Vec.dim eigenvalues then
          eigenvalues.{nocc+1} -. eigenvalues.{nocc} 
        else 0.
      in

      let () = 
        match energy_prev with
        | Some energy_prev ->
          Printf.eprintf "%3d  %16.10f  %16.10f %11.4e  %10.4f\n%!" nSCF energy (energy -. energy_prev) error gap
        | None -> 
          Printf.eprintf "%3d  %16.10f  %16s %11.4e  %10.4f\n%!" nSCF energy "" error gap
      in

      if not converged then
          loop (nSCF+1) ( (energy, error, gap) :: iterations) (Some energy) m_C m_P (Some fock) error diis
      else
          let iterations = 
            List.rev ( (energy, error, gap) :: iterations )
            |> Array.of_list
          in
          HartreeFock_type.(RHF
          { 
            simulation;
            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;
            occupation = Mat.copy_diag m_P;
            })
  in


  (* Guess coefficients *)
  let m_H = 
    match guess with
    | Guess.Hcore m_H -> m_H 
    | Guess.Huckel 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

  let diis = DIIS.make () in
  loop 1 [] None m_C m_C None threshold_SCF diis