open Lacaml.D
open Simulation
open Constants
open Util


type t =
  {
    fock     : Mat.t ;
    core     : Mat.t ;
    coulomb  : Mat.t ;
    exchange : Mat.t ;
  }

module Ao = AOBasis



let make ~density ao_basis =
  let m_P  = density
  and m_T  = Lazy.force ao_basis.Ao.kin_ints |> KinInt.matrix 
  and m_V  = Lazy.force ao_basis.Ao.eN_ints  |> NucInt.matrix
  and m_G  = Lazy.force ao_basis.Ao.ee_ints  
  in
  let nBas = Mat.dim1 m_T
  in

  let m_Hc = Mat.add m_T m_V 
  and m_J = Array.make_matrix nBas nBas 0.
  and m_K = Array.make_matrix nBas nBas 0.
  in
  (*
  let permutations  i j k l  =
      [  [ i ; j ; k ; l ] ; 
         [ k ; j ; i ; l ] ; 
         [ i ; l ; k ; j ] ; 
         [ k ; l ; i ; j ] ; 
         [ j ; i ; l ; k ] ; 
         [ j ; k ; l ; i ] ; 
         [ l ; i ; j ; k ] ; 
         [ l ; k ; j ; i ]
         ] 
  in
  ERI.to_stream m_G
  |> Stream.iter (fun { ERI.i_r1 ; j_r2 ; k_r1 ; l_r2 ; value } ->
     permutations i_r1 j_r2 k_r1 l_r2 
     |> List.iter ( fun ijkl ->
        match ijkl with
        | mu :: lambda :: nu :: sigma :: [] ->
          let p = m_P.{lambda,sigma} in
          if abs_float p > epsilon then
            let m_Jnu = m_J.(nu-1) in
            m_Jnu.(mu-1) <- m_Jnu.(mu-1) +. p *. value 
        | _ -> assert false
    ));
  *)
  for sigma = 1 to nBas do
    for nu = 1 to nBas do
      let m_Jnu = m_J.(nu-1) in
      for lambda = 1 to sigma do
        let p = 
          if lambda < sigma then
            2. *. m_P.{lambda,sigma}
          else
            m_P.{lambda,sigma}
        in
        if abs_float p > epsilon then
          for mu = 1 to nu do
            m_Jnu.(mu-1) <- m_Jnu.(mu-1) +.  p *. 
                          ERI.get_phys m_G mu lambda nu sigma
          done
      done
    done
  done;
  for nu = 1 to nBas do
    for mu = 1 to nu-1 do
      m_J.(mu-1).(nu-1) <- m_J.(nu-1).(mu-1);
    done
  done;

  for nu = 1 to nBas do
    let m_Knu = m_K.(nu-1) in
    for sigma = 1 to nBas do
      for lambda = 1 to nBas do
        let p = 
            0.5 *. m_P.{lambda,sigma}
        in
        if abs_float p > epsilon then
          for mu = 1 to nu do
            m_Knu.(mu-1) <- m_Knu.(mu-1) -. p *.
                          ERI.get_phys m_G mu lambda sigma nu
          done
      done
    done;
    for mu = 1 to nu-1 do
      m_K.(mu-1).(nu-1) <- m_Knu.(mu-1);
    done
  done;
  let m_J =  Mat.of_array m_J
  and m_K =  Mat.of_array m_K
  in
  { fock = Mat.add m_Hc (Mat.add m_J m_K) ;
    core = m_Hc ; coulomb = m_J ; exchange = m_K }


let pp_fock ppf a =
  Format.fprintf ppf "@[<2>";
  Format.fprintf ppf "@[ Fock matrix:@[<2>@[%a@]@.]@]" pp_matrix a.fock;
  Format.fprintf ppf "@[ Core Hamiltonian:@[<2>@[%a@]@.]@]" pp_matrix a.core;
  Format.fprintf ppf "@[ Coulomb matrix:@[<2>@[%a@]@.]@]" pp_matrix a.coulomb;
  Format.fprintf ppf "@[ Exchange matrix:@[<2>@[%a@]@.]@]" pp_matrix a.exchange;
  Format.fprintf ppf "@]"