open Lacaml.D

type ('a, 'b) t = Mat.t

let dim1 t = Mat.dim1 t
let dim2 t = Mat.dim2 t

let out_of_place f t =
  let result = lacpy t in
  f result;
  result

let relabel t = t
  
let make   m n x = Mat.make   m n x
let make0  m n   = Mat.make0  m n
let create m n   = Mat.create m n
let identity m   = Mat.identity m

let fill_inplace t x = Mat.fill t x

let add a b = Mat.add a b
let sub a b = Mat.sub a b
let mul a b = Mat.mul a b
let div a b = Mat.div a b

let add_inplace ~c a b = ignore @@ Mat.add ~c a b
let sub_inplace ~c a b = ignore @@ Mat.sub ~c a b
let mul_inplace ~c a b = ignore @@ Mat.mul ~c a b
let div_inplace ~c a b = ignore @@ Mat.div ~c a b

let add_const_diag_inplace x a =
  Mat.add_const_diag x a

let add_const_diag x a =
  out_of_place (fun t -> add_const_diag_inplace x t) a

let add_const_inplace x a =
  ignore @@ Mat.add_const x ~b:a a

let add_const x a =
  Mat.add_const x a

let at t i j = t.{i,j}

external to_bigarray_inplace : ('a,'b) t -> (float, Stdlib.Bigarray.float64_elt, Stdlib.Bigarray.fortran_layout) Stdlib.Bigarray.Array2.t = "%identity"

external of_bigarray_inplace : (float, Stdlib.Bigarray.float64_elt, Stdlib.Bigarray.fortran_layout) Stdlib.Bigarray.Array2.t -> ('a,'b) t = "%identity"

let to_bigarray t = lacpy t
let of_bigarray t = lacpy t

let reshape a m n =
  assert ( (dim1 a) * (dim2 a) = m*n);
  let b = 
    to_bigarray a 
    |> Bigarray.genarray_of_array2
  in
  Bigarray.reshape_2 b m n

let col_inplace t j =
  Mat.col t j
  |> Vector.of_bigarray_inplace

(*
let col t j =
  Mat.col t j
  |> Vector.of_bigarray
*)


let to_col_vecs t =
  Mat.to_col_vecs t
  |> Array.map Vector.of_bigarray_inplace

let to_col_vecs_list t =
  Mat.to_col_vecs_list t
  |> List.rev_map Vector.of_bigarray_inplace
  |> List.rev

let detri_inplace t = 
  Mat.detri t

let detri t =
  out_of_place detri_inplace t

let as_vec_inplace t =
  Mat.as_vec t
  |> Vector.of_bigarray_inplace

let as_vec t = 
  lacpy t
  |> Mat.as_vec
  |> Vector.of_bigarray_inplace

let random ?rnd_state ?(from= -. 1.0) ?(range=2.0) m n =
  Mat.random ?rnd_state ~from ~range m n

let map_inplace f ~b a = 
  ignore @@ Mat.map f ~b a

let map f a = 
  Mat.map f a

let scale_inplace x t = 
  Mat.scal x t 

let scale x t = 
  out_of_place (fun t -> scale_inplace x t) t


let gemm_inplace ?m ?n ?k ?(beta=0.) ~c ?(transa=`N) ?(alpha=1.0) a ?(transb=`N) b =
  ignore @@ gemm ?m ?n ?k ~beta ~c ~transa ~alpha a ~transb b

let gemm_nn_inplace ?m ?n ?k ?(beta=0.) ~c ?(alpha=1.0) a b =
  gemm_inplace ?m ?n ?k ~beta ~c ~alpha a b ~transa:`N ~transb:`N

let gemm_nt_inplace ?m ?n ?k ?(beta=0.) ~c ?(alpha=1.0) a b =
  gemm_inplace ?m ?n ?k ~beta ~c ~alpha a b ~transa:`N ~transb:`T

let gemm_tn_inplace ?m ?n ?k ?(beta=0.) ~c ?(alpha=1.0) a b =
  gemm_inplace ?m ?n ?k ~beta ~c ~alpha a b ~transa:`T ~transb:`N

let gemm_tt_inplace ?m ?n ?k ?(beta=0.) ~c ?(alpha=1.0) a b =
  gemm_inplace ?m ?n ?k ~beta ~c ~alpha a b ~transa:`T ~transb:`T


let gemm ?m ?n ?k ?beta ?c ?(transa=`N) ?(alpha=1.0) a ?(transb=`N) b =
  let c =
    match c with
    | Some x -> Some (lacpy x)
    | None -> None
  in
  gemm ?m ?n ?k ?beta ?c ~transa ~alpha a ~transb b

let gemm_nn ?m ?n ?k ?(beta=0.) ?c ?(alpha=1.0) a b =
  gemm ?m ?n ?k ~beta ?c ~alpha a b ~transa:`N ~transb:`N

let gemm_nt ?m ?n ?k ?(beta=0.) ?c ?(alpha=1.0) a b =
  gemm ?m ?n ?k ~beta ?c ~alpha a b ~transa:`N ~transb:`T

let gemm_tn ?m ?n ?k ?(beta=0.) ?c ?(alpha=1.0) a b =
  gemm ?m ?n ?k ~beta ?c ~alpha a b ~transa:`T ~transb:`N

let gemm_tt ?m ?n ?k ?(beta=0.) ?c ?(alpha=1.0) a b =
  gemm ?m ?n ?k ~beta ?c ~alpha a b ~transa:`T ~transb:`T


let gemm_trace ?(transa=`N) a ?(transb=`N) b =
  Mat.gemm_trace ~transa a ~transb b
  
let gemm_nn_trace a b = gemm_trace a b ~transa:`N ~transb:`N 
let gemm_nt_trace a b = gemm_trace a b ~transa:`N ~transb:`T 
let gemm_tn_trace a b = gemm_trace a b ~transa:`T ~transb:`N 
let gemm_tt_trace a b = gemm_trace a b ~transa:`T ~transb:`T 

let init_cols = Mat.init_cols

let of_col_vecs a =
  Array.map Vector.to_bigarray_inplace a
  |> Mat.of_col_vecs 

let of_col_vecs_list a =
  List.rev_map Vector.to_bigarray_inplace a
  |> List.rev
  |> Mat.of_col_vecs_list 

let normalize_mat_inplace t = 
  let norm = Mat.as_vec t |> nrm2 in
  Mat.scal norm t

let normalize_mat t = 
  out_of_place normalize_mat_inplace t

let diagonalize_symm m_H =
  let m_V = lacpy m_H in
  let result =
    syevd ~vectors:true m_V
    |> Vector.of_bigarray_inplace
  in
  m_V, result

let sycon t =
  lacpy t
  |> sycon

let xt_o_x ~o ~x =
  gemm o x
  |> gemm ~transa:`T x

let x_o_xt ~o ~x =
  gemm o x ~transb:`T
  |> gemm x

let pp ppf m =
  let rows = Mat.dim1 m
  and cols = Mat.dim2 m
  in
  let rec aux first last =
    if (first <= last) then begin
        Format.fprintf ppf "@[\n\n    %a@]@ " (Lacaml.Io.pp_lfmat
          ~row_labels:
            (Array.init rows (fun i -> Printf.sprintf "%d  " (i + 1)))
          ~col_labels:
            (Array.init (min 5 (cols-first+1)) (fun i -> Printf.sprintf "-- %d --" (i + first) ))
          ~print_right:false
          ~print_foot:false
          () ) (lacpy ~ac:first ~n:(min 5 (cols-first+1)) m);
        (aux [@tailcall]) (first+5) last
    end
  in
  aux 1 cols

let sysv_inplace ~b a =
  sysv a b

let sysv ~b a = 
  out_of_place (fun b -> sysv_inplace ~b a) b

let debug_matrix name a =
  Format.printf "@[%s =\n@[%a@]@]@." name pp a

let outer_product_inplace m ?(alpha=1.0) u v =
  ger ~alpha (Vector.to_bigarray_inplace u) (Vector.to_bigarray_inplace v) m

let outer_product ?(alpha=1.0) u v =
  let m = make0 (Vector.dim u) (Vector.dim v) in
  outer_product_inplace m ~alpha u v;
  m

let matrix_of_file filename =
  let ic = Scanf.Scanning.open_in filename in
  let rec read_line accu =
    let result =
      try
        Some (Scanf.bscanf ic " %d %d %f" (fun i j v ->
          (i,j,v) :: accu))
      with End_of_file -> None
    in
    match result with
    | Some accu -> (read_line [@tailcall]) accu
    | None      -> List.rev accu
  in
  let data = read_line [] in
  Scanf.Scanning.close_in ic;
  let isize, jsize =
    List.fold_left (fun (accu_i,accu_j) (i,j,_) ->
      (max i accu_i, max j accu_j)) (0,0) data
  in
  let result =
    Lacaml.D.Mat.of_array
      (Array.make_matrix isize jsize 0.)
  in
  List.iter (fun (i,j,v) -> result.{i,j} <- v) data;
  result


let sym_matrix_of_file filename =
  let result =
    matrix_of_file filename
  in
  for j=1 to Mat.dim1 result do
    for i=1 to j do
      result.{j,i} <- result.{i,j}
    done;
  done;
  result


let copy ?m ?n ?br ?bc ?ar ?ac a =
  lacpy a ?m ?n ?br ?bc ?ar ?ac 

let copy_inplace ?m ?n ?br ?bc ~b ?ar ?ac a =
  ignore @@ lacpy ?m ?n ?br ?bc ~b ?ar ?ac a 

let scale_cols_inplace a v =
  Vector.to_bigarray_inplace v
  |> Mat.scal_cols a

let scale_cols a v =
  let a' = copy a in
  Vector.to_bigarray_inplace v
  |> Mat.scal_cols a' ;
  a'


let svd  a =
  let d, u, vt = 
    gesvd (lacpy a)
  in
  u, (Vector.of_bigarray_inplace d), vt

let svd' a =
  let d, u, vt = 
    gesvd (lacpy a)
  in
  u, (Vector.of_bigarray_inplace d), vt


let qr a =
  let result = lacpy a in
  let tau = geqrf result in
  let r =
    let r = to_bigarray result in
    for j=1 to Mat.dim2 r do
      for i=j+1 to Mat.dim2 r do
        r.{i,j} <- 0.
      done
    done;
    of_bigarray r
  in
  orgqr ~tau result;
  let q = result in
  q, r