open Lacaml.D

type t = 
  {
    mo_basis  : MOBasis.t ;
    aux_basis : MOBasis.t ;
    det_space : DeterminantSpace.t ;
    ci        : CI.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  =  DeterminantSpace.mo_class @@ det_space t
let eigensystem t =  Lazy.force t.eigensystem


let f12_integrals mo_basis =
  let two_e_ints  = MOBasis.f12_ints mo_basis in
  ( (fun _ _ _ -> 0.),
    (fun i j k l s s' ->
       if (i=k && j<>l) || (j=l && i<>k) then
          0.
       else
        begin
          let ijkl = F12.get_phys two_e_ints i j k l
          in
          (*
          if s' = Spin.other s then
            (* Minus sign because we swap spin variables
               instead of orbital variables *)
            0.375 *. ijkl +. 0.125 *. ijlk 
          else
            0.25 *. (ijkl -. ijlk) 
          *)
          if s' = Spin.other s then
            ijkl 
          else
            let ijlk = F12.get_phys two_e_ints i j l k
            in
            ijkl -. ijlk
        end
    ) )


let h_ij mo_basis ki kj =
  let integrals =
    List.map (fun f -> f mo_basis)
      [ CI.h_integrals ]
  in
  CIMatrixElement.make integrals ki kj 
  |> List.hd


let f_ij mo_basis ki kj =
  let integrals =
    List.map (fun f -> f mo_basis)
      [ f12_integrals ]
  in
  CIMatrixElement.make integrals ki kj 
  |> List.hd



let hf_ij mo_basis ki kj =
  let integrals =
    List.map (fun f -> f mo_basis)
      [ CI.h_integrals ; f12_integrals ]
  in
  CIMatrixElement.make integrals ki kj 


let is_a_double det_space =
    let mo_class      = DeterminantSpace.mo_class det_space in
    let mo_num = Array.length @@ MOClass.mo_class_array mo_class in
    let m l =
      List.fold_left (fun accu i ->
          let j = i-1 in Bitstring.logor accu (Bitstring.shift_left_one mo_num j)
        ) (Bitstring.zero mo_num) l
    in
    let aux_mask   = m (MOClass.auxiliary_mos mo_class) in
    fun k ->
      let alfa =
        Determinant.alfa k
        |> Spindeterminant.bitstring
      in
      let beta =
        Determinant.beta k
        |> Spindeterminant.bitstring
      in
      let a = Bitstring.logand aux_mask alfa
      and b = Bitstring.logand aux_mask beta
      in
      match Bitstring.popcount a + Bitstring.popcount b with
      | 2 -> true
      | _ -> false


let p12 det_space =
    let mo_class      = DeterminantSpace.mo_class det_space in
    let mo_num = Array.length @@ MOClass.mo_class_array mo_class in
    let m l =
      List.fold_left (fun accu i ->
          let j = i-1 in Bitstring.logor accu (Bitstring.shift_left_one mo_num j)
        ) (Bitstring.zero mo_num) l
    in
    let aux_mask = m (MOClass.auxiliary_mos mo_class) in
    let not_aux_mask =
      Bitstring.(shift_left_one mo_num mo_num |> minus_one)
    in
    fun k ->
      let alfa =
        Determinant.alfa k
        |> Spindeterminant.bitstring
      in
      let beta =
        Determinant.beta k
        |> Spindeterminant.bitstring
      in
      let a = Bitstring.logand aux_mask alfa
      and b = Bitstring.logand aux_mask beta
      in
      match Bitstring.popcount a, Bitstring.popcount b with
      | 2, 0 
      | 0, 2 -> Some (Determinant.negate_phase k)
      | 1, 1 -> Some (Determinant.of_spindeterminants 
                  (Spindeterminant.of_bitstring  @@
                    Bitstring.(logor b (logand not_aux_mask alfa)) )
                  (Spindeterminant.of_bitstring  @@
                    Bitstring.(logor a (logand not_aux_mask beta))
                  ) )
      | _ -> None



let dressing_vector ~frozen_core aux_basis f12_amplitudes ci =

  if Parallel.master then
    Printf.printf "Building matrix\n%!";

  (* Determinants of the FCI space as a list *)
  let in_dets =
    DeterminantSpace.determinant_stream ci.CI.det_space 
    |> Util.stream_to_list
  in

  (* Stream that generates only singly and doubly excited determinants
      wrt FCI space *)
  let out_dets_stream =
      (* Stream that generates all determinants of FCI space *)
      let s = 
        DeterminantSpace.fci_of_mo_basis ~frozen_core aux_basis
        |> DeterminantSpace.determinant_stream

      in
      (* Select only doubly excited determinants wrt FCI space *)
      Stream.from (fun _ ->
        try
          let p12 = p12 ci.CI.det_space in
          let rec result () = 
            let ki = Stream.next s in
            match p12 ki with
            | Some ki' -> Some (ki, ki')
            | None -> result ()
          in
          result ()
        with Stream.Failure -> None
      )
  in

  let make_h_and_f alpha_list = 

    let rec col_vecs_list accu_H accu_F = function
    | [] -> 
        List.rev accu_H, 
        List.rev accu_F 
    | (ki, ki') :: rest ->
        begin
        let h, f = 
          List.map (fun kj ->
          match hf_ij aux_basis kj ki with
          | [ a ; b ] -> a, b
          | _ -> assert false ) in_dets
          |> List.split
        in
        let f' = 
          List.map (fun kj -> f_ij aux_basis kj ki') in_dets
        in
        let h  = Vec.of_list h in
        let f  = Vec.of_list f in
        let f' = Vec.of_list f' in
        scal 0.375 f;
        scal 0.125 f';
        let f  = Vec.add f f' in
        col_vecs_list (h::accu_H) (f::accu_F) rest
        end
    in
    let h, f = 
      col_vecs_list [] [] alpha_list
    in
    Mat.of_col_vecs_list h,
    Mat.of_col_vecs_list f
  in


  let m_HF = 
    let batch_size =  1 + 1_000_000 / (Mat.dim1 f12_amplitudes) in 
    let input_stream = 
      Stream.from (fun i ->
        let rec make_batch accu = function
          | 0 -> accu
          | n -> try
                  let alpha = Stream.next out_dets_stream in
                  let accu  = alpha :: accu in
                  make_batch accu (n-1)
                 with Stream.Failure -> accu
        in
        let result = make_batch [] batch_size in
        if result = [] then None else Some result
        )
    in

    let iteration input =
      Printf.printf ".%!";
      let m_H_aux, m_F_aux = make_h_and_f input in
      let m_HF =
        gemm m_H_aux m_F_aux ~transb:`T 
      in
      gemm m_HF f12_amplitudes
    in

    let result = 
      let x = 
        try [ Stream.next out_dets_stream ]
        with Stream.Failure -> failwith "Auxiliary basis set does not produce any excited determinant"
      in
      iteration x
    in
 
    input_stream
    |> Farm.run ~ordered:false ~f:iteration
    |> Stream.iter (fun hf ->
      ignore @@ Mat.add result hf ~c:result );
    Printf.printf "\n";
    Parallel.broadcast (lazy result)
  in

  if Parallel.master then Printf.printf "Done\n%!";
  Matrix.dense_of_mat m_HF





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

  let f12 = Util.of_some @@ Simulation.f12 simulation in
  let mo_num = MOBasis.size mo_basis in

Printf.printf "Add aux basis\n%!";
  (* Add auxiliary basis set *)
  let s =
    let charge        = Charge.to_int @@ Simulation.charge simulation 
    and multiplicity  = Electrons.multiplicity @@ Simulation.electrons simulation
    and nuclei        = Simulation.nuclei simulation
    in
    let general_basis = 
      Basis.general_basis @@ Simulation.basis simulation
    in
    GeneralBasis.combine [
      general_basis ; GeneralBasis.read aux_basis_filename
    ]
    |> Basis.of_nuclei_and_general_basis nuclei
    |> Simulation.make ~f12 ~charge ~multiplicity ~nuclei 
  in

  let aux_basis = 
    MOBasis.of_mo_basis s mo_basis
  in
  let () = 
    ignore @@ MOBasis.f12_ints aux_basis
  in
  let () =
    ignore @@ MOBasis.two_e_ints aux_basis
  in

  let det_space = 
    DeterminantSpace.fci_f12_of_mo_basis  aux_basis  ~frozen_core  mo_num
  in

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

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

  let e_shift = 
    let det =
      DeterminantSpace.determinant_stream det_space 
      |> Stream.next
    in
    h_ij aux_basis det det
  in

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

      
    let rec iteration ~state 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 aux_basis psi ci
         |> Matrix.to_mat
      in

      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 diagonal =
        Vec.init (Matrix.dim1 m_H) (fun i ->
          if i = column_idx then
            Matrix.get m_H i i +. delta.{column_idx,state} *. f 
          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 c11 = Matrix.get c column_idx state in
        Util.list_range 1 (Mat.dim1 w)
        |> List.iter (fun i ->
          let dci = 
            delta.{i,state} *. f ;
          in
          w.{i,state} <- w.{i,state} +. dci *.  c11;
          if (i <> column_idx) then
            w.{column_idx,state} <- w.{column_idx,state} +. dci *. (Matrix.get c i state);
          );
        Matrix.dense_of_mat w
      in

      let eigenvectors, eigenvalues = 
        Parallel.broadcast (lazy (
          Davidson.make ~threshold:1.e-6 ~guess:psi ~n_states:state diagonal matrix_prod
          ))
      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} +. e_shift 
          +. Simulation.nuclear_repulsion simulation);

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

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