open Lacaml.D
open Util
open Constants


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   ;
    density_a       :  Mat.t   option   ;
    density_b       :  Mat.t   option   ;
    fock            :  Fock.t  option   ;
    fock_a          :  Fock.t  option   ;
    fock_b          :  Fock.t  option   ;
  }

type hartree_fock_kind  =
| RHF  (** Restricted Hartree-Fock *)
| ROHF (** Restricted Open-shell Hartree-Fock *)
| UHF  (** Unrestricted Hartree-Fock *)

type t =
  {
    kind              : hartree_fock_kind;
    simulation        : Simulation.t;
    guess             : Guess.t;
    data              : hartree_fock_data option lazy_t array;
    nocc              : int   ;
  }


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






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


let  kind               t  =  t.kind
let  simulation         t  =  t.simulation
let  guess              t  =  t.guess
let  nocc               t  =  t.nocc



let n_iterations t =
  Array.fold_left (fun accu x ->
    match Lazy.force x with
    | Some x -> accu + 1
    | None -> accu
    ) 0 t.data


let last_iteration t =
  of_some @@ Lazy.force (t.data.(n_iterations t - 1))

let eigenvectors t =
  let data = last_iteration t in
  of_some data.coefficients

let eigenvalues t =
  let data = last_iteration t in
  of_some data.eigenvalues

let density t =
  let data = last_iteration t in
  match kind t with
  | RHF  -> of_some data.density
  | ROHF -> Mat.add (of_some data.density_a) (of_some data.density_b)
  | _    -> failwith "Not implemented"

let occupation t =
  density t
  |> Mat.copy_diag ~n:(Mat.dim2 @@ eigenvectors t)


let energy t = 
  let data = last_iteration t in
  of_some data.energy


let nuclear_repulsion t =
  Si.nuclear_repulsion (simulation t)


let ao_basis t =
  Si.ao_basis (simulation t)


let kin_energy t =
  let m_T    =
    ao_basis t
    |> Ao.kin_ints 
    |> KinInt.matrix
  in
  let m_P = density t in
  Mat.gemm_trace m_P m_T


let eN_energy t =
  let m_V =
    ao_basis t
    |> Ao.eN_ints 
    |> NucInt.matrix
  in
  let m_P = density t in
  Mat.gemm_trace m_P m_V


let coulomb_energy t =
  let data =
    last_iteration t
  in
  match kind t with
  | RHF ->  let  m_P   =  of_some         data.density  in
            let  fock  =  of_some         data.fock     in
            let  m_J   =  Fock.coulomb    fock          in
            0.5 *. Mat.gemm_trace m_P m_J

  | ROHF -> let  m_P_a  =  of_some         data.density_a in
            let  m_P_b  =  of_some         data.density_b in
            let  fock_a =  of_some         data.fock_a    in
            let  fock_b =  of_some         data.fock_b    in
            let  m_J_a  =  Fock.coulomb    fock_a         in
            let  m_J_b  =  Fock.coulomb    fock_b         in
            0.5 *. ( (Mat.gemm_trace m_P_a m_J_a) +. (Mat.gemm_trace m_P_b m_J_b) )

  | _    -> failwith "Not implemented"


let exchange_energy t = 
  let data =
    last_iteration t
  in
  match kind t with
  | RHF ->  let  m_P   =  of_some         data.density  in
            let  fock  =  of_some         data.fock     in
            let  m_K   =  Fock.exchange   fock          in
            0.5 *. Mat.gemm_trace m_P m_K

  | ROHF -> let  m_P_a  =  of_some         data.density_a in
            let  m_P_b  =  of_some         data.density_b in
            let  fock_a =  of_some         data.fock_a    in
            let  fock_b =  of_some         data.fock_b    in
            let  m_K_a  =  Fock.exchange   fock_a         in
            let  m_K_b  =  Fock.exchange   fock_b         in
            0.5 *. ( (Mat.gemm_trace m_P_a m_K_a) +. (Mat.gemm_trace m_P_b m_K_b) )

  | _    -> failwith "Not implemented"








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


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

  let nocc = n_alfa in

  let kind = 
    match kind with
    | Some kind -> kind
    | None -> if (n_alfa = n_beta) then RHF else ROHF
  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

  (* 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

  (* Guess coefficients *)
  let m_C =
    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
    gemm  m_X  m_C'
  in


  (* A single SCF iteration *)
  let scf_iteration_rhf 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
    { empty with
      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 scf_iteration_rohf data = 

    let  nSCF         =  data.iteration + 1
    and  m_C          =  of_some data.coefficients
    and  m_P_a_prev   =  data.density_a
    and  m_P_b_prev   =  data.density_b
    and  fock_a_prev  =  data.fock_a
    and  fock_b_prev  =  data.fock_b
    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 *)
    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 @@ of_some m_P_a_prev) ~density_other:(Mat.sub m_P_b @@ of_some 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 @@ of_some m_P_b_prev) ~density_other:(Mat.sub m_P_a @@ of_some 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   ->
                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  ->
                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
    { empty with
      iteration      = nSCF              ;
      eigenvalues    = Some eigenvalues  ;
      coefficients   = Some m_C          ;
      error          = Some error        ;
      diis           = Some diis         ;
      energy         = Some energy       ;
      density_a      = Some m_P_a        ;
      density_b      = Some m_P_b        ;
      fock_a         = Some fock_a       ;
      fock_b         = Some fock_b       ;
    }

  in


  let scf_iteration data =
    match kind with
    | RHF  -> scf_iteration_rhf data
    | ROHF -> scf_iteration_rohf data
    | _    -> failwith "Not implemented"
  in

  let array_data =
      
    let storage =
      Array.make max_scf None
    in

    let rec iteration = function
      | 0 -> Some (scf_iteration { empty with coefficients = Some m_C })
      | n -> begin
          match storage.(n) with
          | Some result -> Some result
          | None ->
            begin
              let data =
                (iteration (n-1)) 
              in
              match data with
              | None -> None
              | Some data -> 
                begin
                  (** Check convergence *)
                  let converged, error =
                    match data.error with
                    | None -> false, 0.
                    | Some error -> (data.iteration = max_scf || error < threshold_SCF), error
                  in
                  if converged then
                    None
                  else
                    begin
                      storage.(n-1) <- Some { empty with
                                              iteration   = data.iteration;
                                              energy      = data.energy ;
                                              eigenvalues = data.eigenvalues ;
                                              error       = data.error ;
                                            };
                      storage.(n) <- Some (scf_iteration data);
                      storage.(n);
                    end
                end
            end
        end
    in
    Array.init max_scf (fun i -> lazy (iteration i))
  in
  { 
      kind;
      simulation;
      guess ;
      data = array_data;
      nocc;
  }










let linewidth = 60

let pp_iterations ppf t =
  Format.fprintf ppf "@[%4s%s@]@." "" (Printing.line linewidth);
  Format.fprintf ppf "@[%4s@[%5s@]@,@[%16s@]@,@[%16s@]@,@[%11s@]@]@." 
                      ""    "#" "HF energy  " "Convergence" "HOMO-LUMO";
  Format.fprintf ppf "@[%4s%s@]@." "" (Printing.line linewidth);
  let nocc = nocc t in
  Array.iter (fun data ->
    let data = Lazy.force data in
    match data with
    | None -> ()
    | Some data ->
      let e = of_some data.eigenvalues in
      let gap = e.{nocc+1} -. e.{nocc} in
      begin
        Format.fprintf ppf "@[%4s@[%5d@]@,@[%16.8f@]@,@[%16.4e@]@,@[%11.4f@]@]@." ""
          (data.iteration) (of_some data.energy) (of_some data.error) gap;
      end
    ) t.data;
  Format.fprintf ppf "@[%4s%s@]@." "" (Printing.line linewidth)


let pp_summary ppf t =
  let print text value = 
    Format.fprintf ppf "@[@[%30s@]@,@[%16.10f@]@]@;" text value;
  and line () = 
    Format.fprintf ppf "@[  %s  @]@;" (Printing.line (linewidth-4));
  in
  let ha_to_ev = Constants.ha_to_ev in
  let e = eigenvalues t in

  Format.fprintf ppf "@[%s@]@;" (Printing.line ~c:'=' linewidth);
  Format.fprintf ppf "@[<v>";
  print "One-electron energy" (kin_energy t +. eN_energy t) ;
  print "Kinetic"   (kin_energy t) ;
  print "Potential" (eN_energy  t) ;
  line () ;
  print "Two-electron energy" (coulomb_energy t +. exchange_energy t) ;
  print "Coulomb"  (coulomb_energy  t) ;
  print "Exchange" (exchange_energy t) ;
  line ();
  print "HF HOMO" (ha_to_ev *. e.{nocc t});
  print "HF LUMO" (ha_to_ev *. e.{nocc t + 1});
  print "HF LUMO-LUMO" (ha_to_ev *. (e.{nocc t + 1} -. e.{nocc t }));
  line ();
  print "Electronic   energy" (energy t -. nuclear_repulsion t) ;
  print "Nuclear   repulsion" (nuclear_repulsion t) ;
  print "Hartree-Fock energy" (energy t) ;
  Format.fprintf ppf "@]";
  Format.fprintf ppf "@[%s@]@;" (Printing.line ~c:'=' linewidth)


let pp_hf ppf t =
  Format.fprintf ppf "@.@[%s@]@." (Printing.line ~c:'=' 70);
  Format.fprintf ppf "@[%34s %-34s@]@." (match t.kind with
  | UHF  -> "Unrestricted"
  | RHF  -> "Restricted"
  | ROHF -> "Restricted Open-shell") "Hartree-Fock";
  Format.fprintf ppf "@[%s@]@.@." (Printing.line ~c:'=' 70);
  Format.fprintf ppf "@[%a@]@." pp_iterations t;
  Format.fprintf ppf "@[<v 4>@;%a@]@." pp_summary    t