open Lacaml.D
open Util

type guess = 
| Hcore  of Mat.t
| Huckel of Mat.t
| Matrix of Mat.t

type t = guess

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

let hcore_guess ao_basis = 
  let eN_ints  = Ao.eN_ints  ao_basis |> NucInt.matrix
  and kin_ints = Ao.kin_ints ao_basis |> KinInt.matrix
  in
  Mat.add eN_ints kin_ints


let huckel_guess ao_basis =
  let c = 0.5 *. 1.75 in
  let eN_ints  = Ao.eN_ints  ao_basis |> NucInt.matrix
  and kin_ints = Ao.kin_ints ao_basis |> KinInt.matrix
  in
  let m_F = 
    Mat.add eN_ints kin_ints
  in
  let ao_num   = Ao.basis    ao_basis |> Basis.size 
  and overlap  = Ao.overlap  ao_basis |> Ov.matrix
  in
  let diag = Vec.init ao_num (fun i -> 
    m_F.{i,i} )
  in

  function 
  | 0 -> invalid_arg "Huckel guess needs a non-zero number of occupied MOs."
  | nocc -> 
    Mat.init_cols ao_num ao_num (fun i j ->
        if (i<>j) then
          if (diag.{i} +. diag.{j}) < 0. then
            c *. overlap.{i,j} *. (diag.{i} +. diag.{j}) +. m_F.{i,j} (*TODO Pseudo *)
          else
            m_F.{i,j} (*TODO Pseudo *)
        else
          diag.{i}
    )


let make ?(nocc=0) ~guess ao_basis = 
  match guess with
  | `Hcore  -> Hcore  (hcore_guess  ao_basis)
  | `Huckel -> Huckel (huckel_guess ao_basis nocc)
  | `Matrix m -> Matrix m



let test_case ao_basis =

  let test_hcore () =
    match make ~guess:`Hcore ao_basis with
    | Hcore matrix -> 
      let a = Lacaml.D.Mat.to_array matrix in
      let reference =
        Lacaml.D.Mat.add 
          (AOBasis.eN_ints  ao_basis |> NucInt.matrix)
          (AOBasis.kin_ints ao_basis |> KinInt.matrix)
        |> Lacaml.D.Mat.to_array
      in
      Array.iteri (fun i x ->
        let message =
          Printf.sprintf "Guess line %d" (i)
        in
        Alcotest.(check (array (float 1.e-15))) message a.(i) x) reference
    | _ -> assert false
  
  in
  [
    "HCore", `Quick, test_hcore;
  ]