;;; 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 "math.h")
(pre-include "gsl/gsl_blas.h")
(pre-include "gsl/gsl_poly.h")
(pre-include "gsl/gsl_randist.h")
(pre-include "oracles.h")
(pre-include "utils.h")

(define (bernoulli-extent-oracle r x -params) (double (const gsl-rng*) (const gsl-vector*) void*)
  (let* ((params (const bernoulli-params*) (convert-type -params bernoulli-params*)))
    (return (if* (gsl-ran-bernoulli r (: params p))
                 (: params r1)
                 (: params r0)))))

(define (bernoulli-true-volume dimension -params) (double (unsigned int) void*)
  (let* ((params (const bernoulli-params*) (convert-type -params bernoulli-params*)))
    (return (* (volume-of-ball dimension)
               (+ (* (: params p) (gsl-pow-uint (: params r1) dimension))
                  (* (- 1 (: params p)) (gsl-pow-uint (: params r0) dimension)))))))

(define (uniform-extent-oracle r x -params) (double (const gsl-rng*) (const gsl-vector*) void*)
  (let* ((params (const uniform-params*) (convert-type -params uniform-params*)))
    (return (gsl-ran-flat r (: params a) (: params b)))))

;; TODO: Verify the accuracy of this function for non-trivial a, b.
(define (uniform-true-volume dimension -params) (double (unsigned int) void*)
  (let* ((params uniform-params* (convert-type -params uniform-params*)))
    (return (* (volume-of-ball dimension)
               (- (/ (gsl-pow-uint (: params b) (+ dimension 1))
                     (+ dimension 1))
                  (/ (gsl-pow-uint (: params a) (+ dimension 1))
                     (+ dimension 1)))))))

(define (beta-extent-oracle r x -params) (double (const gsl-rng*) (const gsl-vector*) void*)
  (let* ((params (const beta-params*) (convert-type -params beta-params*)))
    (return (gsl-ran-beta r (: params alpha) (: params beta)))))

(define (beta-true-volume dimension -params) (double (unsigned int) void*)
  (let* ((params (const beta-params*) (convert-type -params beta-params*))
         (vol double (volume-of-ball dimension)))
    (for-i r dimension
      (set* vol (/ (+ (: params alpha) r)
                   (+ (: params alpha) (: params beta) r))))
    (return vol)))

(define (infinity-norm x) ((static double) (const gsl-vector*))
  (let* ((max double (fabs (gsl-vector-get x 0))))
    ;; TODO: Start this loop from i = 1, not i = 0. That would be
    ;; slightly faster.
    (for-i i (: x size)
      (set max (GSL-MAX max (fabs (gsl-vector-get x i)))))
    (return max)))

(define (cube-extent-oracle r x -params) (double (const gsl-rng*) (const gsl-vector*) void*)
  (let* ((params (const cube-params*) (convert-type -params cube-params*)))
    (return (/ (: params edge) 2 (infinity-norm x)))))

(sc-define-syntax (compute-cube-extent-oracle-minimizand i)
  (/ (- (/ (: params edge) 2)
        (* (GSL-SIGN (gsl-vector-get x i))
           (gsl-vector-get (: params center) i)))
     (fabs (gsl-vector-get x i))))

(define (cube-extent-oracle-with-center r x -params) (double (const gsl-rng*) (const gsl-vector*) void*)
  (let* ((params (const cube-params*) (convert-type -params cube-params*))
         (min double (compute-cube-extent-oracle-minimizand 0)))
    ;; TODO: Start this loop from i = 1, not i = 0. That would be
    ;; slightly faster.
    (for-i i (: (: params center) size)
      (set min (GSL-MIN min (compute-cube-extent-oracle-minimizand i))))
    (return min)))

(define (cube-true-volume dimension -params) (double (unsigned int) void*)
  (let* ((params (const cube-params*) (convert-type -params cube-params*)))
    (return (gsl-pow-uint (: params edge) dimension))))

(define (ellipsoid-extent-oracle r x -params) (double (const gsl-rng*) (const gsl-vector*) void*)
  (let* ((params (const ellipsoid-params*) (convert-type -params ellipsoid-params*))
         (k double 0))
    (for-i i (: (: params axes) size)
      (set+ k (gsl-pow-2 (/ (gsl-vector-get x i)
                            (gsl-vector-get (: params axes) i)))))
    (return (/ (sqrt k)))))

(define (ellipsoid-extent-oracle-with-center r x -params) (double (const gsl-rng*) (const gsl-vector*) void*)
  (let* ((params (const ellipsoid-params*) (convert-type -params ellipsoid-params*))
         ;; a, b and c are the coefficients of a quadratic equation
         ;; ax^2 + bx + c = 0.
         (a double 0)
         (b double 0)
         (c double 0))
    (for-i i (: (: params axes) size)
      (set+ a (gsl-pow-2 (/ (gsl-vector-get x i)
                            (gsl-vector-get (: params axes) i))))
      (set+ b (/ (* (gsl-vector-get x i)
                    (gsl-vector-get (: params center) i))
                 (gsl-pow-2 (gsl-vector-get (: params axes) i))))
      (set+ c (gsl-pow-2 (/ (gsl-vector-get (: params center) i)
                            (gsl-vector-get (: params axes) i)))))
    (set* b 2)
    (set- c 1)
    (declare k1 double
             k2 double)
    (gsl-poly-solve-quadratic a b c (address-of k1) (address-of k2))
    ;; k2 > k1, k1 < 0, k2 > 0. k2 is our desired root.
    (return k2)))

(define (ellipsoid-true-volume dimension -params) (double (unsigned int) void*)
  (let* ((params (const ellipsoid-params*) (convert-type -params ellipsoid-params*))
         (vol double (volume-of-ball (: (: params axes) size))))
    (for-i i (: (: params axes) size)
      (set* vol (gsl-vector-get (: params axes) i)))
    (return vol)))

(define (spheroid-extent-oracle r x -params) (double (const gsl-rng*) (const gsl-vector*) void*)
  (let* ((params (const spheroid-params*) (convert-type -params spheroid-params*))
         (xsub gsl-vector-const-view
               (gsl-vector-const-subvector x 1 (- (: x size) 1))))
    (return (/ (sqrt (+ (gsl-pow-2 (gsl-blas-dnrm2 (address-of (struct-get xsub vector))))
                        (gsl-pow-2 (/ (gsl-vector-get x 0)
                                      (: params eccentricity)))))))))

(define (spheroid-true-volume dimension -params) (double (unsigned int) void*)
  (let* ((params (const spheroid-params*) (convert-type -params spheroid-params*)))
    (return (* (volume-of-ball dimension)
               (: params eccentricity)))))