(sc-include "macros/macros")

(pre-include "gsl/gsl_poly.h")
(pre-include "gsl/gsl_sf.h")
(pre-include "extent-sampling.h")
(pre-include "oracles.h")
(pre-include "integrands.h")
(pre-include "utils.h")

(pre-define TRIALS 1000)
(pre-define MAX-EXTENT 1.0)

(sc-define-syntax* (define-polynomial var)
  (let ((coefficients (list -0.09375 0.6875 -1.5 1))
        (coeff (sc-gensym)))
    `(begin
       (declare ,coeff (array double ,(length coefficients)))
       (array-set* ,coeff ,@coefficients)
       (declare ,var (struct-variable polynomial-integrand-params
                                      (coefficients ,coeff)
                                      (degree ,(1- (length coefficients))))))))

(sc-define-syntax (polynomial-map! polynomial i coeff form)
  (for-i i (+ 1 (struct-get polynomial degree))
    (let* ((coeff double (array-get (struct-get polynomial coefficients)
                                    i)))
      (array-set (convert-type (struct-get polynomial coefficients)
                               double*)
                 i form))))

(define (polyval polynomial value) (double polynomial-integrand-params double)
  (return (gsl-poly-eval (struct-get polynomial coefficients)
                         (+ 1 (struct-get polynomial degree))
                         value)))

(sc-define-syntax* (define-polynomial-integrand integrand)
  (let ((params (sc-gensym)))
    `(begin
       (define-polynomial ,params)
       (declare ,integrand (struct-variable (const integrand-t)
                                            (integrand polynomial-integrand)
                                            (params (address-of ,params)))))))

(sc-define-syntax* (define-uniform-extent-oracle oracle a b)
  (let ((oracle-params (sc-gensym)))
    `(begin
       (declare ,oracle-params (struct-variable uniform-params
                                                (a ,a)
                                                (b ,b)))
       (declare ,oracle (struct-variable extent-oracle-t
                                         (oracle uniform-extent-oracle)
                                         (params (address-of ,oracle-params)))))))

(sc-define-syntax* (experiment integrand true-integral filename)
  (let ((rtols (list 0.2 0.1 0.05)))
    `(begin
       ,@(map (lambda (i rtol)
                `(begin
                   (with-data-file fp ,(format #f "integral-~a-~a.dat" filename rtol) "w"
                     (fprintf fp "dimension\tsamples\n"))))
              (iota (length rtols))
              rtols)
       (for-i-step dimension 10 100 10
         (printf "Dimension %d: " dimension)
         (with-rstats* samples ,(length rtols)
           (for-i trial TRIALS
             (with-rstats stats
               ,@(map (lambda (i rtol)
                        `(begin
                           (integral (address-of ,integrand)
                                     (address-of oracle)
                                     (,true-integral dimension)
                                     rng
                                     dimension
                                     ,rtol
                                     stats)
                           (gsl-rstat-add (gsl-rstat-n stats)
                                          (array-get samples ,i))))
                      (iota (length rtols))
                      rtols))
             (when (= (modulo trial (/ TRIALS 10)) 0)
               (printf "%d " trial))
             (fflush NULL))
           (printf "\n")
           ,@(map (lambda (i rtol)
                    `(with-data-file fp ,(format #f "integral-~a-~a.dat" filename rtol) "a"
                       (fprintf fp "%d\t%g\n" dimension (gsl-rstat-mean (array-get samples ,i)))))
                  (iota (length rtols))
                  rtols))))))

(define (true-polynomial-integral dimension)
  (double (unsigned int))
  (define-polynomial polynomial)
  (polynomial-map! polynomial k coeff
    (/ coeff (+ dimension k) (+ dimension k 1)))
  (return (* (surface-area-of-ball dimension)
             (gsl-pow-uint MAX-EXTENT dimension)
             (polyval polynomial MAX-EXTENT))))

(define (lower-incomplete-gamma s x)
  (double double double)
  (return (* (gsl-sf-gamma s)
             (gsl-sf-gamma-inc-P s x))))

(define (true-gaussian-integral dimension)
  (double (unsigned int))
  (return (/ (* (pow 2.0 (- (* dimension 0.5) 1))
                (surface-area-of-ball dimension)
                (- (* MAX-EXTENT
                      (lower-incomplete-gamma (* 0.5 dimension)
                                              (* 0.5 (gsl-pow-2 MAX-EXTENT))))
                   (* (sqrt 2.0)
                      (lower-incomplete-gamma (* 0.5 (+ dimension 1))
                                              (* 0.5 (gsl-pow-2 MAX-EXTENT))))))
             MAX-EXTENT)))

(define (true-x-coordinate-integral dimension)
  (double (unsigned int))
  (return (/ (* (gsl-pow-uint MAX-EXTENT (1+ dimension))
                (surface-area-of-ball (+ dimension 3)))
             2.0
             (pow M-PI 2)
             (+ dimension 2))))

(define (main) (int)
  (define-polynomial-integrand -polynomial-integrand)
  (declare -gaussian-integrand (struct-variable (const integrand-t)
                                                gaussian-integrand
                                                NULL))
  (declare -x-coordinate-integrand (struct-variable (const integrand-t)
                                                    x-coordinate-integrand
                                                    NULL))
  (define-uniform-extent-oracle oracle 0 MAX-EXTENT)
  (with-rng rng
    (experiment -polynomial-integrand true-polynomial-integral "polynomial")
    (experiment -gaussian-integrand true-gaussian-integral "gaussian")
    (experiment -x-coordinate-integrand true-x-coordinate-integral "x-coordinate"))
  (with-data-file fp "integral-parameters.tex" "w"
    (fprintf fp "\\\\newcommand{\\\\integraltrials}{%d}\n" TRIALS))
  (return 0))