open Lacaml.D

let epsilon = Constants.epsilon

type index_value = 
  {
    index: int;
    value: float
  }

type sparse_vector =
  {
    n: int;
    v: index_value array
  }

type t =
  | Dense  of Vec.t
  | Sparse of sparse_vector


let is_sparse = function
  | Sparse _ -> true
  | Dense  _ -> false


let is_dense = function
  | Sparse _ -> false
  | Dense  _ -> true


exception Found of float

let get = function
  | Dense  v ->  (fun i -> v.{i})
  | Sparse { n ; v } ->  (fun i ->
      if i < 1 || i > n then invalid_arg "index out of bounds";
      try
        Array.iter (fun {index ; value} ->
            if index=i then
              raise (Found value)) v; 
        raise Not_found
      with
      | Not_found -> 0.
      | Found x -> x
    )



let dim = function
  | Dense v -> Vec.dim v
  | Sparse {n ; v} -> n


let sparse_of_dense ?(threshold=epsilon) = function
  | Sparse _ -> invalid_arg "Expected a dense vector"
  | Dense v ->
    let rec aux accu = function
      | 0 -> accu |> Array.of_list
      | i ->
        let x = v.{i} in
        if abs_float x < threshold then
          (aux [@tailcall]) accu (i-1)
        else
          (aux [@tailcall]) ({index=i ; value=x}::accu) (i-1)
    in
    let n = Vec.dim v in
    Sparse { n ; v=aux [] n }


let rec to_assoc_list ?(threshold=epsilon) = function
  | Sparse {n ; v} -> Array.map (fun {index ; value} -> (index, value)) v |> Array.to_list
  | Dense v -> to_assoc_list @@ sparse_of_dense ~threshold (Dense v)


let dense_of_sparse = function
  | Dense _ -> invalid_arg "Expected a sparse vector"
  | Sparse {n ; v} -> 
    let v' = Vec.make0 n in
    Array.iter (fun {index ; value} -> v'.{index} <- value) v;
    Dense v'


let dense_of_vec v = Dense v


let sparse_of_vec ?(threshold=epsilon) v = 
  dense_of_vec v
  |> sparse_of_dense ~threshold 


let sparse_of_assoc_list n v =
  Sparse { n ;
    v = List.map (fun (index, value) -> {index ; value}) v
        |> Array.of_list
  }


let rec to_vec = function
  | Dense  v -> v
  | Sparse v -> dense_of_sparse (Sparse v) |> to_vec 



let scale ?(threshold=epsilon) x = function
  | Dense  v -> let v' = copy v in (scal x v'; Dense v')
  | Sparse {n ; v} ->
    Sparse {n ; v = Array.map (fun {index ; value} -> { index ; value=x *. value} ) v }


let rec neg = function
  | Dense  v -> Dense (Vec.neg v)
  | Sparse {n ; v} ->
      Sparse {n ; v = Array.map (fun {index ; value} -> {index ; value = -. value}) v}



let axpy ?(threshold=epsilon) ?(alpha=1.) x y = 
  if dim x <> dim y then
    invalid_arg "Inconsistent dimensions";

  match x, y with
  | Dense  x      , Dense  y -> Dense (let y = copy y in axpy ~alpha x y ; y)
  | Sparse {n ; v}, Dense  y ->
    begin
      let v' = copy y in
      Array.iter (fun {index ; value} -> v'.{index} <- v'.{index} +. alpha *. value) v;
      sparse_of_vec ~threshold v'
    end
  | Dense  x      , Sparse {n ; v}  ->
    begin
      let v' = copy x in
      scal alpha v';
      Array.iter (fun {index ; value} -> v'.{index} <- v'.{index} +. value) v;
      sparse_of_vec ~threshold v'
    end
  | Sparse {n ; v}, Sparse {n=n' ; v=v'} -> 
    begin
      let rec aux accu k l =
        match k < Array.length v, l < Array.length v' with
        | true, true ->
        begin
          let {index=i ; value=x} = v.(k)
          and {index=j ; value=y} = v'.(l)
          in
          match compare i j with
          | -1 -> 
            let z = alpha *. x in
            let new_accu = 
              if abs_float z > threshold then
                {index=i ; value=z} :: accu
              else
                accu
            in (aux [@tailcall]) new_accu (k+1) l
          | 1 ->
            let new_accu = 
              if abs_float y > threshold then
                {index=j ; value=y} :: accu
              else
                accu
            in (aux [@tailcall]) new_accu k (l+1)
          | 0 ->
            let z = alpha *. x +. y in
            let new_accu =
              if abs_float z > threshold then
                {index=i ; value=z} :: accu
              else
                accu
            in (aux [@tailcall]) new_accu (k+1) (l+1)
          | _ -> assert false
          end

        | true, false ->
          let {index=i ; value=x} = v.(k) in
          (aux [@tailcall]) ({index=i ; value=alpha *. x}::accu) (k+1) l

        | false, true ->
          (aux [@tailcall]) (v'.(l)::accu) k (l+1)

        | false, false -> {n ; v=List.rev accu |> Array.of_list}
      in
      Sparse (aux [] 0 0)
    end

let add = axpy ~alpha:1.

let sub ?(threshold=epsilon) x y = add ~threshold x @@ neg y

let pp ppf = function
  | Dense  m -> Util.pp_float_array ppf @@ Vec.to_array m 
  | Sparse {n ; v} -> 
    begin
      Format.fprintf ppf "@[[ %d | " n;
      Array.iter (fun {index ; value} -> Format.fprintf ppf "@[(%d, %f); @]" index value) v;
      Format.fprintf ppf "]@]"
    end
  
  

let dot v v' = 

  let d_d v v' =
    dot v v'
  in

  let d_sp v' {n ; v} =
    if n <> Vec.dim v' then
      invalid_arg "Inconsistent dimensions";
    Array.fold_left (fun accu {index ; value} -> accu +. value *. v'.{index}) 0. v
  in

  let sp_sp {n ; v} {n=n' ; v=v'} =
    if n <> n' then
      invalid_arg "Inconsistent dimensions";

    let rec aux accu k l = 
      match Array.length v > k, Array.length v' > l with
      | true, true ->
        let {index=i ; value=x} = v.(k)
        and {index=j ; value=y} = v'.(l)
        in
        begin
          match compare i j with
          | -1 -> (aux [@tailcall])  accu (k+1) l
          |  1 -> (aux [@tailcall])  accu k (l+1)
          |  0 -> (aux [@tailcall]) (accu +. x *. y) (k+1) (l+1)
          |  _ -> assert false
        end
      | _ -> accu
    in
    aux 0. 0 0
  in

  match v, v' with 
  | (Dense  v), (Dense  v') -> d_d v v'
  | (Sparse v), (Sparse v') -> sp_sp v v'
  | (Dense  v), (Sparse v') -> d_sp v v'
  | (Sparse v), (Dense  v') -> d_sp v' v


let norm v = sqrt @@ dot v v

let test_case () =

  let x1 = Vec.map (fun x -> if abs_float x < 0.6 then 0. else x) (Vec.random 100)
  and x2 = Vec.map (fun x -> if abs_float x < 0.3 then 0. else x) (Vec.random 100)
  in
  let x3 = Vec.map (fun x -> 2. *. x) x1
  and x4 = Vec.add x1 x2
  and x5 = Vec.sub x1 x2
  and x6 =
    let v = copy x2 in
    Lacaml.D.axpy ~alpha:3. x1 v;
    v
  in

  let v1 = dense_of_vec x1
  and v2 = dense_of_vec x2
  and v3 = dense_of_vec x3
  and v4 = dense_of_vec x4
  and v5 = dense_of_vec x5
  and v6 = dense_of_vec x6
  in

  let v1_s = sparse_of_vec x1
  and v2_s = sparse_of_vec x2
  and v3_s = sparse_of_vec x3
  and v4_s = sparse_of_vec x4
  and v5_s = sparse_of_vec x5
  and v6_s = sparse_of_vec x6
  in

  let zero   = dense_of_vec  (Vec.make0 100)
  and zero_s = sparse_of_vec (Vec.make0 100)
  in

  let test_conversion () =
    Alcotest.(check bool) "sparse -> dense  1" true (dense_of_sparse v1_s = v1  );
    Alcotest.(check bool) "sparse -> dense  2" true (dense_of_sparse v2_s = v2  );
    Alcotest.(check bool) "dense  -> sparse 1" true (sparse_of_dense v1   = v1_s);
    Alcotest.(check bool) "dense  -> sparse 2" true (sparse_of_dense v2   = v2_s);
  in
  let test_operations () =
    Alcotest.(check  bool)  "dense   scale"        true  (scale  2.    v1    =  v3);
    Alcotest.(check  bool)  "sparse  scale"        true  (scale  2.    v1_s  =  v3_s);

    Alcotest.(check  bool)  "dense   dense   add"  true  (add    v1    v2    =  v4);
    Alcotest.(check  bool)  "dense   sparse  add"  true  (add    v1    v2_s  =  v4_s);
    Alcotest.(check  bool)  "sparse  dense   add"  true  (add    v1_s  v2    =  v4_s);
    Alcotest.(check  bool)  "sparse  dense   add"  true  (add    v1    v2_s  =  v4_s);
    Alcotest.(check  bool)  "sparse  sparse  add"  true  (add    v1_s  v2_s  =  v4_s);

    Alcotest.(check  bool)  "dense   dense   sub"  true  (sub    v1    v2    =  v5);
    Alcotest.(check  bool)  "dense   sparse  sub"  true  (sub    v1    v2_s  =  v5_s);
    Alcotest.(check  bool)  "sparse  dense   sub"  true  (sub    v1_s  v2    =  v5_s);
    Alcotest.(check  bool)  "sparse  dense   sub"  true  (sub    v1    v2_s  =  v5_s);
    Alcotest.(check  bool)  "sparse  sparse  sub"  true  (sub    v1_s  v2_s  =  v5_s);

    Alcotest.(check  bool)  "dense   dense   sub"  true  (sub    v1    v1    =  zero);
    Alcotest.(check  bool)  "dense   sparse  sub"  true  (sub    v1    v1_s  =  zero_s);
    Alcotest.(check  bool)  "sparse  dense   sub"  true  (sub    v1_s  v1    =  zero_s);
    Alcotest.(check  bool)  "sparse  sparse  sub"  true  (sub    v1_s  v1_s  =  zero_s);

    Alcotest.(check  bool)  "dense   dense   axpy"  true  (axpy ~alpha:3.  v1    v2    =  v6);
    Alcotest.(check  bool)  "dense   sparse  axpy"  true  (sub ~threshold:1.e-12 (axpy ~alpha:3.  v1    v2_s) v6_s = zero_s);
    Alcotest.(check  bool)  "sparse  dense   axpy"  true  (sub ~threshold:1.e-12 (axpy ~alpha:3.  v1_s  v2) v6_s = zero_s);

    Alcotest.(check  bool)  "sparse  sparse  axpy"  true  (sub ~threshold:1.e-12 (axpy ~alpha:3.  v1_s  v2_s) v6_s = zero_s);
  in
  let test_dot () =
    let d1d2 = Lacaml.D.dot x1 x2 
    and d1d1 = Lacaml.D.dot x1 x1
    and d2d2 = Lacaml.D.dot x2 x2
    in
    Alcotest.(check (float 1.e-10)) "sparse x dense  1" (dot v1_s v2  )  d1d2;
    Alcotest.(check (float 1.e-10)) "sparse x dense  2" (dot v1_s v1  )  d1d1;
    Alcotest.(check (float 1.e-10)) "sparse x dense  3" (dot v2_s v2  )  d2d2;
    Alcotest.(check (float 1.e-10)) "dense  x sparse 1" (dot v1   v2_s)  d1d2;
    Alcotest.(check (float 1.e-10)) "dense  x sparse 2" (dot v1   v1_s)  d1d1;
    Alcotest.(check (float 1.e-10)) "dense  x sparse 3" (dot v2   v2_s)  d2d2;
    Alcotest.(check (float 1.e-10)) "sparse x sparse 1" (dot v1_s v2_s)  d1d2;
    Alcotest.(check (float 1.e-10)) "sparse x sparse 2" (dot v1_s v1_s)  d1d1;
    Alcotest.(check (float 1.e-10)) "sparse x sparse 3" (dot v2_s v2_s)  d2d2;
  in
  [
    "Conversion", `Quick, test_conversion;
    "Operations", `Quick, test_operations;
    "Dot product", `Quick, test_dot;
  ]