;;; 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 (random-vector-in-sphere rng radius x)
  (void (const gsl-rng*) double gsl-vector*)
  "Generate a random vector distributed uniformly in a sphere of
RADIUS. Store the generated vector in X."
  (random-direction-vector rng x)
  (gsl-vector-scale x (* radius
                         (pow (gsl-ran-flat rng 0 1)
                              (/ (: x size))))))

(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 (planar-angle->solid-angle-fraction theta dimension)
  (double double (unsigned int))
  (let* ((alpha double (* 0.5 (- dimension 1)))
         (beta double 0.5)
         (x double (gsl-pow-2 (sin theta))))
    (cond
     ((or (< theta 0) (> theta M-PI))
      (GSL-ERROR "planar-angle->solid-angle-fraction only allows theta in [0,pi]" GSL-EDOM))
     ;; gsl-sf-beta-inc underflows when x = 0. Work around that. TODO:
     ;; Fix it upstream in gsl.
     ((< (rerror theta M-PI) DBL-EPSILON)
      (return 1))
     ((< theta M-PI-2)
      (return (* 0.5 (gsl-sf-beta-inc alpha beta x))))
     (else
      (return (- 1 (* 0.5 (gsl-sf-beta-inc alpha beta x))))))))

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

  (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-fraction 0) (return 0))
   ((= solid-angle-fraction 1) (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-fraction->planar-angle
                        (gsl-ran-flat r
                                      (planar-angle->solid-angle-fraction theta-min n)
                                      (planar-angle->solid-angle-fraction 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))