open Lacaml.D
open Util

module Ds = DeterminantSpace
module Sd = Spindeterminant

type t =
  {
    e_shift     : float ; (* Diagonal energy shift for increasing numerical precision *)
    det_space   : Ds.t ;
    m_H         : Matrix.t lazy_t ;
    m_S2        : Matrix.t lazy_t ;
    eigensystem : (Mat.t * Vec.t) lazy_t;
    n_states    : int;
  }

let det_space    t = t.det_space

let n_states     t = t.n_states

let m_H          t = Lazy.force t.m_H

let m_S2         t = Lazy.force t.m_S2

let eigensystem  t = Lazy.force t.eigensystem

let mo_class     t = DeterminantSpace.mo_class t.det_space

let eigenvectors t =
  let (x,_) = eigensystem t in x

let eigenvalues t =
  let (_,x) = eigensystem t in x

let h_integrals mo_basis =
  let one_e_ints    = MOBasis.one_e_ints mo_basis
  and two_e_ints    = MOBasis.two_e_ints mo_basis
  in
  ( (fun i j _ -> one_e_ints.{i,j}),
    (fun i j k l s s' ->
       if s' <> s then
         ERI.get_phys two_e_ints i j k l
       else
         (ERI.get_phys two_e_ints i j k l) -.
         (ERI.get_phys two_e_ints i j l k)
    ), None )



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


let h_ij_non_zero mo_basis deg_a deg_b ki kj =
  let integrals =
    List.map (fun f -> f mo_basis)
      [ h_integrals ]
  in
  CIMatrixElement.non_zero integrals deg_a deg_b ki kj
  |> List.hd


let create_matrix_arbitrary f det_space =
  lazy (
    let ndet = Ds.size det_space in
    let data =
      match Ds.determinants det_space with
      | Ds.Arbitrary a -> a
      | _ -> assert false
    in
    let det_alfa    = data.Ds.det_alfa
    and det_beta    = data.Ds.det_beta
    and det         = data.Ds.det
    and index_start = data.Ds.index_start
    in


    (** Array of (list of singles, list of doubles) in the beta spin *)
    let degree_bb =
      Array.map (fun det_i ->
          let deg = Spindeterminant.degree det_i in
          let doubles =
            Array.mapi (fun i det_j ->
                let d = deg det_j in
                if d < 3 then
                  Some (i,d,det_j)
                else
                  None
              ) det_beta
            |> Array.to_list
            |> Util.list_some
          in
          let singles =
            List.filter (fun (i,d,det_j) -> d < 2) doubles
            |> List.rev_map (fun (i,_,det_j) -> (i,det_j))
            |> List.rev
          in
          let doubles =
            List.rev_map (fun (i,_,det_j) -> (i,det_j)) doubles
            |> List.rev
          in
          (singles, doubles)
        ) det_beta
    in

    let task (i,i_dets) =
      let shift = index_start.(i) in

      let result =
        Array.init (index_start.(i+1) - shift)
          (fun _ -> [])
      in

      (** Update function when ki and kj are connected *)
      let update i j deg_a deg_b ki kj =
        let x = f deg_a deg_b ki kj in
        if abs_float x > Constants.epsilon then
          result.(i - shift) <- (j, x) :: result.(i - shift) ;
      in

      let i_alfa = det_alfa.(i) in
      let deg_a = Spindeterminant.degree i_alfa in

      Array.iteri (fun j j_dets ->
          let j_alfa = det_alfa.(j) in
          let degree_a = deg_a j_alfa in

          begin
            match degree_a with
            | 2 ->
              Array.iteri (fun i' i_b ->
                  try
                    Array.iteri (fun j' j_b ->
                        if j_b >= i_b then
                          ( if j_b = i_b then
                              ( let i_beta = det_beta.(i_b) in
                                let ki = Determinant.of_spindeterminants i_alfa i_beta in
                                let kj = Determinant.of_spindeterminants j_alfa i_beta in
                                update (index_start.(i) + i')
                                  (index_start.(j) + j' + 1) 2 0 ki kj);
                            raise Exit)
                      ) j_dets
                  with Exit -> ()
                ) i_dets
            | 1 ->
              Array.iteri (fun i' i_b ->
                  let i_beta = det_beta.(i_b) in
                  let ki = Determinant.of_spindeterminants i_alfa i_beta in
                  let singles, _ = degree_bb.(i_b) in
                  let rec aux singles j' =
                    match singles with
                    | [] -> ()
                    | (js, j_beta)::r_singles ->
                      begin
                        match compare js j_dets.(j') with
                        | -1 -> (aux [@tailcall]) r_singles j'
                        |  0 ->
                          let kj =
                            Determinant.of_spindeterminants j_alfa j_beta
                          in (update
                                (index_start.(i) + i') (index_start.(j) + j' + 1)
                                1 (Determinant.degree_beta ki kj) ki kj;
                              (aux [@tailcall]) r_singles (j'+1);)
                        |  1 -> if (j' < Array.length j_dets) then (aux [@tailcall]) singles (j'+1)
                        |  _ -> assert false
                      end
                  in aux singles 0
                ) i_dets
            | 0 ->
              Array.iteri (fun i' i_b ->
                  let i_beta = det_beta.(i_b) in
                  let ki = Determinant.of_spindeterminants i_alfa i_beta in
                  let _, doubles = degree_bb.(i_b) in
                  let rec aux doubles j' =
                    match doubles with
                    | [] -> ()
                    | (js, j_beta)::r_doubles ->
                      begin
                        match compare js j_dets.(j') with
                        | -1 -> (aux [@tailcall]) r_doubles j'
                        |  0 ->
                          let kj =
                            Determinant.of_spindeterminants j_alfa j_beta
                          in (update
                                (index_start.(i) + i') (index_start.(j) + j' + 1)
                                0 (Determinant.degree_beta ki kj) ki kj;
                              (aux [@tailcall]) r_doubles (j'+1);)
                        |  1 -> if (j' < Array.length j_dets) then (aux [@tailcall]) doubles (j'+1)
                        |  _ -> assert false
                      end
                  in aux doubles 0
                ) i_dets
            | _ -> ();
          end
        ) det;
      let r =
        Array.map (fun l ->
            List.rev l
            |> Vector.sparse_of_assoc_list ndet
          ) result
      in (i,r)
    in


    let result =
      if Parallel.master then
        Array.init ndet (fun _ -> Vector.sparse_of_assoc_list ndet [])
      else
        Array.init ndet (fun _ -> Vector.sparse_of_assoc_list ndet [])
    in

    let n_det_alfa =
      Array.length det_alfa
    in
    Array.mapi (fun i i_dets -> i, i_dets) det
    |> Array.to_list
    |> Stream.of_list
    |> Farm.run ~ordered:false ~f:task
    |> Stream.iter (fun (k, r) ->
        let shift = index_start.(k) in
        let det_k = det.(k) in
        Array.iteri (fun j r_j -> result.(shift+det_k.(j)) <- r_j) r;
        Printf.eprintf "%8d / %8d\r%!" (k+1) n_det_alfa;
      ) ;
    Matrix.sparse_of_vector_array result
  )


(* Create a matrix using the fact that the determinant space is made of
   the outer product of spindeterminants.  *)
let create_matrix_spin ?(nmax=2) f det_space =
  lazy (
    let ndet = Ds.size det_space in
    let a, b =
      match Ds.determinants det_space with
      | Ds.Spin (a,b) -> (a,b)
      | _ -> assert false
    in
    let n_beta = Array.length b in


    (** Array of (list of singles, list of doubles, list of triples) in the beta spin *)
    let degree_bb =
      match nmax with
      | 2 -> begin
                Array.map (fun det_i ->
                    let deg = Spindeterminant.degree det_i in
                    let triples = [] in
                    let doubles =
                      Array.mapi (fun i det_j ->
                          let d = deg det_j in
                          if d < 3 then
                            Some (i,d,det_j)
                          else
                            None
                        ) b
                      |> Array.to_list
                      |> Util.list_some
                    in
                    let singles =
                      List.filter (fun (i,d,det_j) -> d < 2) doubles
                      |> List.rev_map (fun (i,_,det_j) -> (i,det_j))
                      |> List.rev
                    in
                    let doubles =
                      List.rev_map (fun (i,_,det_j) -> (i,det_j)) doubles
                      |> List.rev
                    in
                    (singles, doubles, triples)
                  ) b
             end
      | 3 -> begin
                Array.map (fun det_i ->
                    let deg = Spindeterminant.degree det_i in
                    let triples =
                      Array.mapi (fun i det_j ->
                          let d = deg det_j in
                          if d < 4 then
                            Some (i,d,det_j)
                          else
                            None
                        ) b
                      |> Array.to_list
                      |> Util.list_some
                    in
                    let doubles =
                      List.filter (fun (i,d,det_j) -> d < 3) triples
                    in
                    let singles =
                      List.filter (fun (i,d,det_j) -> d < 2) doubles
                    in

                    let triples =
                      List.rev_map (fun (i,_,det_j) -> (i,det_j)) triples
                      |> List.rev
                    in
                    let doubles =
                      List.rev_map (fun (i,_,det_j) -> (i,det_j)) doubles
                      |> List.rev
                    in
                    let singles =
                      List.rev_map (fun (i,_,det_j) -> (i,det_j)) singles
                      |> List.rev
                    in
                    (singles, doubles, triples)
                  ) b
             end
      | _ -> assert false
    in
    let a = Array.to_list a
    and b = Array.to_list b
    in

    let task =
      let result =
        Array.init n_beta (fun _ -> [])
      in

      (** Update function when ki and kj are connected *)
      let update i j deg_a deg_b ki kj =
        let x = f deg_a deg_b ki kj in
        if abs_float x > Constants.epsilon then
          result.(i) <- (j, x) :: result.(i) ;
      in

      fun (i,i_alfa) ->
        begin match nmax with
        | 2 -> begin
                  let j = ref 1 in

                  let deg_a = Spindeterminant.degree i_alfa in
                  List.iter (fun j_alfa ->
                      let degree_a = deg_a j_alfa in
                      begin
                        match degree_a with
                        | 2 ->
                          let i' = ref 0 in
                          List.iteri (fun ib i_beta ->
                              let ki = Determinant.of_spindeterminants i_alfa i_beta in
                              let kj = Determinant.of_spindeterminants j_alfa i_beta in
                              update !i' (ib + !j) 2 0 ki kj;
                              incr i';
                            ) b;
                        | 1 ->
                          let i' = ref 0 in
                          List.iteri (fun ib i_beta ->
                              let ki = Determinant.of_spindeterminants i_alfa i_beta in
                              let singles, _doubles, _triples = degree_bb.(ib) in
                              List.iter (fun (j', j_beta) ->
                                  let kj = Determinant.of_spindeterminants j_alfa j_beta in
                                  update !i' (j' + !j) 1 (Determinant.degree_beta ki kj) ki kj
                                ) singles;
                              incr i';
                            ) b;
                        | 0 ->
                          let i' = ref 0 in
                          List.iteri (fun ib i_beta ->
                              let ki = Determinant.of_spindeterminants i_alfa i_beta in
                              let _singles, doubles, _triples = degree_bb.(ib) in
                              List.iter (fun (j', j_beta) ->
                                  let kj = Determinant.of_spindeterminants j_alfa j_beta in
                                  update !i' (j' + !j) 0 (Determinant.degree_beta ki kj) ki kj
                                ) doubles;
                              incr i';
                            ) b;
                        | _ -> ();
                      end;
                      j := !j + n_beta
                    ) a
                  end
        | 3 -> begin
                  let j = ref 1 in

                  let deg_a = Spindeterminant.degree i_alfa in
                  List.iter (fun j_alfa ->
                      let degree_a = deg_a j_alfa in
                      begin
                        match degree_a with
                        | 3 ->
                          let i' = ref 0 in
                          List.iteri (fun ib i_beta ->
                              let ki = Determinant.of_spindeterminants i_alfa i_beta in
                              let kj = Determinant.of_spindeterminants j_alfa i_beta in
                              update !i' (ib + !j) 3 0 ki kj;
                              incr i';
                            ) b;
                        | 2 ->
                          let i' = ref 0 in
                          List.iteri (fun ib i_beta ->
                              let ki = Determinant.of_spindeterminants i_alfa i_beta in
                              let singles, _doubles, _triples = degree_bb.(ib) in
                              List.iter (fun (j', j_beta) ->
                                  let kj = Determinant.of_spindeterminants j_alfa j_beta in
                                  update !i' (j' + !j) 2 (Determinant.degree_beta ki kj) ki kj
                                ) singles;
                              incr i';
                            ) b;
                        | 1 ->
                          let i' = ref 0 in
                          List.iteri (fun ib i_beta ->
                              let ki = Determinant.of_spindeterminants i_alfa i_beta in
                              let _singles, doubles, _triples = degree_bb.(ib) in
                              List.iter (fun (j', j_beta) ->
                                  let kj = Determinant.of_spindeterminants j_alfa j_beta in
                                  update !i' (j' + !j) 1 (Determinant.degree_beta ki kj) ki kj
                                ) doubles;
                              incr i';
                            ) b;
                        | 0 ->
                          let i' = ref 0 in
                          List.iteri (fun ib i_beta ->
                              let ki = Determinant.of_spindeterminants i_alfa i_beta in
                              let _singles, _doubles, triples = degree_bb.(ib) in
                              List.iter (fun (j', j_beta) ->
                                  let kj = Determinant.of_spindeterminants j_alfa j_beta in
                                  update !i' (j' + !j) 0 (Determinant.degree_beta ki kj) ki kj
                                ) triples;
                              incr i';
                            ) b;
                        | _ -> ();
                      end;
                      j := !j + n_beta
                    ) a
                  end
          | _ -> assert false
          end;
          let r =
            Array.map (fun l ->
                List.rev l
                |> Vector.sparse_of_assoc_list ndet
              ) result
          in (i,r)
    in


    let result =
      Array.init ndet (fun _ -> Vector.sparse_of_assoc_list ndet [])
    in


    List.mapi (fun i i_alfa -> i*n_beta, i_alfa) a
    |> Stream.of_list
    |> Farm.run ~ordered:false ~f:task
    |> Stream.iter (fun (k, r) ->
        Array.iteri (fun j r_j -> result.(k+j) <- r_j) r;
        Printf.eprintf "%8d / %8d\r%!" (k+Array.length r) ndet;
      ) ;
    Matrix.sparse_of_vector_array result
  )



(* Create a matrix using the fact that the determinant space is made of
   the outer product of spindeterminants.  *)
let create_matrix_spin_computed ?(nmax=2) f det_space =
  lazy (
    let ndet = Ds.size det_space in
    let a, b =
      match Ds.determinants det_space with
      | Ds.Spin (a,b) -> (a,b)
      | _ -> assert false
    in
    let n_beta = Array.length b in

    let h2 i_alfa j_alfa =
        let deg_a = Spindeterminant.degree a.(i_alfa) a.(j_alfa) in
        match deg_a with
        | 2 ->
          let ai, aj = a.(i_alfa), a.(j_alfa) in
          (fun i_beta j_beta ->
              if i_beta <> j_beta then 0. else
                let deg_b = 0 in
                let ki = Determinant.of_spindeterminants ai b.(i_beta) in
                let kj = Determinant.of_spindeterminants aj b.(j_beta) in
                f deg_a deg_b ki kj
          )
        | 1 ->
          let ai, aj = a.(i_alfa), a.(j_alfa) in
          (fun i_beta j_beta ->
              let deg_b = Spindeterminant.degree b.(i_beta) b.(j_beta) in
              match deg_b with
              | 0 | 1 ->
                let ki = Determinant.of_spindeterminants ai b.(i_beta) in
                let kj = Determinant.of_spindeterminants aj b.(j_beta) in
                  f deg_a deg_b ki kj
              | _ -> 0.
          )
        | 0 ->
          let ai, aj = a.(i_alfa), a.(j_alfa) in
          (fun i_beta j_beta ->
              let deg_b = Spindeterminant.degree b.(i_beta) b.(j_beta) in
              match deg_b with
              | 0 | 1 | 2 ->
                let deg_b = Spindeterminant.degree b.(i_beta) b.(j_beta) in
                let ki = Determinant.of_spindeterminants ai b.(i_beta) in
                let kj = Determinant.of_spindeterminants aj b.(j_beta) in
                f deg_a deg_b ki kj
              | _ -> 0.
          )
        | _ -> (fun _ _ -> 0.)
    in

    let h3 i_alfa j_alfa =
        let deg_a = Spindeterminant.degree a.(i_alfa) a.(j_alfa) in
        match deg_a with
        | 3 ->
          let ai, aj = a.(i_alfa), a.(j_alfa) in
          (fun i_beta j_beta ->
              if i_beta <> j_beta then 0. else
                let deg_b = 0 in
                let ki = Determinant.of_spindeterminants ai b.(i_beta) in
                let kj = Determinant.of_spindeterminants aj b.(j_beta) in
                f deg_a deg_b ki kj
          )
        | 2 ->
          let ai, aj = a.(i_alfa), a.(j_alfa) in
          (fun i_beta j_beta ->
              let deg_b = Spindeterminant.degree b.(i_beta) b.(j_beta) in
              match deg_b with
              | 0 | 1 ->
                let ki = Determinant.of_spindeterminants ai b.(i_beta) in
                let kj = Determinant.of_spindeterminants aj b.(j_beta) in
                  f deg_a deg_b ki kj
              | _ -> 0.
          )
        | 1 ->
          let ai, aj = a.(i_alfa), a.(j_alfa) in
          (fun i_beta j_beta ->
              let deg_b = Spindeterminant.degree b.(i_beta) b.(j_beta) in
              match deg_b with
              | 0 | 1 | 2 ->
                let deg_b = Spindeterminant.degree b.(i_beta) b.(j_beta) in
                let ki = Determinant.of_spindeterminants ai b.(i_beta) in
                let kj = Determinant.of_spindeterminants aj b.(j_beta) in
                f deg_a deg_b ki kj
              | _ -> 0.
          )
        | 0 ->
          let ai, aj = a.(i_alfa), a.(j_alfa) in
          (fun i_beta j_beta ->
              let deg_b = Spindeterminant.degree b.(i_beta) b.(j_beta) in
              match deg_b with
              | 0 | 1 | 2 | 3 ->
                let deg_b = Spindeterminant.degree b.(i_beta) b.(j_beta) in
                let ki = Determinant.of_spindeterminants ai b.(i_beta) in
                let kj = Determinant.of_spindeterminants aj b.(j_beta) in
                f deg_a deg_b ki kj
              | _ -> 0.
          )
        | _ -> (fun _ _ -> 0.)
    in

    let i_prev = ref (-10)
    and result = ref (fun _ -> 0.)
    in
    let h123 = ref (fun _ -> 0.)  in

    let h = if nmax = 2 then h2 else h3 in

    let g i =
      (*
        i      : index of the i-th determinant. 1 <= i <= ndet
        i_prev : index of the i-th determinant in the previous function call.
                  1 <= i_prev <= ndet
        i_a    : index of the i_a-th alpha determinant. 0 <= i_a < n_alfa
        i_b    : index of the i_b-th beta determinant. 0 <= i_b < n_beta
        j0     : index - 1 of the first determinant with the same alfa component
                  0 <= j0 < n_beta*(n_alfa-1)
        j1     : index - 1 of the next determinant with the 1st beta determinant
                  n_beta <= j1 <= ndet
        j_a    : index of the j_a-th alpha determinant. 0 <= j_a < n_alfa
        j_b    : index of the j_b-th beta determinant. 0 <= j_b < n_beta
      *)
      if i <> !i_prev then
        begin
          i_prev := i;
          let i_a = (i-1)/n_beta in
          let h1 = h i_a in
          let i_b = i - i_a*n_beta - 1  in
          let j0 = ref (2*ndet) in
          let j1 = ref (2*ndet) in
          let j_a = ref 0 in
          result := fun j ->
              if  (!j0 < j) && (j <= !j1) then
                  ()
              else
                begin
                  if (!j1 < j) && (j <= !j1 + n_beta) then
                    begin
                      incr j_a;
                      j0 := !j1;
                    end
                  else
                    begin
                      j_a := (j-1)/n_beta;
                      j0 := !j_a * n_beta;
                    end;
                  j1 := !j0 + n_beta;
                  h123 := h1 !j_a i_b;
                end;
                let j_b = j - !j0 - 1 in
                !h123 j_b
        end;
      !result
    in

    Matrix.of_fun ndet ndet g
  )



let make ?(n_states=1) ?(algo=`Direct) det_space =

  let mo_basis = Ds.mo_basis det_space in

  (* While in a sequential region, initiate the parallel
     4-idx transformation to avoid nested parallel jobs
  *)
  ignore @@ MOBasis.two_e_ints mo_basis;

  let e_shift =
    let d0 =
      Ds.determinant_stream det_space
      |> Stream.next
    in
    h_ij_non_zero mo_basis 0 0 d0 d0
  in

  let m_H =

    let f =
      match Ds.determinants det_space with
      | Ds.Arbitrary _ -> create_matrix_arbitrary
      | Ds.Spin      _ ->
        if algo = `Direct then
          create_matrix_spin_computed ~nmax:2
        else
          create_matrix_spin  ~nmax:2
    in
    f (fun deg_a deg_b ki kj ->
         if deg_a + deg_b > 0 then
            h_ij_non_zero mo_basis deg_a deg_b ki kj
         else
            h_ij_non_zero mo_basis 0 0 ki ki -. e_shift
      ) det_space
  in

  let m_S2 =
    let f =
      match Ds.determinants det_space with
      | Ds.Arbitrary _ -> create_matrix_arbitrary
      | Ds.Spin      _ -> create_matrix_spin  ~nmax:2
    in
    f (fun _deg_a _deg_b ki kj -> CIMatrixElement.make_s2 ki kj) det_space
    (*TODO*)
  in

  let eigensystem = lazy (

    let eigensystem_incore () =
      let m_H =
        Lazy.force m_H
      in
      let diagonal =
        Parallel.broadcast (lazy (
          Vec.init (Matrix.dim1 m_H) (fun i -> Matrix.get m_H i i)
        ))
      in
      let matrix_prod psi =
        let result =
          Matrix.mm ~transa:`T psi m_H
          |> Matrix.transpose
        in
        Parallel.broadcast (lazy result)
      in
      let eigenvectors, eigenvalues =
        let result =
          Davidson.make ~threshold:1.e-6  ~n_states diagonal matrix_prod
        in
        Parallel.broadcast (lazy result)
      in
      let eigenvalues = Vec.map (fun x -> x +. e_shift) eigenvalues in
      (Conventions.rephase @@ Util.remove_epsilons eigenvectors), eigenvalues
    in

    let eigensystem_direct () =
      let m_H =
        Lazy.force m_H
      in
      let diagonal =
        Parallel.broadcast (lazy (
          Vec.init (Matrix.dim1 m_H) (fun i -> Matrix.get m_H i i)
        ))
      in
      let matrix_prod psi =
        let result =
          Matrix.parallel_mm ~transa:`T ~transb:`N psi m_H
          |> Matrix.transpose
        in
        Parallel.broadcast (lazy result)
      in
      let eigenvectors, eigenvalues =
        let result =
          Davidson.make ~threshold:1.e-6  ~n_states diagonal matrix_prod
        in
        Parallel.broadcast (lazy result)
      in
      let eigenvalues = Vec.map (fun x -> x +. e_shift) eigenvalues in
      (Conventions.rephase @@ Util.remove_epsilons eigenvectors), eigenvalues
    in

    match algo with
    | `Direct -> eigensystem_direct ()
    | `InCore -> eigensystem_incore ()
  )
  in
  { det_space ; e_shift ; m_H ; m_S2 ; eigensystem ; n_states }




let second_order_sum { det_space ; m_H ; m_S2 ; eigensystem ; n_states }
    list_holes  list_particles ?(unique=true) is_internal
    i_o1_alfa  alfa_o2_i  w_alfa  psi0 =

  let list_holes     = Array.of_list list_holes in
  let list_particles = Array.of_list list_particles in

  let psi0 =
    let stream =
      Ds.determinant_stream det_space
    in
    Array.init (Ds.size det_space) (fun i ->
        (Stream.next stream), (Mat.copy_row psi0 (i+1)) )
  in



  let symmetric = i_o1_alfa == alfa_o2_i in


  let det_contribution i =

    let already_generated =
      if unique then
        (fun alfa ->
          if is_internal alfa then
            true
          else
            let rec aux = function
              | -1 -> false
              | j  -> Determinant.degree (fst psi0.(j)) alfa <= 2
                      || aux (j-1)
            in
            aux (i-1)
        )
      else
        is_internal
    in

    let psi_filtered_idx =
      let rec aux accu =  function
        | j when j < i -> List.rev accu
        | j ->
          if Determinant.degree (fst psi0.(i)) (fst psi0.(j)) < 4
          then
            aux (j::accu) (j-1)
          else
            aux accu (j-1)
      in aux [] (Array.length psi0 - 1)
    in

    let psi_filtered =
      List.rev_map (fun i -> psi0.(i)) psi_filtered_idx
      |> List.rev
    in

    let psi_h_alfa alfa =
      List.fold_left (fun accu (det, coef) ->
(* Single state here *)
          accu +. coef.{1} *. (i_o1_alfa det alfa)) 0. psi_filtered
    in

    let alfa_h_psi alfa =
      List.fold_left (fun accu (det, coef) ->
(* Single state here *)
              accu +. coef.{1} *. (alfa_o2_i alfa det)) 0. psi_filtered
    in

    let psi_h_alfa_alfa_h_psi alfa =
      if symmetric then
        let x = psi_h_alfa alfa in x *. x
      else
        (psi_h_alfa alfa) *. (alfa_h_psi alfa)
    in

    let det_i = fst psi0.(i) in
    let w_alfa = w_alfa det_i in

    let same_spin =
      List.fold_left (fun accu spin ->
          accu +.
          Array.fold_left (fun accu particle ->
              accu +.
              Array.fold_left (fun accu hole ->
                  if hole = particle then accu else
                    let alfa =
                      Determinant.single_excitation spin hole particle det_i
                    in
                    if Determinant.is_none alfa then accu else

                      let single =
                        if already_generated alfa then 0. else
                          w_alfa alfa *. psi_h_alfa_alfa_h_psi alfa
                      in

                      let double =
                        Array.fold_left (fun accu particle' ->
                            if particle' >= particle || particle' = hole then
                              accu
                            else
                              accu +.
                              Array.fold_left (fun accu hole' ->
                                  if hole' = particle' || hole' = particle || hole' <= hole then
                                    accu
                                  else
                                      let alfa =
                                        Determinant.single_excitation
                                          spin  hole' particle' alfa
                                      in
                                      if Determinant.is_none alfa ||
                                         already_generated alfa then
                                        accu
                                      else
                                        accu +. w_alfa alfa *. psi_h_alfa_alfa_h_psi alfa
                                ) 0. list_holes
                          ) 0. list_particles
                      in
                      accu +. single +. double
                ) 0. list_holes
            ) 0. list_particles
        ) 0. [ Spin.Alfa ; Spin.Beta ]
    in

    let opposite_spin =
      Array.fold_left (fun accu particle ->
          accu +.
          Array.fold_left (fun accu hole ->
              if hole = particle then accu else
                let alfa =
                  Determinant.single_excitation Spin.Alfa hole particle det_i
                in
                if Determinant.is_none alfa then accu else

                  let double_ab =
                    Array.fold_left (fun accu particle' ->
                          accu +.
                          Array.fold_left (fun accu hole' ->
                              if hole' = particle' then accu else
                                let alfa =
                                  Determinant.double_excitation
                                    Spin.Alfa hole  particle
                                    Spin.Beta hole' particle' det_i
                                in
                                if Determinant.is_none alfa ||
                                  already_generated alfa then
                                  accu
                                else
                                  accu +. w_alfa alfa *. psi_h_alfa_alfa_h_psi alfa
                          ) 0. list_holes
                      ) 0. list_particles
                  in

                  accu +. double_ab
            ) 0. list_holes
        ) 0. list_particles
    in
    same_spin +. opposite_spin
  in

  Util.stream_range 0 (Array.length psi0 - 1)
  |> Farm.run ~ordered:true ~f:det_contribution
  |> Util.stream_to_list



let second_order_sum2 { det_space ; m_H ; m_S2 ; eigensystem ; n_states }
    list_holes  list_particles i_o1_alfa e0 psi0 =

  let psi0 =
    let stream =
      Ds.determinant_stream det_space
    in
    Array.init (Ds.size det_space) (fun i ->
        (Stream.next stream), (Mat.copy_row psi0 (i+1)) )
  in

  let determinants =

    Ds.determinants_array det_space
    |> Array.to_list
    |> List.concat_map (fun det_i ->
        [ Spin.Alfa ; Spin.Beta ]
        |>  List.concat_map (fun spin ->
            List.concat_map (fun particle ->
              List.map (fun hole ->
                  [ [ Determinant.single_excitation spin hole particle det_i ] ;
                    List.concat_map (fun particle' ->
                        List.map (fun hole' ->
                          Determinant.double_excitation
                            spin  hole particle
                            spin  hole' particle' det_i
                        ) list_holes
                    ) list_particles
                    ;
                    List.concat_map (fun particle' ->
                        List.map (fun hole' ->
                          Determinant.double_excitation
                            spin  hole particle
                            (Spin.other spin)  hole' particle' det_i
                        ) list_holes
                    ) list_particles
                  ]
                  |> List.concat
                ) list_holes
            ) list_particles
          )
        )
        |> List.concat
    |> List.filter (fun alfa -> not (Determinant.is_none alfa))
    |> List.sort_uniq compare
  in

  List.fold_left (fun accu alfa ->
    let alfa_o2 = i_o1_alfa alfa in
    let a_h_psi =
      Array.fold_left (fun accu (det,ci) -> ci.{1} *. (alfa_o2 det)) 0. psi0
    in
    accu +. (a_h_psi *. a_h_psi) /. (e0 -. (alfa_o2 alfa))
    ) 0. determinants




let is_internal det_space =
    let mo_class      = DeterminantSpace.mo_class det_space in
    let numbits = 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 numbits j)
        ) (Bitstring.zero numbits) l
    in
    let active_mask   = m (MOClass.active_mos   mo_class) in
    let occ_mask      = m (MOClass.core_mos     mo_class) in
    let inactive_mask = m (MOClass.inactive_mos mo_class) in
    let occ_mask = Bitstring.logor occ_mask inactive_mask in
    let neg_active_mask = Bitstring.lognot active_mask in
    fun a ->
      let alfa =
        Determinant.alfa a
        |> Spindeterminant.bitstring
      in
      if Bitstring.logand neg_active_mask alfa <> occ_mask then
        false
      else
        let beta =
          Determinant.beta a
          |> Spindeterminant.bitstring
        in
        Bitstring.logand neg_active_mask beta = occ_mask


let is_external det_space =
    let mo_class      = DeterminantSpace.mo_class det_space in
    let numbits = 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 numbits j)
        ) (Bitstring.zero numbits) l
    in
    let inactive_mask = m (MOClass.inactive_mos mo_class) in
    fun a ->
      let alfa =
        Determinant.alfa a
        |> Spindeterminant.bitstring
      in
      let n_a =
	Bitstring.(popcount @@ logand inactive_mask alfa)
      in
      match n_a with
      | 0 | 1 | 2 ->
	let beta =
	  Determinant.beta a
	  |> Spindeterminant.bitstring
	in
	n_a + Bitstring.(popcount @@ logand inactive_mask beta) < 3
      | _ ->   false


let pt2_en ci =

  let mo_basis  = Ds.mo_basis ci.det_space in
  let psi0, e0 = Parallel.broadcast ci.eigensystem in

  let i_o1_alfa = h_ij mo_basis in

  let w_alfa det_i =
    let one_e, two_e, _ = h_integrals mo_basis in
    let fock_diag_alfa, fock_diag_beta =
      Ds.fock_diag ci.det_space det_i
    in
    let h_aa alfa =
      match Excitation.of_det det_i alfa with
      | Excitation.Double (_,
            {hole = h ; particle = p ; spin = s },
            {hole = h'; particle = p'; spin = s'}) ->
              let fock_diag1 =
                if s = Spin.Alfa then
                  fock_diag_alfa
                else
                  fock_diag_beta
              in
              let fock_diag2 =
                if s' = Spin.Alfa then
                  fock_diag_alfa
                else
                  fock_diag_beta
              in
              fock_diag1.(0) -. fock_diag1.(h)
              +. (fock_diag1.(p ) -. two_e h p h p s s)
              -. (fock_diag2.(h') -. two_e h h' h h' s s' +. two_e p h' p h' s s')
              +. (fock_diag2.(p') -. two_e h p' h p' s s'
                  +. two_e p p' p p' s s' -. two_e h' p' h' p' s' s')
      | Excitation.Single (_,
            {hole = h ; particle = p ; spin = s }) ->
              let fock_diag =
                if s = Spin.Alfa then
                  fock_diag_alfa
                else
                  fock_diag_beta
              in
              fock_diag.(0) -. fock_diag.(h)
              +. (fock_diag.(p) -. two_e h p h p s s)
              |> ignore;
              h_ij mo_basis alfa alfa
      | _ -> e0.{1} -. 1.0
    in
    let e0 = e0.{1} in
    fun alfa ->
      1. /. (e0 -. h_aa alfa)
  in

  let mo_class = mo_class ci in
  let list_holes = List.concat
    [ MOClass.inactive_mos mo_class ; MOClass.active_mos mo_class ]
  and list_particles = List.concat
    [ MOClass.active_mos mo_class ; MOClass.virtual_mos mo_class ]
  in

  second_order_sum  ci  list_holes  list_particles
    (is_internal ci.det_space) i_o1_alfa  i_o1_alfa  w_alfa  psi0
  |> List.fold_left (+.) 0.


let _pt2_en ci =

  let mo_basis  = Ds.mo_basis ci.det_space in
  let psi0, e0 = Parallel.broadcast ci.eigensystem in

  let i_o1_alfa = h_ij mo_basis in

  let mo_class = mo_class ci in
  let list_holes = List.concat
    [ MOClass.inactive_mos mo_class ; MOClass.active_mos mo_class ]
  and list_particles = List.concat
    [ MOClass.active_mos mo_class ; MOClass.virtual_mos mo_class ]
  in
  second_order_sum2 ci list_holes  list_particles i_o1_alfa e0.{1} psi0



let pt2_mp ci =

  let mo_basis  = Ds.mo_basis ci.det_space in

  let i_o1_alfa = h_ij mo_basis in

  let eps = MOBasis.mo_energies mo_basis in
  let w_alfa det_i alfa=
    match Excitation.of_det  det_i  alfa with
    | Excitation.Single (_, { hole ; particle ; spin })->
        1./.(eps.{hole} -. eps.{particle})
    | Excitation.Double (_, { hole=h ; particle=p ; spin=s },
                            { hole=h'; particle=p'; spin=s'})->
        1./.(eps.{h} +. eps.{h'} -. eps.{p} -. eps.{p'})
    | _ -> assert false
  in

  let mo_class = mo_class ci in
  let list_holes = List.concat
    [ MOClass.inactive_mos mo_class ; MOClass.active_mos mo_class ]
  and list_particles = List.concat
    [ MOClass.active_mos mo_class ; MOClass.virtual_mos mo_class ]
  in

  let psi0, _ = Parallel.broadcast ci.eigensystem in
  second_order_sum  ci  list_holes  list_particles
    (is_internal ci.det_space) i_o1_alfa  i_o1_alfa  w_alfa  psi0
  |> List.fold_left (+.) 0.


let variance ci =

  let mo_basis  = Ds.mo_basis ci.det_space in
  let psi0, _ = Parallel.broadcast ci.eigensystem in

  let i_o1_alfa = h_ij mo_basis in

  let w_alfa _ _ = 1. in

  let mo_class = mo_class ci in
  let list_holes = List.concat
    [ MOClass.inactive_mos mo_class ; MOClass.active_mos mo_class ]
  and list_particles = List.concat
    [ MOClass.active_mos mo_class ; MOClass.virtual_mos mo_class ]
  in

  second_order_sum  ci  list_holes  list_particles
    (is_internal ci.det_space) i_o1_alfa  i_o1_alfa  w_alfa  psi0
  |> List.fold_left (+.) 0.




let pt2_en_reference ci =

  let mo_basis  = Ds.mo_basis ci.det_space in
  let psi0, e0 = Parallel.broadcast ci.eigensystem in

  let aux_basis = mo_basis in
  let ds =
    DeterminantSpace.fci_of_mo_basis ~frozen_core:false aux_basis
  in

  let det_stream =
    let e =
      let f = is_external ci.det_space  in
      function
      | None -> false
      | Some d -> f d
    in
    let stream =
      DeterminantSpace.determinant_stream ds
    in
    let rec next i =
      let det =
        try
          Some (Stream.next stream)
        with Stream.Failure -> None
      in
      if det = None then
        None
      else
        if e det then
          det
        else
          (next [@tailcall]) i
    in
    Stream.from next
  in

  let out_dets =
    det_stream
    |> stream_to_list
    |> Array.of_list
  in

  let in_dets =
    DeterminantSpace.determinants_array ci.det_space
  in

  let m_H_aux =
    let h_aa =
      Array.map (fun ki -> h_ij aux_basis ki ki) out_dets
    in
    Array.map (fun ki ->
      Array.mapi (fun j kj ->
        (h_ij aux_basis ki kj) /. (e0.{1} -. h_aa.(j))
      ) out_dets
    ) in_dets
    |> Mat.of_array
  in

  let m_F_aux =
    Array.map (fun ki ->
      Array.map (fun kj ->
        h_ij aux_basis ki kj
      ) out_dets
    ) in_dets
    |> Mat.of_array
  in

  let m_HF =
    gemm m_H_aux m_F_aux ~transb:`T
  in
  (gemm ~transa:`T psi0 @@ gemm m_HF psi0).{1,1}