type t
type ws

let () = Error.init ()

type kind =
  | DAUBECHIES
  | DAUBECHIES_CENTERED
  | HAAR
  | HAAR_CENTERED
  | BSPLINE
  | BSPLINE_CENTERED

type direction =
  | FORWARD
  | BACKWARD

external _alloc : kind -> int -> t = "ml_gsl_wavelet_alloc"
external _free : t -> unit = "ml_gsl_wavelet_free"
let make kind size =
  let w = _alloc kind size in
  Gc.finalise _free w ;
  w

external name : t -> string = "ml_gsl_wavelet_name"

external _workspace_alloc : int -> ws = "ml_gsl_wavelet_workspace_alloc"
external _workspace_free  : ws -> unit = "ml_gsl_wavelet_workspace_free"
let workspace_make size =
  let ws = _workspace_alloc size in
  Gc.finalise _workspace_free ws ;
  ws
external workspace_size : ws -> int = "ml_gsl_wavelet_workspace_size"

external _transform : 
  t -> direction -> Vector_flat.vector -> ws -> unit 
  = "ml_gsl_wavelet_transform"
external _transform_bigarray : 
  t -> direction -> Vector.vector -> ws -> unit 
  = "ml_gsl_wavelet_transform_bigarray"

let with_workspace ws length f arg =
  let workspace = 
    match ws with 
    | Some ws -> ws
    | None -> _workspace_alloc (length arg) in
  try 
    f arg workspace ;
    if ws = None then _workspace_free workspace
  with exn ->
    if ws = None then _workspace_free workspace ;
    raise exn

let transform_vector_flat w dir ?ws =
  with_workspace ws 
    Vector_flat.length
    (_transform w dir)

let transform_vector w dir ?ws =
  with_workspace ws 
    Vector.length
    (_transform_bigarray w dir)

let transform_gen w dir ?ws = function
  | `V v  -> transform_vector      w dir ?ws v
  | `VF v -> transform_vector_flat w dir ?ws v

let transform_array w dir ?ws ?stride ?off ?len arr =
  transform_vector_flat w dir ?ws
    (Vector_flat.view_array ?stride ?off ?len arr)

let transform_forward w = transform_array w FORWARD
let transform_inverse w = transform_array w BACKWARD

type ordering =
  | STANDARD
  | NON_STANDARD

external _transform_2d :
  t -> ordering -> direction -> Matrix_flat.matrix -> ws -> unit
  = "ml_gsl_wavelet2d_transform_matrix"
external _transform_2d_bigarray :
  t -> ordering -> direction -> Matrix.matrix -> ws -> unit
  = "ml_gsl_wavelet2d_transform_matrix"
external _transform_2d_gen :
  t -> ordering -> direction -> [< Vectmat.mat] -> ws -> unit
  = "ml_gsl_wavelet2d_transform_matrix"

let transform_matrix_flat w order dir ?ws =
  with_workspace ws
    (fun m -> fst (Matrix_flat.dims m))
    (_transform_2d w order dir)

let transform_matrix w order dir ?ws =
  with_workspace ws
    (fun m -> fst (Matrix.dims m))
    (_transform_2d_bigarray w order dir)

let transform_matrix_gen w order dir ?ws =
  with_workspace ws
    (fun m -> fst (Vectmat.dims m))
    (_transform_2d_gen w order dir)
