open Lacaml.D
open Common

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

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_inplace m n a =
  assert ( (dim1 a) * (dim2 a) = m*n);
  let b = 
    to_bigarray a 
    |> Bigarray.genarray_of_array2
  in
  Bigarray.reshape_2 b m n

let reshape_vec_inplace m n v =
  assert ( Vector.dim v = m*n);
  let b = 
    Vector.to_bigarray_inplace v 
    |> Bigarray.genarray_of_array1
  in
  Bigarray.reshape_2 b m n

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

let transpose t =
  Mat.transpose_copy t

let transpose_inplace t =
  let mat = to_bigarray_inplace t in
  let d1 = dim1 t in
  let d2 = dim2 t in
  assert (d1 = d2);
  for j=1 to d1 do
    for i=1 to (j-1) do
      let ij, ji = mat.{i,j}, mat.{j,i} in
      mat.{i,j} <- ji;
      mat.{j,i} <- ij
    done;
  done

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 of_diag v =
  Vector.to_bigarray_inplace v
  |> Mat.of_diag 

let diag t =
  Mat.copy_diag t
  |> Vector.of_bigarray_inplace
       
let gemv_n_inplace ?m ?n ?(beta=0.) y ?(alpha=1.) ?(ar=1) ?(ac=1) t v =
  let y = Vector.to_bigarray_inplace y in
  let v = Vector.to_bigarray_inplace v in
  ignore @@ gemv ?m ?n ~beta ~trans:`N ~y ~alpha ~ar ~ac t v

let gemv_t_inplace ?m ?n ?(beta=0.) y ?(alpha=1.) ?(ar=1) ?(ac=1) t v =
  let y = Vector.to_bigarray_inplace y in
  let v = Vector.to_bigarray_inplace v in
  ignore @@ gemv ?m ?n ~beta ~trans:`T ~y ~alpha ~ar ~ac t v

let gemv_n ?m ?n ?(beta=0.) ?y ?(alpha=1.) ?(ar=1) ?(ac=1) t v =
  let v = Vector.to_bigarray_inplace v in
  let y =
    match y with
    | None -> None
    | Some y -> Some (Vector.to_bigarray_inplace y)
  in
  gemv ?m ?n ~beta ?y ~trans:`N ~alpha ~ar ~ac t v
  |> Vector.of_bigarray_inplace

let gemv_t ?m ?n ?(beta=0.) ?y ?(alpha=1.) ?(ar=1) ?(ac=1) t v =
  let v = Vector.to_bigarray_inplace v in
  let y =
    match y with
    | None -> None
    | Some y -> Some (Vector.to_bigarray_inplace y)
  in
  gemv ?m ?n ~beta ?y ~trans:`T ~alpha ~ar ~ac t v
  |> Vector.of_bigarray_inplace

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 ?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 ?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 ?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 ?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 of_array a =
  Bigarray.Array2.of_array Bigarray.Float64 Bigarray.fortran_layout a
    
let to_array a =
  let result = Array.make_matrix (Mat.dim1 a) (Mat.dim2 a) 0. in
  for i=1 to Mat.dim1 a do
    let result_i = result.(i-1) in
    for j=1 to Mat.dim2 a do
      result_i.(j-1) <- a.{i,j}
    done
  done;
  result
    
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 amax t =
  Mat.as_vec t |> amax 
    
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 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_of_file filename =
  let result =
    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 scale_rows_inplace v a =
  let v' = Vector.to_bigarray_inplace v in
  Mat.scal_rows v' a

let scale_rows v a =
  let a' = copy a in
  let v' = Vector.to_bigarray_inplace v in
  Mat.scal_rows v' 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

  
let exponential_iterative a =
  assert (dim1 a = dim2 a);
  let rec loop result accu n =
    let b = scale (1./.n) a in
    let new_accu = gemm accu b in
    let residual =
      sub new_accu accu
      |> amax
      |> abs_float
    in
    let result = add result new_accu in
    if residual > Constants.epsilon then
      loop result new_accu (n+.1.)
    else
      result
  in
  let id = identity (dim1 a) in
  loop id id 1.


let exponential a =
  
  let n = dim1 a in
  assert (n = dim2 a);
  let a2 = gemm a a in
  let (u, w, vt) = svd a2 in
  let tau = Vector.map (fun x -> -. sqrt x) w in
  let cos_tau = Vector.cos tau in
  let sin_tau_tau = Vector.mul (Vector.sin tau) (Vector.reci tau) in
  let result = 
    add (gemm (scale_cols u cos_tau) vt) (gemm (scale_cols u sin_tau_tau) @@ gemm vt a)
  in
  
  (* Post-condition: Check if exp(-A) * exp(A) = I *)
  let id = identity n in
  let test = 
    gemm_tn ~beta:(-.1.0) ~c:id result result
    |> amax
    |> abs_float
  in
  assert (test < Constants.epsilon);
  result


let to_file ~filename ?(sym=false) ?(cutoff=0.) t =

  let oc = open_out filename in
  let n = Mat.dim1 t in
  let m = Mat.dim2 t in

  if sym then 
    for j=1 to n  do
      for i=1 to j do
        if (abs_float (t.{i,j}) > cutoff) then
          Printf.fprintf oc "%4d %4d %20.12e\n" i j (t.{i,j})
      done;
    done
  else
    for j=1 to n  do
      for i=1 to m do
        if (abs_float (t.{i,j}) > cutoff) then
          Printf.fprintf oc "%4d %4d %20.12e\n" i j (t.{i,j})
      done;
    done;
  close_out oc



  
  
let (%:) t (i,j) = t.{i,j}

let set  t i j v = t.{i,j} <- v