open Lacaml.D

module Ds = DeterminantSpace

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 }