open Lacaml.D

module Ds = DeterminantSpace
module Sd = Spindeterminant

type t =
  {
    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 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' = Spin.other 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) 
    ) )



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 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.map (fun (i,_,det_j) -> (i,det_j))
          in
          let doubles = 
            List.map (fun (i,_,det_j) -> (i,det_j)) doubles
          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 ki kj =
        let x = f 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) 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 r_singles j'
                        |  0 ->
                          let kj =
                            Determinant.of_spindeterminants j_alfa j_beta
                          in (update
                                (index_start.(i) + i') (index_start.(j) + j' + 1)
                                ki kj;
                              aux r_singles (j'+1);)
                        |  1 -> if (j' < Array.length j_dets) then aux 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 r_doubles j'
                        |  0 ->
                          let kj =
                            Determinant.of_spindeterminants j_alfa j_beta
                          in (update
                                (index_start.(i) + i') (index_start.(j) + j' + 1)
                                ki kj;
                              aux r_doubles (j'+1);)
                        |  1 -> if (j' < Array.length j_dets) then aux 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 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) 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
              ) b
            |> Array.to_list
            |> Util.list_some
          in
          let singles = 
            List.filter (fun (i,d,det_j) -> d < 2) doubles
            |> List.map (fun (i,_,det_j) -> (i,det_j))
          in
          let doubles = 
            List.map (fun (i,_,det_j) -> (i,det_j)) doubles
          in
          (singles, doubles)
        ) b
    in
    let a = Array.to_list a
    and b = Array.to_list b
    in

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

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

      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) 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, _ = degree_bb.(ib) in
                  List.iter (fun (j', j_beta) ->
                      let kj = Determinant.of_spindeterminants j_alfa j_beta in
                      update !i' (j' + !j) 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 = degree_bb.(ib) in
                  List.iter (fun (j', j_beta) ->
                      let kj = Determinant.of_spindeterminants j_alfa j_beta in
                      update !i' (j' + !j) ki kj
                    ) doubles;
                  incr i';
                ) b;
            | _ -> ();
          end;
          j := !j + n_beta
        ) a;
      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

    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+1) ndet;
      ) ;
    Matrix.sparse_of_vector_array result
  )






let make ?(n_states=1) det_space = 

  let m_H = 

    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 f = 
      match Ds.determinants det_space with
      | Ds.Arbitrary _ -> create_matrix_arbitrary
      | Ds.Spin      _ -> create_matrix_spin
    in
    f (fun ki kj -> h_ij mo_basis ki kj) det_space 
  in

  let m_S2 =
    let f = 
      match Ds.determinants det_space with
      | Ds.Arbitrary _ -> create_matrix_arbitrary
      | Ds.Spin      _ -> create_matrix_spin
    in
    f (fun ki kj -> CIMatrixElement.make_s2 ki kj) det_space 
  in

  let eigensystem = lazy (
    let m_H =
      Lazy.force m_H
    in
    let diagonal =
      Vec.init (Matrix.dim1 m_H) (fun i -> Matrix.get m_H i i)
    in
    let matrix_prod psi = 
      Matrix.mm ~transa:`T m_H psi
    in
    Davidson.make ~n_states diagonal matrix_prod
  )
  in
  { det_space ; m_H ; m_S2 ; eigensystem ; n_states }




let second_order_sum { det_space ; m_H ; m_S2 ; eigensystem ; n_states }
    i_o1_alfa  alfa_o2_i  w_alfa =

  let mo_basis = Ds.mo_basis det_space in

  let mo_class = DeterminantSpace.mo_class det_space in
  let mo_indices =
    let cls = MOClass.mo_class_array mo_class in
    Util.list_range 1 (MOBasis.size mo_basis)
    |> List.filter (fun i -> match cls.(i) with
        | MOClass.Deleted  _ 
        | MOClass.Core     _ -> false
        | _ -> true
      )
  in

  let psi0 =
    let psi0, _ = Lazy.force eigensystem in

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


    (*
  let is_internal =
    let m l =
      List.fold_left (fun accu i ->
          let j = i-1 in Z.(logor accu (shift_left one j))
        ) Z.zero 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 = Z.logor occ_mask inactive_mask in
    let neg_active_mask = Z.lognot active_mask in
    fun a ->
      let alfa =
        Determinant.alfa a
        |> Spindeterminant.bitstring
      in
      if Z.logand neg_active_mask alfa <> occ_mask then
        false
      else
        let beta =
          Determinant.beta a
          |> Spindeterminant.bitstring
        in
        Z.logand neg_active_mask beta = occ_mask
  in
        *)


  let symmetric = i_o1_alfa == alfa_o2_i in

  let is_internal alfa =
    let rec aux = function
      | -1 -> false
      | j  -> if (Determinant.degree (fst psi0.(j)) alfa = 0) then true
              else aux (j-1)
    in
    aux (Array.length psi0 - 1)
  in


  let det_contribution i =

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

    let psi_h_alfa alfa =
      List.fold_left (fun accu (det, coef) ->
          accu +. coef *. (i_o1_alfa det alfa)) 0. psi_filtered
    in

    let alfa_h_psi =
      if symmetric then 
        psi_h_alfa 
      else
        fun alfa ->
          List.fold_left (fun accu (det, coef) ->
              accu +. coef *. (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 +.
          List.fold_left (fun accu particle ->
              accu +.
              List.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 =
                        List.fold_left (fun accu particle' ->
                            if particle' > particle then
                              accu
                            else
                              accu +.
                              List.fold_left (fun accu hole' ->
                                  if hole' = particle' || hole' < hole then
                                    accu
                                  else
                                      let alfa = 
                                        Determinant.double_excitation
                                          spin  hole  particle
                                          spin  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. mo_indices
                          ) 0. mo_indices
                      in
                      accu +. single +. double
                ) 0. mo_indices
            ) 0. mo_indices
        ) 0. [ Spin.Alfa ; Spin.Beta ]
    in

    let opposite_spin =
      List.fold_left (fun accu particle ->
          accu +.
          List.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 =
                    List.fold_left (fun accu particle' ->
                        accu +.
                        List.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. mo_indices
                      ) 0. mo_indices
                  in
                  accu +. double
            ) 0. mo_indices
        ) 0. mo_indices
    in
    same_spin +. opposite_spin
  in

  Array.mapi (fun i (_,c_i) -> det_contribution i) psi0
  |> Array.fold_left (+.) 0.



let pt2_en ci =

  let mo_basis  = Ds.mo_basis ci.det_space in
  let _psi0, e0 = Lazy.force ci.eigensystem in

  let i_o1_alfa = h_ij mo_basis in

  let w_alfa _ =
    let e0 = e0.{1} in 
    let h_aa alfa = h_ij mo_basis alfa alfa in
    fun alfa ->
      1. /. (e0 -. h_aa alfa)
  in

  second_order_sum  ci  i_o1_alfa  i_o1_alfa  w_alfa



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

  second_order_sum  ci  i_o1_alfa  i_o1_alfa  w_alfa


let variance ci =

  let mo_basis  = Ds.mo_basis ci.det_space in

  let i_o1_alfa = h_ij mo_basis in

  let w_alfa _ _ = 1. in

  second_order_sum  ci  i_o1_alfa  i_o1_alfa  w_alfa