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

(pre-include "math.h")
(pre-include "gsl/gsl_blas.h")
(pre-include "gsl/gsl_randist.h")
(pre-include "gsl/gsl_sf_gamma.h")
(pre-include "gsl/gsl_vector.h")
(pre-include "utils.h")

(define (random-direction-vector r x)
  (void (const gsl-rng*) gsl-vector*)
  "Generate a random vector distributed uniformly on the unit
sphere. Write the result to X."
  (gsl-ran-dir-nd r (: x size) (: x data)))

(define (rotate-from-nth-canonical x orient)
  ((static void) gsl-vector* (const gsl-vector*))
  (let* ((n size-t (: x size))
         (mun double (gsl-vector-get orient (- n 1))))
    ;; If the orient vector is already the nth canonical axis, do
    ;; nothing.
    (cond
     ((not (= mun 1))
      (let* ((xn double (gsl-vector-get x (- n 1)))
             (orient-sub gsl-vector-const-view
                         (gsl-vector-const-subvector orient 0 (- n 1)))
             (b double (gsl-blas-dnrm2 (address-of (struct-get orient-sub vector))))
             (a double (/ (- (dot-product orient x)
                             (* xn mun))
                          b))
             (s double (sqrt (- 1 (gsl-pow-2 mun)))))
        (gsl-blas-daxpy (/ (+ (* xn s)
                              (* a (- mun 1)))
                           b)
                        orient
                        x)
        (gsl-vector-set x
                        (- n 1)
                        (+ (gsl-vector-get x (- n 1))
                           (* xn (- mun 1))
                           (- (* a s))
                           (- (/ (* mun (+ (* xn s)
                                           (* a (- mun 1))))
                                 b)))))))))

(define (beta-inc-unnormalized a b x) ((static double) double double double)
  (return (* (gsl-sf-beta-inc a b x)
             (gsl-sf-beta a b))))

(define (incomplete-wallis-integral theta m) ((static double) double (unsigned int))
  "Return the incomplete Wallis integral \\int_0^\\theta \\sin^m x
dx. THETA should be in [0,pi]."
  (cond
   ((or (< theta 0) (> theta M-PI))
    (GSL-ERROR "Incomplete Wallis integral only allows theta in [0,pi]" GSL-EDOM))
   ((< theta M-PI-2)
    (return (/ (beta-inc-unnormalized (* 0.5 (+ m 1)) 0.5 (gsl-pow-2 (sin theta)))
               2)))
   (else
    (return (/ (+ (gsl-sf-beta (* 0.5 (+ m 1)) 0.5)
                  (beta-inc-unnormalized 0.5 (* 0.5 (+ m 1)) (gsl-pow-2 (cos theta))))
               2)))))

(define (planar-angle->solid-angle planar-angle dimension) (double double (unsigned int))
  (return (/ (* 2
                (pow M-PI (* 0.5 (- dimension 1)))
                (incomplete-wallis-integral planar-angle (- dimension 2)))
             (gsl-sf-gamma (* 0.5 (- dimension 1))))))

(define (solid-angle->planar-angle solid-angle dimension) (double double (unsigned int))
  (define (f planar-angle params) (double double void*)
    (return (- (planar-angle->solid-angle planar-angle dimension)
               solid-angle)))

  (declare gsl-f (struct-variable gsl-function (function (address-of f))))
  ;; TODO: Equality comparisons to floating point values may be
  ;; problematic. Fix it.
  (cond
   ((= solid-angle 0) (return 0))
   ((= solid-angle (surface-area-of-ball dimension)) (return M-PI))
   (else (return (bisection (address-of gsl-f) 0 M-PI)))))

(define (hollow-cone-random-vector r mean theta-min theta-max x)
  (void (const gsl-rng*) (const gsl-vector*) double double gsl-vector*)
  ;; Generate random vector around the nth canonical basis vector.
  (let* ((n size-t (: x size))
         (theta double (solid-angle->planar-angle
                        (gsl-ran-flat r
                                      (planar-angle->solid-angle theta-min n)
                                      (planar-angle->solid-angle theta-max n))
                        n)))
    (gsl-ran-dir-nd r (- n 1) (: x data))
    (gsl-vector-scale x (sin theta))
    (gsl-vector-set x (- n 1) (cos theta)))

  ;; Rotate to arbitrary basis.
  (rotate-from-nth-canonical x mean))

(define (cone-random-vector r mean theta-max x)
  (void (const gsl-rng*) (const gsl-vector*) double gsl-vector*)
  (hollow-cone-random-vector r mean 0 theta-max x))