open Lacaml.D

module Ds = DeterminantSpace
module De = Determinant
module Sp = Spindeterminant

type t = 
  {
    mo_basis  : MOBasis.t ;
    det_space : DeterminantSpace.t ;
    ci        : CI.t ;
    hf12_integrals : HF12.t ;
    eigensystem : (Mat.t * Vec.t) lazy_t;
  }

let ci t        =  t.ci
let mo_basis t  =  t.mo_basis
let det_space t =  t.det_space
let mo_class t  =  Ds.mo_class @@ det_space t
let eigensystem t =  Lazy.force t.eigensystem


(*

let single_matrices hf12_integrals density =
    let nocc = Mat.dim1 density in
    let nvir = Mat.dim2 density in
    let { HF12.
          simulation ; aux_basis ;
          hf12 ; hf12_anti ;
          hf12_single ; hf12_single_anti } = hf12_integrals
    in
    let f12 = MOBasis.f12_ints   aux_basis in
    let eri = MOBasis.two_e_ints aux_basis in

    let d = Mat.as_vec density in
    let v_s = Mat.create nocc nvir in
    let v_o = Mat.create nocc nvir in
    let h_o, h_s, f_o, f_s = 
      Mat.create nocc nvir ,
      Mat.create nocc nvir ,
      Mat.create nocc nvir ,
      Mat.create nocc nvir ,
    in
    for a=1 to nvir do
      for m=1 to occ do
       for u=1 to nocc do
        for t=1 to nocc do
          let hmtau = ERI.get_phys eri m t a u
          and hmtua = ERI.get_phys eri m t u a
	  in
          v_o.{t,u} <- hmtau;
          v_s.{t,u} <- hmtau -. hmtua;
        done
       done;
       h_o.{m,a} <- dot d_o @@ Mat.as_vec v_o;
       h_s.{m,a} <- dot d_s @@ Mat.as_vec v_s
       for u=1 to nocc do
        for t=1 to nocc do
          let fmtau = ERI.get_phys f12 m t a u
          and fmtua = ERI.get_phys f12 m t u a
	  in
          v_o.{t,u} <- 0.375 *. fmtau +. 0.125 *. fmtua;
          v_s.{t,u} <- 0.25 *, (fmtau -. fmtua);
        done
       done;
       f_o.{m,a} <- dot d_o @@ Mat.as_vec v_o;
       f_s.{m,a} <- dot d_s @@ Mat.as_vec v_s
      done
    done;
*)
  
(*---
let hf_ij_non_zero mo_basis hf12_integrals deg_a deg_b ki kj =
  let integrals = [

    let { HF12.
          simulation ; aux_basis ;
          hf12 ; hf12_anti ;
          hf12_single ; hf12_single_anti } = hf12_integrals
    in

    let kia = De.alfa ki and kib = De.beta ki in
    let kja = De.alfa kj and kjb = De.beta kj in

    let mo_a =
      Bitstring.logand (Sp.bitstring kia) (Sp.bitstring kja)
      |> Bitstring.to_list
      |> Array.of_list
    and mo_b =
      Bitstring.logand (Sp.bitstring kib) (Sp.bitstring kjb)
      |> Bitstring.to_list
      |> Array.of_list
    in

    let aux_mos = 
      Util.list_range (MOBasis.size mo_basis) (MOBasis.size aux_basis)
      |> Array.of_list
    in

    let one_e _ _ _ = 0. in

    let h = MOBasis.ee_ints aux_basis in
    let two_e_h i j k l s s' =
      if s' <> s then
        ERI.get_phys h i j k l
      else
        (ERI.get_phys h i j k l) -.  (ERI.get_phys h i j l k)
    in

    let f = MOBasis.f12_ints aux_basis in
    let two_e_f i j k l s s' =
      let fijkl = F12.get_phys f i j k l
      and fijlk = F12.get_phys f i j l k
      in
      if s' <> s then
        0.325 *. fijkl +. 0.125 *. fijlk
      else
        0.25 *. (fijkl -. fijlk)
    in

    (*
    let mo_of_s = function
    | Spin.Alfa -> mo_a
    | Spin.Beta -> mo_b
    in
    *)

    let two_e i j k l s s' =
      (
      if s = s' then
	hf12_anti.{i,j,k,l} -. (
	(Array.fold_left (fun accu m -> accu +. hf12_single_anti.{m,i,j,k,l}) 0. mo_a) +.
	(Array.fold_left (fun accu m -> accu +. hf12_single_anti.{m,j,i,l,k}) 0. mo_b) )
      else
	hf12.{i,j,k,l}  -. (
	(Array.fold_left (fun accu m -> accu +. hf12_single.{m,i,j,k,l}) 0. mo_a) +.
	(Array.fold_left (fun accu m -> accu +. hf12_single.{m,j,i,l,k}) 0. mo_b) )
      )
      (*
      +. Array.fold_left ( fun accu a -> accu +.
         (Array.fold_left ( fun accu m -> accu +. two_e_h m i m a s s) 0. (mo_of_s s) +.
          Array.fold_left ( fun accu m -> accu +. two_e_h m i m a s s) 0. (mo_of_s @@ Spin.other s) )
         *. (two_e_f a j k l s s') ) 0. aux_mos 
         *)
    in

    let three_e i j k l m n s s' s'' =
         Array.fold_left (fun accu a -> accu +. two_e_h i j l a s   s'  *. two_e_f a k m n s'  s'') 0. aux_mos
      -. Array.fold_left (fun accu a -> accu +. two_e_h j i m a s'  s   *. two_e_f a k l n s   s'') 0. aux_mos
      +. Array.fold_left (fun accu a -> accu +. two_e_h j k m a s'  s'' *. two_e_f a i n l s'' s  ) 0. aux_mos
      -. Array.fold_left (fun accu a -> accu +. two_e_h k j n a s'' s'  *. two_e_f a i m l s'  s  ) 0. aux_mos
      +. Array.fold_left (fun accu a -> accu +. two_e_h k i n a s'' s   *. two_e_f a j l m s   s' ) 0. aux_mos
      -. Array.fold_left (fun accu a -> accu +. two_e_h i k l a s   s'' *. two_e_f a j n m s'' s' ) 0. aux_mos
    in
    (one_e, two_e, Some three_e)
  ]
  in
  CIMatrixElement.non_zero integrals deg_a deg_b ki kj
  |> List.hd


let dressing_vector ~frozen_core hf12_integrals f12_amplitudes ci =

  if Parallel.master then
    Printf.printf "Building matrix\n%!";
  let det_space = 
    ci.CI.det_space 
  in
  let mo_basis =
    Ds.mo_basis det_space
  in

  let m_HF = 
    
    let f =
      match Ds.determinants det_space with
      | Ds.Arbitrary _ -> CI.create_matrix_arbitrary
      | Ds.Spin      _ -> CI.create_matrix_spin_computed ~nmax:3
    in
    f (fun deg_a deg_b ki kj ->
         hf_ij_non_zero mo_basis hf12_integrals deg_a deg_b ki kj
      ) det_space 

  in

  Matrix.mm (Lazy.force m_HF) (Matrix.dense_of_mat f12_amplitudes)

--- *)


let hf_ij_non_zero hf12_integrals deg_a deg_b ki kj =
  let integrals = [

    let { HF12.
          simulation ; aux_basis ;
          hf12 ; hf12_anti ;
          hf12_single ; hf12_single_anti } = hf12_integrals
    in

    let kia = De.alfa ki and kib = De.beta ki in
    let kja = De.alfa kj and kjb = De.beta kj in

    let mo_a =
      Bitstring.logand (Sp.bitstring kia) (Sp.bitstring kja)
      |> Bitstring.to_list
    and mo_b =
      Bitstring.logand (Sp.bitstring kib) (Sp.bitstring kjb)
      |> Bitstring.to_list
    in

    let one_e _ _ _ = 0. in

    let two_e i j k l s s' =
      if s = s' then
	hf12_anti.{i,j,k,l} -. (
	(List.fold_left (fun accu m -> accu +. hf12_single_anti.{m,i,j,k,l}) 0. mo_a) +.
	(List.fold_left (fun accu m -> accu +. hf12_single_anti.{m,j,i,l,k}) 0. mo_b) 
        )
      else
	hf12.{i,j,k,l}  -. (
	(List.fold_left (fun accu m -> accu +. hf12_single.{m,i,j,k,l}) 0. mo_a) +.
	(List.fold_left (fun accu m -> accu +. hf12_single.{m,j,i,l,k}) 0. mo_b) 
        )
    in

    let h = MOBasis.ee_ints aux_basis in
    let two_e_h i j k l s s' =
      if s' <> s then
        ERI.get_phys h l k j i
      else
        (ERI.get_phys h l k j i) -.  (ERI.get_phys h k l j i)
    in

    let f = MOBasis.f12_ints aux_basis in
    let two_e_f i j k l s s' =
      if s' <> s then
        F12.get_phys f i j k l
      else
        (F12.get_phys f i j k l) -.  (F12.get_phys f i j l k)
    in

    let mo_of_s = function
    | Spin.Alfa -> mo_a
    | Spin.Beta -> mo_b
    in

    let three_e i j k l m n s s' s'' =
         List.fold_left (fun accu a -> accu +. two_e_h i j l a s   s'  *. two_e_f a k m n s'  s'') 0. (mo_of_s s' )
      -. List.fold_left (fun accu a -> accu +. two_e_h j i m a s'  s   *. two_e_f a k l n s   s'') 0. (mo_of_s s  )
      +. List.fold_left (fun accu a -> accu +. two_e_h j k m a s'  s'' *. two_e_f a i n l s'' s  ) 0. (mo_of_s s'')
      -. List.fold_left (fun accu a -> accu +. two_e_h k j n a s'' s'  *. two_e_f a i m l s'  s  ) 0. (mo_of_s s' )
      +. List.fold_left (fun accu a -> accu +. two_e_h k i n a s'' s   *. two_e_f a j l m s   s' ) 0. (mo_of_s s  )
      -. List.fold_left (fun accu a -> accu +. two_e_h i k l a s   s'' *. two_e_f a j n m s'' s' ) 0. (mo_of_s s'')
    in

    (one_e, two_e, Some three_e)
  ]
  in
  CIMatrixElement.non_zero integrals deg_a deg_b ki kj
  |> List.hd


let dressing_vector ~frozen_core hf12_integrals f12_amplitudes ci =

  if Parallel.master then
    Printf.printf "Building matrix\n%!";
  let det_space =
    ci.CI.det_space
  in

  let m_HF =

    let f =
      match Ds.determinants det_space with
      | Ds.Arbitrary _ -> CI.create_matrix_arbitrary
      | Ds.Spin      _ -> CI.create_matrix_spin_computed ~nmax:3
    in
    f (fun deg_a deg_b ki kj ->
         hf_ij_non_zero hf12_integrals deg_a deg_b ki kj
      ) det_space

  in

  Matrix.mm (Lazy.force m_HF) (Matrix.dense_of_mat f12_amplitudes)


let make ~simulation ?(threshold=1.e-12) ~frozen_core ~mo_basis ~aux_basis_filename ?(state=1) () =

  let det_space = 
    DeterminantSpace.fci_of_mo_basis  mo_basis  ~frozen_core 
  in

  let ci = CI.make ~n_states:state det_space in

  let hf12_integrals =
    HF12.make ~simulation ~mo_basis ~aux_basis_filename ()
  in

  let ci_coef, ci_energy =
    let x = Lazy.force ci.eigensystem in
    Parallel.broadcast (lazy x)
  in

  let eigensystem = lazy (
    let m_H =
      Lazy.force ci.CI.m_H
    in

      
    let rec iteration ~state psi = 
(*
Format.printf "%a@." DeterminantSpace.pp_det_space @@ CI.det_space ci;
Format.printf "%a@." Matrix.pp_matrix @@ Matrix.dense_of_mat psi;
*)
      let column_idx = iamax (Mat.to_col_vecs psi).(state-1) in

      let delta = 
         (* delta_i = {% $\sum_j c_j H_{ij}$ %} *)
         dressing_vector ~frozen_core hf12_integrals psi ci
         |> Matrix.to_mat
      in
      (*
      Format.printf "%a@." Matrix.pp_matrix @@ Matrix.dense_of_mat delta;
      *)


      Printf.printf "Cmax : %e\n" psi.{column_idx,state};
      Printf.printf "Norm : %e\n" (sqrt (gemm ~transa:`T delta delta).{state,state});

      let f = 1.0 /. psi.{column_idx,state} in
      let delta_00 =
        (* Delta_00 = {% $\sum_{j \ne x} delta_j c_j / c_x$ %} *)
        f *. ( (gemm ~transa:`T delta psi).{state,state} -.
          delta.{column_idx,state} *. psi.{column_idx,state} )
      in
      Printf.printf "Delta_00 : %e %e\n" delta.{column_idx,state} delta_00;

      delta.{column_idx,state} <- delta.{column_idx,state} -. delta_00;


      let eigenvectors, eigenvalues = 

        let delta = lacpy delta in
        Mat.scal f delta;
        for k=1 to state-1 do
          for i=1 to Mat.dim1 delta do
                delta.{i,k} <- delta.{i,state}
          done;
        done;
        let diagonal =
          Vec.init (Matrix.dim1 m_H) (fun i ->
            if i = column_idx then
              Matrix.get m_H i i +. delta.{column_idx,state}
            else
              Matrix.get m_H i i
            )
        in

        let matrix_prod c = 
          let w = 
            Matrix.mm ~transa:`T m_H c
            |> Matrix.to_mat
          in
          let c = Matrix.to_mat c in

          for k=1 to state do
            for i=1 to (Mat.dim1 w) do
              w.{i,k} <- w.{i,k} +.  delta.{i,k} *. c.{column_idx, k} ;
              w.{column_idx,k} <- w.{column_idx,k} +.  delta.{i,k} *. c.{i,k};
            done;
            w.{column_idx,k} <- w.{column_idx,k} -.
                delta.{column_idx,k} *. c.{column_idx,k};
          done;

          Matrix.dense_of_mat w
        in


        Parallel.broadcast (lazy (
          Davidson.make ~threshold:1.e-10 ~guess:psi ~n_states:state diagonal matrix_prod
          ))

      in
      let eigenvectors = 
        Conventions.rephase @@ Util.remove_epsilons eigenvectors
      in


      Vec.iter (fun energy -> Printf.printf "%g\t" energy) eigenvalues;
      print_newline ();

      let conv =
        1.0 -. abs_float ( dot
          (Mat.to_col_vecs psi).(0)
          (Mat.to_col_vecs eigenvectors).(0) )
      in
      if Parallel.master then
        Printf.printf "F12 Convergence : %e  %f\n" conv (eigenvalues.{state} 
          +. Simulation.nuclear_repulsion simulation);

(*
      let cabs_singles = 
        let f =
          Fock.make_rhf ~density ~ao_basis:large_ao_basis
        in
      in
*)

      if conv > threshold then
        iteration ~state eigenvectors
      else
        let eigenvalues =
          Vec.map (fun x -> x +. ci.CI.e_shift) eigenvalues
        in
        eigenvectors, eigenvalues
    in
    iteration ~state ci_coef

  )
  in
  { mo_basis ; det_space ; ci ; hf12_integrals ; eigensystem }