open Lacaml.D
open Util

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

type t = guess

module Si = Simulation
module El = Electrons

let hcore_guess simulation = 
  let eN_ints  = Lazy.force simulation.Si.eN_ints
  and kin_ints = Lazy.force simulation.Si.kin_ints
  in
  Mat.add eN_ints kin_ints


let huckel_guess simulation =
  let c = 0.5 *. 1.75 in
  let ao_num = Basis.size simulation.Si.basis in
  let eN_ints  = Lazy.force simulation.Si.eN_ints
  and kin_ints = Lazy.force simulation.Si.kin_ints
  and overlap  = Lazy.force simulation.Si.overlap
  and m_X = Lazy.force simulation.Si.overlap_ortho
  in
  let diag = Array.init (ao_num+1) (fun i -> if i=0 then 0. else
    eN_ints.{i,i} +. kin_ints.{i,i})
  in

  let nocc = 
    simulation.Si.electrons.El.n_alpha
  in
  let m_F = (Fock.make ~density:(gemm ~alpha:2. ~transb:`T ~k:nocc m_X m_X) simulation).Fock.fock in
  for j=1 to ao_num do
    for i=1 to ao_num do
      if (i <> j) then
        m_F.{i,j} <- c *. overlap.{i,j} *. (diag.(i) +. diag.(j)) (*TODO Pseudo *)
    done;
  done;
  m_F

let make ~guess simulation = 
  match guess with
  | `Hcore  -> Hcore  (hcore_guess  simulation)
  | `Huckel -> Huckel (huckel_guess simulation)