(* gsl-ocaml - OCaml interface to GSL                       *)
(* Copyright (©) 2002-2012 - Olivier Andrieu                *)
(* Distributed under the terms of the GPL version 3         *)

let () = Error.init ()

open Vectmat

type ws
external alloc_ws : int -> int -> ws
    = "ml_gsl_multifit_linear_alloc"

external free_ws : ws -> unit
    = "ml_gsl_multifit_linear_free"

let make ~n ~p =
  let ws = alloc_ws n p in
  Gc.finalise free_ws ws ;
  ws

external _linear : ?weight:vec -> x:mat -> y:vec -> c:vec -> cov:mat -> 
      ws -> float
	  = "ml_gsl_multifit_linear_bc" "ml_gsl_multifit_linear"

let linear ?weight x y =
  let (n,p) = Vectmat.dims x in
  let dy = Vectmat.length y in
  if dy <> n
  then invalid_arg "Multifit.linear: wrong dimensions" ;
  Misc.may weight 
    (fun w -> 
      if Vectmat.length w <> n
      then invalid_arg "Multifit.linear: wrong dimensions") ;
  let c = Vector.create p in
  let cov = Matrix.create p p in
  let ws = alloc_ws n p in
  try
    let chisq = _linear ?weight ~x ~y ~c:(`V c) ~cov:(`M cov) ws in
    free_ws ws ;
    (c, cov, chisq)
  with exn ->
    free_ws ws ; raise exn

external linear_est : x:vec -> c:vec -> cov:mat -> Fun.result
    = "ml_gsl_multifit_linear_est"

let fit_poly ?weight ~x ~y order =
  let n = Array.length y in
  let x_mat = Matrix.create n (succ order) in
  for i=0 to pred n do
    let xi = x.(i) in
    for j=0 to order do
      x_mat.{i, j} <- Math.pow_int xi j
    done
  done ;
  let weight = match weight with
  | None -> None
  | Some a -> Some (vec_convert (`A a)) in
  let (c, cov, chisq) = 
    linear ?weight (`M x_mat) (vec_convert (`A y)) in
  (Vector.to_array c,
   Matrix.to_arrays cov,
   chisq)
