(sc-include "macros/macros")

(pre-include "gsl/gsl_blas.h")
(pre-include "gsl/gsl_cblas.h")
(pre-include "gsl/gsl_eigen.h")
(pre-include "gsl/gsl_matrix.h")
(pre-include "extent-sampling.h")
(pre-include "nd-random.h")
(pre-include "oracles.h")
(pre-include "utils.h")

(pre-define ECCENTRICITY 2)
(pre-define DIMENSION 2)

(pre-define SAMPLES 100)

(declare covariance-accumulator
         (type (struct (square gsl-matrix*)
                       (sum gsl-vector*)
                       (n int))))

(define (covariance-accumulator-alloc dimension) (covariance-accumulator int)
  (declare accumulator covariance-accumulator)
  (set (struct-get accumulator square) (gsl-matrix-alloc dimension dimension))
  (set (struct-get accumulator sum) (gsl-vector-alloc dimension))
  (set (struct-get accumulator n) 0)
  (return accumulator))

(define (covariance-accumulator-add accumulator sample) (void covariance-accumulator* gsl-vector*)
  (define square gsl-matrix* (: accumulator square))
  (cblas-dsyr CblasRowMajor CblasLower (: sample size) 1
              (: sample data) (: sample stride)
              (: square data) (: square tda))
  (gsl-vector-add (: accumulator sum) sample)
  (set+ (: accumulator n) 1))

(define (covariance-accumulator-covariance accumulator result) (void covariance-accumulator gsl-matrix*)
  (let* ((dim int (: (struct-get accumulator sum) size))
         (n int (struct-get accumulator n))
         (sum gsl-vector* (struct-get accumulator sum))
         (square gsl-matrix* (struct-get accumulator square)))
    (gsl-matrix-memcpy result square)
    (gsl-matrix-scale result (/ 1.0 (- n 1)))
    (cblas-dsyr CblasRowMajor CblasLower dim (/ -1.0 n n)
                (: sum data) (: sum stride)
                (: result data) (: result tda))))

(define (covariance-accumulator-free accumulator) (void covariance-accumulator)
  (gsl-matrix-free (struct-get accumulator square))
  (gsl-vector-free (struct-get accumulator sum)))

(sc-define-syntax (with-covariance-accumulator var dimension body ...)
  (with-alloc var covariance-accumulator
              (covariance-accumulator-alloc dimension)
              covariance-accumulator-free
              body ...))

(define (matrix-vector-multiply A x y)
  (void (const gsl-matrix*) (const gsl-vector*) gsl-vector*)
  (cblas-dgemv CblasRowMajor CblasNoTrans (: A size1) (: A size2)
               1 (: A data) (: A tda)
               (: x data) (: x stride) 0 (: y data) (: y stride)))

(sc-define-syntax (with-eigen-symmv workspace dimension body ...)
  (with-alloc workspace gsl-eigen-symmv-workspace*
              (gsl-eigen-symmv-alloc dimension)
              gsl-eigen-symmv-free
              body ...))

(define (eigen-symmetric A eval evec) (void gsl-matrix* gsl-vector* gsl-matrix*)
  (with-eigen-symmv workspace (: A size1)
    (gsl-eigen-symmv A eval evec workspace)))

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

(define (print-vector x) (void gsl-vector*)
  (for-i i (: x size)
    (printf "%g\t" (gsl-vector-get x i)))
  (printf "\n"))

(define (print-matrix A) (void gsl-matrix*)
  (for-i i (: A size1)
    (for-i j (: A size2)
      (printf "%g\t" (gsl-matrix-get A i j)))
    (printf "\n")))

(pre-let* (COVARIANCE-SAMPLES 10000)
  (define (volume-whitening extent-oracle rng dimension stats)
    (void (const extent-oracle-t*) (const gsl-rng*) (unsigned int) gsl-rstat-workspace*)
    (let* ((vn double (ln-volume-of-ball dimension)))
      (with-square-matrix transform dimension
        ;; (with-covariance-accumulator accumulator dimension
        ;;   (with-vector x dimension
        ;;     (for-i i COVARIANCE-SAMPLES
        ;;       (random-direction-vector rng x)
        ;;       (gsl-vector-scale x (invoke-extent-oracle extent-oracle rng x))
        ;;       (covariance-accumulator-add (address-of accumulator) x)))
        ;;   (covariance-accumulator-covariance accumulator transform))
        ;; (with-vector eval dimension
        ;;   (with-square-matrix evec dimension
        ;;     (eigen-symmetric transform eval evec)
        ;;     (print-vector eval)
        ;;     (printf "\n")
        ;;     (print-matrix evec)))
        ;; TODO: Replace hardcoded eccentricity.
        (gsl-matrix-set-identity transform)
        (gsl-matrix-set transform 0 0 (/ 1.0 ECCENTRICITY))
        ;; find surface vector as usual using the extent function
        ;; transform this surface vector to a virtual body
        ;; find the volume of this virtual body as usual
        ;; use this volume to find the volume of the real body
        ;;
        ;; The idea is to transform the real body to a virtual body
        ;; whose volume can be estimated very accurately. Then, scale
        ;; up the virtual volume to get the real volume. Hopefully,
        ;; the enlargement of error will be small.

        ;; So, the simple determinant relation between the volumes of
        ;; the real and virtual bodies is incorrect. What then is
        ;; correct?
        (with-vector x dimension
          (with-vector y dimension
            (for-i i SAMPLES
              (random-direction-vector rng x)
              (gsl-vector-scale x (invoke-extent-oracle extent-oracle rng x))
              (matrix-vector-multiply transform x y)
              (gsl-rstat-add (exp (+ vn (* dimension (log (gsl-blas-dnrm2 y)))))
                             stats)
              (printf "%g (%g/%g)\t" (gsl-rstat-mean stats)
                      (gsl-blas-dnrm2 x)
                      (gsl-blas-dnrm2 y)))))))
    (printf "\n")))

(define (main) (int)
  (declare params (struct-variable spheroid-params
                                   (eccentricity ECCENTRICITY)))
  (declare oracle (struct-variable extent-oracle-t
                                   (oracle spheroid-extent-oracle)
                                   (params (address-of params))))
  (printf "True volume: %g\n" (spheroid-true-volume DIMENSION (address-of params)))
  (with-rng rng
    (with-rstats stats
      (volume-whitening (address-of oracle) rng DIMENSION stats)))
  (return 0))