;;; nsmc --- n-sphere Monte Carlo method
;;; Copyright © 2021 Arun I <arunisaac@systemreboot.net>
;;; Copyright © 2021 Murugesan Venkatapathi <murugesh@iisc.ac.in>
;;;
;;; This file is part of nsmc.
;;;
;;; nsmc is free software: you can redistribute it and/or modify it
;;; under the terms of the GNU General Public License as published by
;;; the Free Software Foundation, either version 3 of the License, or
;;; (at your option) any later version.
;;;
;;; nsmc is distributed in the hope that it will be useful, but
;;; WITHOUT ANY WARRANTY; without even the implied warranty of
;;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
;;; General Public License for more details.
;;;
;;; You should have received a copy of the GNU General Public License
;;; along with nsmc.  If not, see <https://www.gnu.org/licenses/>.

(define-module (nsmc wrap)
  #:use-module (srfi srfi-1)
  #:use-module (srfi srfi-26)
  #:use-module (ice-9 match)
  #:use-module (system foreign))

(include "load-libs.scm")

(define (make-gsl-alloc allocator-name allocator-args freer-name)
  (lambda args
    (let ((obj
           (apply (pointer->procedure
                   '* (dynamic-func allocator-name lib-gsl) allocator-args)
                  args)))
      (set-pointer-finalizer! obj (dynamic-func freer-name lib-gsl))
      obj)))

(define-public pi
  ((pointer->procedure double
                       (dynamic-func "pi" lib-nsmc)
                       (list))))

(define-public volume-of-ball
  (pointer->procedure double
                      (dynamic-func "volume_of_ball" lib-nsmc)
                      (list int)))

(define-public surface-area-of-ball
  (pointer->procedure double
                      (dynamic-func "surface_area_of_ball" lib-nsmc)
                      (list int)))

(define-public (log-gamma x)
  ((pointer->procedure double
                       (dynamic-func "gsl_sf_lngamma" lib-gsl)
                       (list double))
   x))

(define-public (lower-incomplete-gamma s x)
  (* ((pointer->procedure double
                          (dynamic-func "gsl_sf_gamma" lib-gsl)
                          (list double))
      s)
     ((pointer->procedure double
                          (dynamic-func "gsl_sf_gamma_inc_P" lib-gsl)
                          (list double double))
      s x)))

(define-public angle-between-vectors
  (pointer->procedure double
                      (dynamic-func "angle_between_vectors" lib-nsmc)
                      (list '* '*)))

(define-public planar-angle->solid-angle-fraction
  (pointer->procedure double
                      (dynamic-func "planar_angle_to_solid_angle_fraction" lib-nsmc)
                      (list double int)))

;; Random state

(define %gsl-random-state
  ((make-gsl-alloc "gsl_rng_alloc" (list '*) "gsl_rng_free")
   (dereference-pointer
    (dynamic-pointer "gsl_rng_default" lib-gsl))))

(define-public (set-gsl-random-state! seed)
  ((pointer->procedure void
                       (dynamic-func "gsl_rng_set" lib-gsl)
                       (list '* unsigned-long))
   %gsl-random-state seed))

(define-public random-flat
  (cut (pointer->procedure double
                           (dynamic-func "gsl_ran_flat" lib-gsl)
                           (list '* double double))
    %gsl-random-state <> <>))

;; Vector functions

(define vector-alloc
  (make-gsl-alloc "gsl_vector_alloc" (list int) "gsl_vector_free"))

(define-public (vector-size vector)
  (match (parse-c-struct vector
                         (list size_t size_t '* '* int))
    ((size _ _ _ _) size)))

(define-public (vector-scale! vector k)
  "Scale VECTOR by a factor K."
  ((pointer->procedure int
                       (dynamic-func "gsl_vector_scale" lib-gsl)
                       (list '* double))
   vector k)
  vector)

(define-public (vector-norm2 vector)
  ((pointer->procedure double
                       (dynamic-func "gsl_blas_dnrm2" lib-gsl)
                       (list '*))
   vector))

(define-public (vector-add! u v)
  "Return the vector sum of vectors U and V. The sum is stored in U."
  ((pointer->procedure int
                       (dynamic-func "gsl_vector_add" lib-gsl)
                       (list '* '*))
   u v)
  u)

(define-public (basis n dimension)
  "Return the basis vector in the Nth canonical direction."
  (let ((vector (vector-alloc dimension)))
    ((pointer->procedure void
                         (dynamic-func "gsl_vector_set_basis" lib-gsl)
                         (list '* int))
     vector n)
    vector))

;; Histogram functions

(define histogram-struct
  (list size_t '* '*))

(define (histogram-alloc bins min max)
  (let ((histogram
         ((make-gsl-alloc "gsl_histogram_alloc" (list size_t) "gsl_histogram_free")
          bins)))
    ((pointer->procedure int
                         (dynamic-func "gsl_histogram_set_ranges_uniform" lib-gsl)
                         (list '* double double))
     histogram min max)
    histogram))

(define-public (histogram->points histogram)
  (match (parse-c-struct histogram histogram-struct)
    ((n range bin)
     (let ((range (array->list (pointer->bytevector range (1+ n) 0 'f64))))
       (map cons
            (drop-right range 1)
            (array->list (pointer->bytevector bin n 0 'f64)))))))

(define-public (histogram->pdf histogram)
  (match (parse-c-struct histogram histogram-struct)
    ((n range bin)
     (let ((range (pointer->bytevector range (1+ n) 0 'f64)))
       ((pointer->procedure int
                            (dynamic-func "gsl_histogram_scale" lib-nsmc)
                            (list '* double))
        histogram
        (/ n
           (- (array-ref range n)
              (array-ref range 0))
           (fold + 0 (array->list (pointer->bytevector bin n 0 'f64)))))
       histogram))))

(define-public (rhistogram bins min max)
  (case-lambda
    (() (histogram-alloc bins min max))
    ((histogram) histogram)
    ((histogram input)
     ((pointer->procedure int
                          (dynamic-func "gsl_histogram_increment" lib-nsmc)
                          (list '* double))
      histogram input)
     histogram)))

;; Running statistics

(define-public rstat-alloc
  (make-gsl-alloc "gsl_rstat_alloc" (list) "gsl_rstat_free"))

(define-public rstat-n
  (pointer->procedure size_t
                      (dynamic-func "gsl_rstat_n" lib-nsmc)
                      (list '*)))

;; Polynomial functions

(define-public (polyval coefficients x)
  ((pointer->procedure double
                       (dynamic-func "gsl_poly_eval" lib-gsl)
                       (list '* int double))
   (bytevector->pointer (list->typed-array 'f64 1 coefficients))
   (length coefficients)
   x))

;; nd-random

(define-public (random-direction-vector dimension)
  (let ((vector (vector-alloc dimension)))
    ((pointer->procedure void
                         (dynamic-func "random_direction_vector" lib-nsmc)
                         (list '* '*))
     %gsl-random-state vector)
    vector))

(define-public (cone-random-vector mean max-theta)
  (let ((vector (vector-alloc (vector-size mean))))
    ((pointer->procedure void
                         (dynamic-func "cone_random_vector" lib-nsmc)
                         (list '* '* double '*))
     %gsl-random-state mean max-theta vector)
    vector))

(define-public (shifted-gaussian-random-vector mean max-theta standard-deviation)
  (let* ((vector (vector-alloc (vector-size mean)))
         (cost
          ((pointer->procedure int
                               (dynamic-func "shifted_gaussian_random_vector" lib-nsmc)
                               (list '* '* double double '*))
           %gsl-random-state mean max-theta standard-deviation vector)))
    vector))

(define-public (shifted-gaussian-random-vector-cost mean max-theta standard-deviation)
  ((pointer->procedure int
                         (dynamic-func "shifted_gaussian_random_vector" lib-nsmc)
                         (list '* '* double double '*))
   %gsl-random-state mean max-theta standard-deviation (vector-alloc (vector-size mean))))

(define %integration-workspace
  ((make-gsl-alloc "gsl_integration_workspace_alloc" (list size_t) "gsl_integration_workspace_free")
   1000))

(define-public (shifted-gaussian-pdf theta mean max-theta standard-deviation)
  ((pointer->procedure double
                       (dynamic-func "shifted_gaussian_pdf" lib-nsmc)
                       (list double double double double unsigned-int '*))
   theta (vector-norm2 mean) max-theta standard-deviation (vector-size mean) %integration-workspace))

;; oracles

(define (make-extent-oracle oracle params)
  (make-c-struct (list '* '*)
                 (list oracle params)))

(define (make-bernoulli-params p r0 r1)
  (make-c-struct (list double double double)
                 (list p r0 r1)))

(define-public (make-bernoulli-oracle p r0 r1)
  (make-extent-oracle (dynamic-func "bernoulli_extent_oracle" lib-nsmc)
                      (make-bernoulli-params p r0 r1)))

(define (true-volume-procedure name)
  (pointer->procedure double
                      (dynamic-func name lib-nsmc)
                      (list unsigned-int '*)))

(define-public (bernoulli-true-volume p r0 r1 dimension)
  ((true-volume-procedure "bernoulli_true_volume")
   dimension (make-bernoulli-params p r0 r1)))

(define (make-uniform-params a b)
  (make-c-struct (list double double)
                 (list a b)))

(define-public (make-uniform-oracle a b)
  (make-extent-oracle (dynamic-func "uniform_extent_oracle" lib-nsmc)
                      (make-uniform-params a b)))

(define-public (uniform-true-volume a b dimension)
  ((true-volume-procedure "uniform_true_volume")
   dimension (make-uniform-params a b)))

(define (make-beta-params alpha beta)
  (make-c-struct (list double double)
                 (list alpha beta)))

(define-public (make-beta-oracle alpha beta)
  (make-extent-oracle (dynamic-func "beta_extent_oracle" lib-nsmc)
                      (make-beta-params alpha beta)))

(define-public (beta-true-volume alpha beta dimension)
  ((true-volume-procedure "beta_true_volume")
   dimension (make-beta-params alpha beta)))

(define (make-cube-params edge)
  (make-c-struct (list double) (list edge)))

(define-public (make-cube-oracle edge)
  (make-extent-oracle (dynamic-func "cube_extent_oracle" lib-nsmc)
                      (make-cube-params edge)))

(define-public (cube-true-volume edge dimension)
  ((true-volume-procedure "cube_true_volume")
   dimension (make-cube-params edge)))

(define (make-spheroid-params eccentricity)
  (make-c-struct (list double) (list eccentricity)))

(define-public (make-spheroid-oracle eccentricity)
  (make-extent-oracle (dynamic-func "spheroid_extent_oracle" lib-nsmc)
                      (make-spheroid-params eccentricity)))

(define-public (spheroid-true-volume eccentricity dimension)
  ((true-volume-procedure "spheroid_true_volume")
   dimension (make-spheroid-params eccentricity)))

(define (make-ellipsoid-params axes)
  (make-c-struct (list '*) (list axes)))

(define-public (make-ellipsoid-oracle axes)
  (make-extent-oracle (dynamic-func "ellipsoid_extent_oracle" lib-nsmc)
                      (make-ellipsoid-params axes)))

(define-public (ellipsoid-true-volume axes)
  ((true-volume-procedure "ellipsoid_true_volume")
   (vector-size axes) (make-ellipsoid-params axes)))

;; integrands

(define (make-integrand integrand params)
  (make-c-struct (list '* '*)
                 (list integrand params)))

(define-public (make-polynomial-integrand polynomial)
  (make-integrand (dynamic-func "polynomial_integrand" lib-nsmc)
                  (make-c-struct (list '* int)
                                 (list (bytevector->pointer
                                        (list->typed-array 'f64 1 polynomial))
                                       (1- (length polynomial))))))

(define-public gaussian-integrand
  (make-integrand (dynamic-func "gaussian_integrand" lib-nsmc)
                  %null-pointer))

(define-public x-coordinate-integrand
  (make-integrand (dynamic-func "x_coordinate_integrand" lib-nsmc)
                  %null-pointer))

;; extent-sampling

(define maybe-procedure->extent-oracle
  (match-lambda
    ((? procedure? proc)
     (make-extent-oracle
      (procedure->pointer double
                          (lambda (r x params)
                            (proc r x))
                          (list '* '* '*))
      %null-pointer))
    (extent-oracle extent-oracle)))

(define maybe-procedure->integrand
  (match-lambda
    ((?  procedure? integrand)
     (make-integrand
      (procedure->pointer double
                          (lambda (r x params)
                            (integrand r x))
                          (list double '* '*))
      %null-pointer))
    (integrand integrand)))

(define-public (volume extent-oracle true-volume dimension rtol stats)
  ((pointer->procedure void
                       (dynamic-func "volume" lib-nsmc)
                       (list '* double '* unsigned-int double '*))
   (maybe-procedure->extent-oracle extent-oracle)
   true-volume %gsl-random-state dimension rtol stats))

(define-public (integral integrand extent-oracle true-integral dimension rtol stats)
  ((pointer->procedure void
                       (dynamic-func "integral" lib-nsmc)
                       (list '* '* double '* unsigned-int double '*))
   (maybe-procedure->integrand integrand)
   (maybe-procedure->extent-oracle extent-oracle)
   true-integral %gsl-random-state dimension rtol stats))

(define-public (volume-importance extent-oracle mean samples-per-cone
                                  solid-angle-factor solid-angle-threshold-exponent-factor)
  (let ((samples (make-c-struct (list unsigned-int) (list 0))))
    (values
     ((pointer->procedure double
                          (dynamic-func "volume_importance" lib-nsmc)
                          (list '* '* '* unsigned-int double double '*))
      (maybe-procedure->extent-oracle extent-oracle)
      %gsl-random-state mean samples-per-cone solid-angle-factor
      solid-angle-threshold-exponent-factor samples)
     (match (parse-c-struct samples (list unsigned-int))
       ((samples) samples)))))