open Util
open Constants
open Lacaml.D

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

let make ?guess:(guess=`Huckel) ?max_scf:(max_scf=64) ?level_shift:(level_shift=0.1)
    ?threshold_SCF:(threshold_SCF=1.e-5) simulation =

  (* Number of occupied MOs *)
  let n_alfa =
    El.n_alfa @@ Si.electrons simulation
  in

  let nocc = n_alfa in

  let n_beta = 
    El.n_beta @@ 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 > n_alfa 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_a_prev m_P_b_prev fock_a_prev fock_b_prev threshold diis =


    (* Density matrix *)
    let m_P_a =
      gemm  ~alpha:1.  ~transb:`T  ~k:n_alfa  m_C  m_C
    in

    let m_P_b =
      gemm  ~alpha:1.  ~transb:`T  ~k:n_beta  m_C  m_C
    in

    let m_P =
      Mat.add m_P_a m_P_b
    in

    (* Fock matrix in AO basis *)
    let fock_a = 
      match fock_a_prev, threshold > 100. *. threshold_SCF with
      | Some fock_a_prev, true -> 
        let threshold = 1.e-8 in
        Fock.make_uhf  ~density_same:(Mat.sub m_P_a m_P_a_prev) ~density_other:(Mat.sub m_P_b m_P_b_prev) ~threshold ao_basis
        |> Fock.add fock_a_prev
      | _ -> Fock.make_uhf  ~density_same:m_P_a  ~density_other:m_P_b ao_basis
    in

    let fock_b = 
      match fock_b_prev, threshold > 100. *. threshold_SCF with
      | Some fock_b_prev, true -> 
        let threshold = 1.e-8 in
        Fock.make_uhf  ~density_same:(Mat.sub m_P_b m_P_b_prev) ~density_other:(Mat.sub m_P_a m_P_a_prev) ~threshold ao_basis
        |> Fock.add fock_b_prev
      | _ -> Fock.make_uhf  ~density_same:m_P_b  ~density_other:m_P_a ao_basis
    in

    let m_F_a, m_Hc_a, m_J_a, m_K_a =
      let x = fock_a in
      Fock.(fock x, core x, coulomb x, exchange x)
    in

    let m_F_b, m_Hc_b, m_J_b, m_K_b =
      let x = fock_b in
      Fock.(fock x, core x, coulomb x, exchange x)
    in


    let m_F_mo_a = 
      xt_o_x ~o:m_F_a ~x:m_C
    in

    let m_F_mo_b = 
      xt_o_x ~o:m_F_b ~x:m_C
    in

    let m_F_mo  =
      let space k =
        if k <= n_beta then
          `Core
        else if k <= n_alfa then
          `Active
        else
          `Virtual
      in
      Array.init (Mat.dim2 m_F_mo_a) (fun i ->
          let i = i+1 in
          Array.init (Mat.dim1 m_F_mo_a) (fun j ->
              let j = j+1 in
              match (space i), (space j) with
              |  `Core     ,  `Core     ->  
                0.5 *. (m_F_mo_a.{i,j} +. m_F_mo_b.{i,j}) -.
                (m_F_mo_b.{i,j} -. m_F_mo_a.{i,j})

              |  `Active   ,  `Core
              |  `Core     ,  `Active   ->
                (*
                0.5 *. (m_F_mo_a.{i,j} +. m_F_mo_b.{i,j}) +.
                0.5 *. (m_F_mo_b.{i,j} -. m_F_mo_a.{i,j})
                *)
                m_F_mo_b.{i,j}

              |  `Core     ,  `Virtual
              |  `Virtual  ,  `Core     
              |  `Active   ,  `Active   ->
                0.5 *. (m_F_mo_a.{i,j} +. m_F_mo_b.{i,j})

              |  `Virtual  ,  `Active   
              |  `Active   ,  `Virtual  ->
                (*
                0.5 *. (m_F_mo_a.{i,j} +. m_F_mo_b.{i,j}) -.
                0.5 *. (m_F_mo_b.{i,j} -. m_F_mo_a.{i,j})
                *)
                m_F_mo_a.{i,j}

              |  `Virtual  ,  `Virtual  ->
                0.5 *. (m_F_mo_a.{i,j} +. m_F_mo_b.{i,j}) +.
                (m_F_mo_b.{i,j} -. m_F_mo_a.{i,j})
            ) )
      |> Mat.of_array
    in

    let m_SC =
      gemm  m_S  m_C
    in

    let m_F = 
      x_o_xt ~x:m_SC ~o:m_F_mo
    in


    (* Add level shift in AO basis *)
    let m_F =
      x_o_xt ~x:m_SC ~o:m_LSmo
      |> 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_a  (Mat.add  m_Hc_a  m_F_a) +.
                                     Mat.gemm_trace  m_P_b  (Mat.add  m_Hc_b  m_F_b) )
    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_a m_P_b (Some fock_a) (Some fock_b) error diis
    else
      let iterations = 
        List.rev ( (energy, error, gap) :: iterations )
        |> Array.of_list
      in
      HartreeFock_type.(ROHF
                          { 
                            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_a m_J_a) +. 
                                              0.5 *. (Mat.gemm_trace m_P_b m_J_b);
                            exchange_energy = 0.5 *. (Mat.gemm_trace m_P_a m_K_a) +.
                                              0.5 *. (Mat.gemm_trace m_P_b m_K_b);
                            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 m_C None None threshold_SCF diis