open Lacaml.D

let full_ldl m_A =

  let n = Mat.dim1 m_A in

  assert (Mat.dim2 m_A = n);

  let v_D  = Vec.make0 n in

  let m_Lt = Mat.identity n in

  let v_D_inv  = Vec.make0 n in


  let compute_d j =
    let l_jk = 
      Mat.col m_Lt j
    in
    let l_jk__d_k = 
      Vec.mul ~n:(j-1) l_jk v_D
    in
    m_A.{j,j} -. dot ~n:(j-1) l_jk  l_jk__d_k
  in


  let compute_l i =
    let l_ik__d_k = 
      Mat.col m_Lt i
      |> Vec.mul ~n:(i-1) v_D
    in
    fun j ->
      assert (i > j);
      let l_jk = 
        Mat.col m_Lt j
      in
      v_D_inv.{j} *. ( m_A.{j,i} -. dot ~n:(j-1) l_ik__d_k l_jk )
  in


  for i=1 to n do
    for j=1 to (i-1) do
      m_Lt.{j,i} <- compute_l i j;
    done;
    let d_i = compute_d i in
    v_D.{i} <- d_i;
    v_D_inv.{i} <- 1. /. d_i;
  done;
  m_Lt, v_D




let pivoted_ldl threshold m_A =
(** {% $P A P^\dagger = L D L^\dagger$. %}
 Input : Matrix $A$
 Output : Matrices $L, D, P$.
*)

  let n = Mat.dim1 m_A in
  assert (Mat.dim2 m_A = n);

  let pi = Array.init (n+1) (fun i -> i) in

  let v_D  = Vec.init n (fun i -> abs_float m_A.{i,i}) in
  let m_Lt = Mat.identity n in

  let v_D_inv  = Vec.make0 n in

  let compute_d i =
    let l_ik = 
      Mat.col m_Lt i
    in
    let l_ik__d_k = 
      Vec.mul ~n:(i-1) l_ik v_D
    in
    m_A.{pi.(i),pi.(i)} -. dot ~n:(i-1) l_ik  l_ik__d_k
  in

  let compute_l i =
    let l_ik__d_k = 
      Mat.col m_Lt i
      |> Vec.mul ~n:(i-1) v_D
    in
    fun j ->
      assert (i > j);
      let l_jk = 
        Mat.col m_Lt j
      in
      v_D_inv.{j} *. ( m_A.{pi.(j),pi.(i)} -. dot ~n:(j-1) l_ik__d_k l_jk )
  in


  let maxloc v i =
    let rec aux pos value i =
      if i > n then
        pos
      else if v_D.{i} > value then
        (aux [@tailcall]) i v_D.{i}  (i+1)
      else
        (aux [@tailcall]) pos value (i+1)
    in
    aux i v.{i} (i+1)
  in


  let pivot i =
    let j = maxloc v_D i in
    let p_i, p_j = pi.(i), pi.(j) in
    pi.(i) <- p_j;
    pi.(j) <- p_i;
  in                                                                                          




  let () = 
    try
      for i=1 to n do
        pivot i;
        for j=1 to (i-1) do
          m_Lt.{j,i} <- compute_l i j;
        done;
        let d_i = compute_d i in
        if abs_float d_i < threshold then
          raise Exit;
        v_D.{i} <- d_i;
        v_D_inv.{i} <- 1. /. d_i;
      done
    with Exit -> ()
  in
  m_Lt, v_D, pi




let make_ldl ?(threshold=Constants.epsilon) m_A =
  pivoted_ldl m_A








let test_case () =

  let matrix_diff m_A m_B = 
    Mat.syrk_trace (Mat.sub m_A m_B) 
  in

  let vector_diff v_A v_B = 
    let v = Vec.sub v_A v_B in
    dot v v
  in

  let m_A = Mat.random 1000 1000 in
  let m_A = Mat.add m_A (Mat.transpose_copy m_A) in

  let test_full_small () =
    let m_A =
      Mat.of_array [| [|   4. ;  12. ; -16. |] ;
                      [|  12. ;  37. ; -43. |] ;
                      [| -16. ; -43. ;  98. |] |]
    in

    let m_L_ref = 
      Mat.of_array [| [|   1. ;   0. ;   0. |] ;
                      [|   3. ;   1. ;   0. |] ;
                      [|  -4. ;   5. ;   1. |] |]
    in

    let m_Lt_ref = 
      Mat.transpose_copy m_L_ref
    in

    let v_D_ref = 
      Vec.of_array [| 4. ; 1. ; 9. |]
    in

    let m_Lt, v_D = full_ldl m_A in
    Alcotest.(check (float 1.e-15)) "full L" 0. (matrix_diff m_Lt m_Lt_ref);
    Alcotest.(check (float 1.e-15)) "full D" 0. (vector_diff v_D v_D_ref);
    let m_D = Mat.of_diag v_D in
    let m_B = gemm ~transa:`T m_Lt @@ gemm m_D m_Lt in
    Alcotest.(check (float 1.e-15)) "full L" 0. (matrix_diff m_A m_B);
    ()
  in

  let test_full () =
    let m_Lt, v_D = full_ldl m_A in
    let m_D = Mat.of_diag v_D in
    let m_B = gemm ~transa:`T m_Lt @@ gemm m_D m_Lt in
    Alcotest.(check (float 1.e-15)) "full D" 0. (matrix_diff m_A m_B);
    ()
  in

  let test_pivoted () =
    let m_Lt, v_D, pi = pivoted_ldl 0. m_A  in
    let n = Mat.dim1 m_A in
    let m_P = Mat.make0 n n in
    for i=1 to n do
      m_P.{i,pi.(i)} <- 1.
    done;
    let m_D = Mat.of_diag v_D in
    let m_B =
      gemm ~transa:`T m_P  @@
      gemm ~transa:`T m_Lt @@
      gemm m_D @@
      gemm m_Lt m_P
    in
    Alcotest.(check (float 1.e-14)) "pivoted D" 0. (matrix_diff m_A m_B);
    ()
  in

  let test_truncated () =
    let m_Lt, v_D, pi = pivoted_ldl 0.001 m_A  in
    let n = Mat.dim1 m_Lt in
    let m_P = Mat.make0 n n in
    for i=1 to n do
      m_P.{i,pi.(i)} <- 1.
    done;
    let m_D = Mat.of_diag v_D in
    let m_B =
      gemm ~transa:`T m_P  @@
      gemm ~transa:`T m_Lt @@
      gemm m_D @@
      gemm m_Lt m_P
    in
    Alcotest.(check (float 1.e-3)) "full D" 0. (matrix_diff m_A m_B);
    ()
  in
  [
    "Small", `Quick, test_full_small;
    "Full", `Quick, test_full;
    "Pivoted", `Quick, test_pivoted ;
    "Truncated", `Quick, test_truncated ;
  ]