open Common

let (%.) = Vector.(%.)
let (%:) = Matrix.(%:)

type t

let make
    ?guess
    ?(n_states=1)
    ?(n_iter=8)
    ?(threshold=1.e-7)
    diagonal
    matrix_prod
  =

  (* Size of the matrix to diagonalize *)
  let n = Vector.dim diagonal in

  let m = n_states in

  (* Create guess vectors u, with unknown vectors initialized to unity. *)
  let init_vectors m =
    let result = Matrix.make n m 0. in
    for i=1 to m do
      Matrix.set result i i 1.;
    done;
    result
  in

  let guess =
    match guess with
    | Some vectors -> vectors
    | None -> init_vectors m
  in

  let guess =
    if Matrix.dim2 guess = n_states then guess
    else
      (Matrix.to_col_vecs_list guess) @
        (Matrix.to_col_vecs_list (init_vectors (m-(Matrix.dim2 guess))) )
      |> Matrix.of_col_vecs_list
  in

  let pick_new u =
    Matrix.to_col_vecs_list u
    |> Util.list_pack m
    |> List.rev
    |> List.hd
  in

  let u_new = Matrix.to_col_vecs_list guess in

  let rec iteration u u_new w iter macro =
    (* u is a list of orthonormal vectors, on which the operator has
       been applied : w = op.u
       u_new is a list of vectors which will increase the size of the
       space.
    *)
    (* Orthonormalize input vectors u_new *)
    let u_new_ortho =
      List.concat [u ; u_new]
      |> Matrix.of_col_vecs_list
      |> Orthonormalization.qr_ortho
      |> pick_new
    in

    (* Apply the operator to the m last vectors *)
    let w_new =
      matrix_prod (
        u_new_ortho
        |> Matrix.of_col_vecs_list)
      |> Matrix.to_col_vecs_list
    in

    (* Data for the next iteration *)
    let u_next =
      List.concat [ u ; u_new_ortho ]
    and w_next =
      List.concat [ w ; w_new ]
    in

    (* Build the small matrix h = <U_k | W_l> *)
    let m_U =
      Matrix.of_col_vecs_list u_next
    and m_W =
      Matrix.of_col_vecs_list w_next
    in
    let m_h =
      Matrix.gemm_tn m_U m_W
    in

    (* Diagonalize h *)
    let y, lambda =
      Matrix.diagonalize_symm m_h
    in

    (* Express m lowest eigenvectors of h in the large basis *)
    let m_new_U =
      Matrix.gemm  ~n:m  m_U  y
    and m_new_W =
      Matrix.gemm  ~n:m  m_W  y
    in

    (* Compute the residual as proposed new vectors *)
    let u_proposed =
      Matrix.init_cols n m (fun i k ->
          let delta = lambda%.(k) -. diagonal%.(i) in
          let delta =
            if abs_float delta > 1.e-2  then delta
            else if delta > 0. then 1.e-2
            else (-1.e-2)
          in
            (lambda%.(k) *. (m_new_U%:(i,k)) -. (m_new_W%:(i,k)) ) /. delta
          )
      |> Matrix.to_col_vecs_list
    in


    let residual_norms = List.rev @@ List.rev_map Vector.norm u_proposed in
    let residual_norm  =
      List.fold_left (fun accu i -> accu +. i *. i) 0. residual_norms
      |> sqrt
    in

    Printf.printf "%3d " iter;
    Vector.iteri (fun i x -> if (i<=m) then Printf.printf "%16.10f  " x) lambda;
    Printf.printf "%16.8e%!\n" residual_norm;

    (* Make new vectors sparse *)

    let u_proposed =
      Matrix.of_col_vecs_list u_proposed
    in
    let maxu = Matrix.norm ~l:`Linf u_proposed in
    let thr = maxu *. 0.00001 in
    let u_proposed =
      Matrix.map (fun x -> if abs_float x < thr then 0. else x) u_proposed
      |> Matrix.to_col_vecs_list
    in
    (*
    Format.printf "%a@." Matrix.pp_matrix @@ m_new_U;
    *)


    if residual_norm > threshold && macro > 0 then
      let u_next, w_next, iter, macro =
        if iter = n_iter then
          m_new_U |> pick_new,
          m_new_W |> pick_new,
          0, (macro-1)
        else
          u_next, w_next, (iter+1), macro
      in
      (iteration [@tailcall]) u_next u_proposed w_next iter macro
    else
      (m_new_U |> pick_new |> Matrix.of_col_vecs_list), lambda

  in
  iteration [] u_new [] 1 30