;;; 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/>.

(sc-include "macros/macros")

(pre-include "stddef.h")
(pre-include "gsl/gsl_integration.h")
(pre-include "gsl/gsl_rstat.h")
(pre-include "gsl/gsl_vector.h")
(pre-include "nd-random.h")
(pre-include "utils.h")
(pre-include "extent-sampling.h")

(sc-define-syntax (with-integration-workspace var intervals body ...)
  (with-alloc var gsl-integration-workspace*
              (gsl-integration-workspace-alloc intervals)
              gsl-integration-workspace-free
              body ...))

(sc-define-syntax (invoke-extent-oracle extent-oracle r x)
  ((: extent-oracle oracle) r x (: extent-oracle params)))

(pre-let* (WINDOW-LENGTH 1000)
  (define (volume extent-oracle true-volume r dimension rtol stats)
    (void (const extent-oracle-t*) double (const gsl-rng*) (unsigned int) double gsl-rstat-workspace*)
    (define accurate-estimates int 0)
    (let* ((vn double (ln-volume-of-ball dimension)))
      (with-vector x dimension
        (do-while (< accurate-estimates WINDOW-LENGTH)
          (random-direction-vector r x)
          (gsl-rstat-add (exp (+ vn (* dimension (log (invoke-extent-oracle extent-oracle r x)))))
                         stats)
          (cond
           ((rtol? (gsl-rstat-mean stats) true-volume rtol)
            (set+ accurate-estimates 1))
           (else
            (set accurate-estimates 0))))))))

(sc-define-syntax (invoke-integrand integrand r x)
  ((: integrand integrand) r x (: integrand params)))

(pre-let* (INTEGRATION-INTERVALS 1000)
  (define (integral-per-direction integrand direction n radius rtol)
    (double (const integrand-t*) (const gsl-vector*) (unsigned int) double double)
    (define (f r params) (double double void*)
      (return (* (gsl-pow-uint r (- n 1))
                 (invoke-integrand integrand r direction))))
    (declare gsl-f (struct-variable gsl-function (function (address-of f)))
             result double
             error double)
    (with-integration-workspace w INTEGRATION-INTERVALS
      (gsl-integration-qag (address-of gsl-f)
                           0
                           radius
                           0
                           rtol
                           INTEGRATION-INTERVALS
                           GSL-INTEG-GAUSS15
                           w
                           (address-of result)
                           (address-of error)))
    (return (* (surface-area-of-ball n)
               result))))

(pre-let* (WINDOW-LENGTH 1000)
  (define (integral integrand extent-oracle true-integral r dimension rtol stats)
    (void (const integrand-t*) (const extent-oracle-t*) double (const gsl-rng*) (unsigned int) double gsl-rstat-workspace*)
    (define accurate-estimates int 0)
    (with-vector x dimension
      (do-while (< accurate-estimates WINDOW-LENGTH)
        (random-direction-vector r x)
        (gsl-rstat-add (integral-per-direction integrand x dimension
                                               (invoke-extent-oracle extent-oracle r x) rtol)
                       stats)
        (cond
         ((rtol? (gsl-rstat-mean stats) true-integral rtol)
          (set+ accurate-estimates 1))
         (else
          (set accurate-estimates 0)))))))

(define (volume-cone extent-oracle r mean omega-min omega-max number-of-samples variance)
  (double (const extent-oracle-t*) (const gsl-rng*) (const gsl-vector*) double double (unsigned int) double*)
  (declare volume double)
  (let* ((dimension (unsigned int) (: mean size))
         (vn double (ln-volume-of-ball dimension))
         (theta-min double (solid-angle-fraction->planar-angle omega-min dimension))
         (theta-max double (solid-angle-fraction->planar-angle omega-max dimension)))
    (with-rstats stats
      (with-vector x dimension
        (for-i i number-of-samples
          (hollow-cone-random-vector r mean theta-min theta-max x)
          (gsl-rstat-add (exp (+ vn (* dimension (log (invoke-extent-oracle extent-oracle r x)))))
                         stats)))
      (cond
       (variance (set (pointer-get variance) (gsl-rstat-variance stats))))
      (set volume (* (gsl-rstat-mean stats)
                     (- omega-max omega-min)))))
  (return volume))

(define (volume-importance extent-oracle r mean samples-per-cone solid-angle-factor
                           solid-angle-threshold-exponent-factor number-of-samples)
  (double (const extent-oracle-t*) (const gsl-rng*) (const gsl-vector*)
          (unsigned int) double double (unsigned int*))
  (define
    volume double 0
    omega-max double 0.5
    N int 0)
  (let* ((dimension int (: mean size)))
    (do-while (> omega-max (pow 2 (- (* dimension solid-angle-threshold-exponent-factor))))
      (let* ((omega-min double (/ omega-max solid-angle-factor)))
        (set+ volume (volume-cone extent-oracle r mean omega-min omega-max samples-per-cone NULL))
        (set+ N samples-per-cone)
        (set omega-max omega-min))))
  (set+ volume (volume-cone extent-oracle r mean 0 omega-max samples-per-cone NULL))
  (set+ N samples-per-cone)
  (cond
   (number-of-samples (set (pointer-get number-of-samples) N)))
  (return volume))