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

let () = Error.init ()

open Fun

(* PLAIN algorithm *)
type plain_state

external _alloc_plain : int -> plain_state = "ml_gsl_monte_plain_alloc"
external _free_plain : plain_state -> unit = "ml_gsl_monte_plain_free"

let make_plain_state s =
  let state = _alloc_plain s in
  Gc.finalise _free_plain state;
  state

external init_plain : plain_state -> unit = "ml_gsl_monte_plain_init"

external integrate_plain :
  monte_fun ->
  lo:float array ->
  up:float array ->
  int ->
  Rng.t ->
  plain_state ->
  Fun.result = "ml_gsl_monte_plain_integrate_bc" "ml_gsl_monte_plain_integrate"

(* MISER algorithm *)
type miser_state

type miser_params = {
  estimate_frac : float; (* 0.1 *)
  min_calls : int; (* 16 * dim *)
  min_calls_per_bisection : int; (* 32 * min_calls *)
  miser_alpha : float; (* 2. *)
  dither : float; (* 0. *)
}

external _alloc_miser : int -> miser_state = "ml_gsl_monte_miser_alloc"
external _free_miser : miser_state -> unit = "ml_gsl_monte_miser_free"

let make_miser_state s =
  let state = _alloc_miser s in
  Gc.finalise _free_miser state;
  state

external init_miser : miser_state -> unit = "ml_gsl_monte_miser_init"

external integrate_miser :
  monte_fun ->
  lo:float array ->
  up:float array ->
  int ->
  Rng.t ->
  miser_state ->
  Fun.result = "ml_gsl_monte_miser_integrate_bc" "ml_gsl_monte_miser_integrate"

external get_miser_params : miser_state -> miser_params
  = "ml_gsl_monte_miser_get_params"

external set_miser_params : miser_state -> miser_params -> unit
  = "ml_gsl_monte_miser_set_params"

(* VEGAS algorithm *)
type vegas_state
type vegas_info = { result : float; sigma : float; chisq : float }
type vegas_mode = STRATIFIED | IMPORTANCE_ONLY | IMPORTANCE

type vegas_params = {
  vegas_alpha : float; (* 1.5 *)
  iterations : int; (* 5 *)
  stage : int;
  mode : vegas_mode;
  verbose : int; (* 0 *)
  ostream : out_channel option; (* stdout *)
}

external _alloc_vegas : int -> vegas_state = "ml_gsl_monte_vegas_alloc"
external _free_vegas : vegas_state -> unit = "ml_gsl_monte_vegas_free"

let make_vegas_state s =
  let state = _alloc_vegas s in
  Gc.finalise _free_vegas state;
  state

external init_vegas : vegas_state -> unit = "ml_gsl_monte_vegas_init"

external integrate_vegas :
  monte_fun ->
  lo:float array ->
  up:float array ->
  int ->
  Rng.t ->
  vegas_state ->
  Fun.result = "ml_gsl_monte_vegas_integrate_bc" "ml_gsl_monte_vegas_integrate"

external get_vegas_info : vegas_state -> vegas_info
  = "ml_gsl_monte_vegas_get_info"

external get_vegas_params : vegas_state -> vegas_params
  = "ml_gsl_monte_vegas_get_params"

external set_vegas_params : vegas_state -> vegas_params -> unit
  = "ml_gsl_monte_vegas_set_params"

(* High-level version *)
type kind = PLAIN | MISER | VEGAS

let integrate kind f ~lo ~up calls rng =
  let dim = Array.length lo in
  let with_state alloc free integ =
    let state = alloc dim in
    try
      let res = integ f ~lo ~up calls rng state in
      free state;
      res
    with exn ->
      free state;
      raise exn
  in
  match kind with
  | PLAIN -> with_state _alloc_plain _free_plain integrate_plain
  | MISER -> with_state _alloc_miser _free_miser integrate_miser
  | VEGAS -> with_state _alloc_vegas _free_vegas integrate_vegas
