open Lacaml.D

module Ds = Determinant_space

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;
  }

let det_space    t = t.det_space

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 make det_space = 
  let ndet = Ds.size det_space in
  let det  = Ds.determinants det_space in
  let mo_basis = Ds.mo_basis det_space in
  

  (*
  let m_H = lazy (
    let ntasks = int_of_float @@ sqrt @@ float_of_int ndet in
    List.init ntasks (fun i -> 
      let m = 
        if i = (ntasks-1) then
          (ndet - i*ntasks) 
        else
          ntasks
      in
      List.init m (fun j -> i*ntasks + j)
      )
    |> Stream.of_list
    |> Farm.run ~ordered:true
      ~f:(fun l ->
          Printf.eprintf "Start\n%!";
      List.map (fun i ->
          let ki = det.(i) in
          Printf.eprintf "%d / %d\n%!" i ndet;
          Array.init ndet (fun j -> let kj = det.(j) in
                            h_ij mo_basis ki kj)
          ) l)
    |> Util.stream_to_list
    |> List.concat
    |> List.map Vec.of_array
    |> Mat.of_col_vecs_list
  ) in
  *)

  (* Parallel
  let m_H = lazy (
    let h =
      if Parallel.master then
        Array.make_matrix ndet ndet 0.
        |> Mat.of_array
      else
        Array.make_matrix 1 1 0.
        |> Mat.of_array
    in
    List.init ndet (fun i -> i)
    |> Stream.of_list
    |> Farm.run ~ordered:false
      ~f:(fun i ->
          let ki = det.(i) in
          Printf.eprintf "%8d / %8d\r%!" i ndet;
          List.init ndet (fun j ->
            let kj = det.(j) in
            let x = h_ij mo_basis ki kj in
            if x <> 0. then Some (i,j,x) else None
          )
          |> Util.list_some
          )
    |> Util.stream_to_list
    |> List.iter (fun l -> if Parallel.master then
      List.iter (fun (i,j,x) -> h.{i+1,j+1} <- x) l);
    Parallel.broadcast (lazy h)
  ) in
  *)


  let m_H = lazy (
    let v = Vec.make0 ndet in
    Array.init ndet (fun i -> let ki = det.(i) in
      Printf.eprintf "%8d / %8d\r%!" i ndet;
      let j = ref 1 in
      Ds.determinant_stream det_space
      |> Stream.iter (fun kj ->
        v.{!j} <- h_ij mo_basis ki kj ; incr j);
      Vector.sparse_of_vec v)
    |> Matrix.sparse_of_vector_array
  )
  in
  let m_S2 = lazy (
    let v = Vec.make0 ndet in
    Array.init ndet (fun i -> let ki = det.(i) in
      Array.iteri (fun j kj ->
        v.{j+1} <- CIMatrixElement.make_s2 ki kj) det;
      Vector.sparse_of_vec v)
    |> Matrix.sparse_of_vector_array
  )
  in
  let eigensystem = lazy (
    let m_H = Lazy.force m_H in
    (*
    Parallel.broadcast @@
      lazy (Util.diagonalize_symm m_H)
      *)
    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 diagonal matrix_prod
  )
  in
  { det_space ; m_H ; m_S2 ; eigensystem }