open Lacaml.D
open Util

type s =
  {
    simulation        : Simulation.t;
    guess             : Guess.t;
    eigenvectors      : Mat.t ;
    eigenvalues       : Vec.t ;
    occupation        : Vec.t ;
    iterations        : (float * float * float) array;
    energy            : float ;
    nuclear_repulsion : float ;
    kin_energy        : float ;
    eN_energy         : float ;
    coulomb_energy    : float ;
    exchange_energy   : float ;
    nocc              : int;
  }


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


let iterations_to_string hf_calc =
  " #       HF energy     Convergence    HOMO-LUMO   
   ---------------------------------------------------" ::
  Array.to_list (Array.mapi (fun i (e, c, g) ->
  Printf.sprintf "    %5d   %13.8f   %11.4e   %11.4f  " (i+1) e c g) hf_calc.iterations)
  @ [ "  -----------------------------------------------------" ]
  |> String.concat "\n"


let summary hf_calc = 
  let ha_to_ev = Constants.ha_to_ev in
  String.concat "\n" [
    "  =====================================================";
    Printf.sprintf "    One-electron energy  %16.10f" (hf_calc.kin_energy +. hf_calc.eN_energy) ;
    Printf.sprintf "    Kinetic      energy  %16.10f"  hf_calc.kin_energy ;
    Printf.sprintf "    Potential    energy  %16.10f"  hf_calc.eN_energy ;
    "   ---------------------------------------------------";
    Printf.sprintf "    Two-electron energy  %16.10f" (hf_calc.coulomb_energy +. hf_calc.exchange_energy) ;
    Printf.sprintf "    Coulomb      energy  %16.10f"  hf_calc.coulomb_energy ;
    Printf.sprintf "    Exchange     energy  %16.10f"  hf_calc.exchange_energy ;
    "   ---------------------------------------------------";
    Printf.sprintf "    Electronic   energy  %16.10f" (hf_calc.energy -. hf_calc.nuclear_repulsion);
    Printf.sprintf "    Nuclear   repulsion  %16.10f"  hf_calc.coulomb_energy ;
    Printf.sprintf "    Hartree-Fock energy  %16.10f"  hf_calc.energy ;
    "   ---------------------------------------------------";
    Printf.sprintf "    HF HOMO      energy  %16.10f eV" (ha_to_ev *. hf_calc.eigenvalues.{hf_calc.nocc} );
    Printf.sprintf "    HF LUMO      energy  %16.10f eV" (ha_to_ev *. hf_calc.eigenvalues.{hf_calc.nocc+1} );
    Printf.sprintf "    HF HOMO-LUMO    gap  %16.10f eV"
          (ha_to_ev *.(hf_calc.eigenvalues.{hf_calc.nocc+1} -. hf_calc.eigenvalues.{hf_calc.nocc}));
    "  =====================================================" ]


let mos_to_string hf_calc =
  let open Lacaml.Io in
  let rows = Mat.dim1 hf_calc.eigenvectors
  and cols = Mat.dim2 hf_calc.eigenvectors
  in
  let rec aux accu first last = 
    
    if (first > last) then String.concat "\n" (List.rev accu)
    else
      let nw = 
    "\n    Eigenvalues : " ^ (
    Array.mapi (fun i x ->
      if (i+1 >= first) && (i+1 <= first+4 ) then
        Some (Printf.sprintf " %f " x)
        else None
      )
      (Vec.to_array hf_calc.eigenvalues)
    |> Array.to_list
    |> list_some
    |> String.concat " "
    )^ 
        Format.asprintf "\n\n    %a\n" (Lacaml.Io.pp_lfmat
          ~row_labels:
            (Array.init rows (fun i -> Printf.sprintf "%d  " (i + 1)))
          ~col_labels:
            (Array.init (min 5 (cols-first+1)) (fun i -> Printf.sprintf "-- %d --" (i + first) ))
          ~print_right:false 
          ~print_foot:false
          () ) (lacpy ~ac:first ~n:(min 5 (cols-first+1)) hf_calc.eigenvectors) 
      in
      aux (nw :: accu) (first+5) last
  in
  String.concat "\n" [ "
  =========================================================================
                            Molecular Orbitals
  =========================================================================
  
    Occupied
   -----------------------------------------------------------------------
  " ;
  aux [] 1 hf_calc.nocc ;
  "
    Virtual 
   -----------------------------------------------------------------------
  " ;
  aux [] (hf_calc.nocc+1) (Mat.dim2 hf_calc.eigenvectors) ;
  "
  =========================================================================
  " ]


let to_string hf =
  let aux hf_calc r =
    String.concat "\n" [ Printf.sprintf "
    =====================================================
                  %s Hartree-Fock                     
    =====================================================" r ; "" ;
    iterations_to_string hf_calc ; "" ;
    summary hf_calc ; "" ;
    mos_to_string hf_calc ; "" ;
    ]
  in
  match hf with
  | RHF  hf_calc -> aux hf_calc "Restricted"
  | UHF  hf_calc -> aux hf_calc "Unrestricted"
  | ROHF hf_calc -> aux hf_calc "Restricted Open-shell"