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 =

  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 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.{pi.(j),pi.(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.{pi.(j),pi.(i)} -. dot ~n:(j-1) l_ik__d_k l_jk )
  in

  let pivot i =
    let rec aux (pos,value) i = 
      if i > n then
        pos
      else if v_D.{i} > value then
        aux i v_D.{i}  (i+1)
      else
        aux pos value (i+1)
    in
    let j = aux i v_D.{i} (i+1) in
    let p_i, p_j = pi.(i), pi.(j) in
    pi.(i) <- pj;
    pi.(j) <- pi;
  in

  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
    v_D.{i} <- d_i;
    v_D_inv.{i} <- 1. /. d_i;
  done;
  m_Lt, v_D





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 test_full () =
    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)
  in
  [
    "Full", `Quick, test_full;
  ]