open Util
open Constants
open Lacaml.D

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


type hartree_fock_data =
  {
    iteration       :  int              ;
    coefficients    :  Mat.t   option   ;
    eigenvalues     :  Vec.t   option   ;
    error           :  float   option   ;
    diis            :  DIIS.t  option   ;
    energy          :  float   option   ;
    density         :  Mat.t   option   ;
    fock            :  Fock.t  option   ;
  }

let empty =
  {
    iteration       =  0                ;
    coefficients    =  None             ;
    eigenvalues     =  None             ;
    error           =  None             ;
    diis            =  None             ;
    energy          =  None             ;
    density         =  None             ;
    fock            =  None             ;
  }

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


  (* A single SCF iteration *)
  let scf_iteration data =

    let  nSCF         =  data.iteration + 1
    and  m_C          =  of_some data.coefficients
    and  m_P_prev     =  data.density
    and  fock_prev    =  data.fock
    and  diis         =  
      match data.diis with
      | Some diis -> diis
      | None      -> DIIS.make ()
    and  threshold    =  
      match data.error with
      | Some error -> error
      | None       -> threshold_SCF *. 2.
    in

    (* 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, m_P_prev, threshold > 100. *. threshold_SCF with
      | Some fock_prev, Some m_P_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
    {
      iteration      = nSCF              ;
      eigenvalues    = Some eigenvalues  ;
      coefficients   = Some m_C          ;
      error          = Some error        ;
      diis           = Some diis         ;
      energy         = Some energy       ;
      density        = Some m_P          ;
      fock           = Some fock         ;
    }

  in


  let rec make_iterations_list data =

    let energy_prev = data.energy in

    (** Perform SCF iteration *)
    let data = scf_iteration data in

    (** Check convergence *)
    let converged, error =
      match data.error with
      | None -> false, 0.
      | Some error -> (data.iteration = max_scf || error < threshold_SCF), error
    in

    (** Print values *)
    let nSCF = data.iteration in

    let energy = of_some data.energy in

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

    if converged then
      [ data ]
    else
      { empty  with 
        iteration   = data.iteration;
        energy      = data.energy ;
        eigenvalues = data.eigenvalues ;
        error       = data.error ;
      } :: (make_iterations_list data)
  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 iterations_list =
    make_iterations_list { empty with coefficients = Some m_C }
  in

  let iterations, data = 
    List.map (fun data -> 
        let gap =
          let eigenvalues = of_some data.eigenvalues in
          if nocc < Vec.dim eigenvalues then
            eigenvalues.{nocc+1} -. eigenvalues.{nocc} 
          else 0.
        and energy = of_some data.energy
        and error  = of_some data.error
        in
        (energy, error, gap)
      ) iterations_list
    |> Array.of_list,
    List.hd (List.rev iterations_list)
  in


  let energy = of_some data.energy  in
  let m_P    = of_some data.density in
  let fock   = of_some data.fock    in
  let m_J    = Fock.coulomb fock    in
  let m_K    = Fock.exchange fock   in

  HartreeFock_type.(
    RHF { 
      simulation;
      nocc;
      guess ;
      eigenvectors = of_some data.coefficients ;
      eigenvalues  = of_some data.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;
    }
  )