open Lacaml.D
open Util

type t =
{
  p : Vec.t list;
  e : Vec.t list;
  m : int;
  mmax : int;
}

let make ?mmax:(mmax=15) () = 
  assert (mmax > 1);
  {
    p = [];
    e = [];
    m = 0 ;
    mmax;
   }


let append ~p ~e diis =
  let update v l =  
    if diis.m < diis.mmax then
      v :: l
    else
      match List.rev l with
      | [] -> assert false
      | _ :: rest -> v :: List.rev rest
  in
  { diis with
    p = update p diis.p;
    e = update e diis.e;
    m = min diis.mmax (diis.m+1);
  }


let next diis =
  let e = Mat.of_col_vecs_list diis.e 
  and p = Mat.of_col_vecs_list diis.p
  in
  let a =
    let rec aux m = 
      let a = Mat.make (m+1) (m+1) 1. in
      a.{m+1,m+1} <- 0.;
      ignore @@ lacpy ~b:a (gemm ~transa:`T ~m ~n:m e e);
      if m > 1 && sycon (lacpy a) > 1.e-14 then 
        (aux [@tailcall]) (m-1)
      else a
    in
    aux diis.m
  in
  let m = Mat.dim1 a - 1 in
  let c = Mat.make0 (m+1) 1 in
  c.{m+1,1} <- 1.;
  sysv a c;
  gemm p c ~k:m
  |> Mat.as_vec