;;
;; Multivariate test of sampled-pdf.
;;
;; This test creates a NormalPdf with a random mean and
;; covariance. It then generates uniform samples within 3 standard
;; deviations of the mean and weights them according to the density
;; function of this NormalPDF, adding them to a SampledPdf
;; distribution. The mean and covariance of the SampledPdf can then be
;; compared with those of the original Gaussian. Furthermore, the
;; SampledPdf is resampled and the new mean and covariance calculated,
;; which should be the same as before the resampling, or very close
;; to.
;;


(use srfi-1 srfi-4 srfi-4-utils matrix-utils blas random-mtzig rb-tree
     normal-pdf sampled-pdf  test)

;; Number of samples 
;(define N 100000)
(define N 10000)

;; Number of dimensions (sample size)
(define S 2)


(define (f64vector-scale a x)
  (let ((n (f64vector-length x)))
    (blas:dscal n a x)))

(define (f64vector-sum x y)
  (let ((n (f64vector-length x)))
    (blas:daxpy n 1.0 x y)))

(define (f64vector-sub x y)
  (let ((n (f64vector-length x)))
    (blas:daxpy n -1.0 y x)))

(define (f64vector-mul a b)
  (f64vector-map (lambda (v1 v2) (* v1 v2)) a b))

(define f64matrix-map (make-matrix-map f64vector-ref f64vector-set!))
(define f64matrix-fold-partial (make-matrix-fold-partial f64vector-ref))

(define (f64vector< x y) 
  (let ((d (f64vector-sub y x)))
    (f64vector-fold (lambda (x ax) (and (< 0 x) ax)) #t d)))

(define (lower M N A)
  (f64matrix-map blas:RowMajor M N A (lambda (i j v) (if (>= i j) v 0.0))))

(define (diag M N A)
  (list->f64vector
   (reverse (f64matrix-fold-partial
	     blas:RowMajor M N A 
	     (lambda (x ax) (cons x ax)) 
	     (lambda (i j) (fx= i j)) (list)))))

;; Given a vector of values in the range [0..1], scale all values in
   ;; the vector to the specified range.
(define (f64vector-urange x lo hi)
  (if (f64vector< lo hi) 
      (let ((d (f64vector-sub hi lo)))
	(f64vector-sum lo (f64vector-mul d x)))))

(define rng      (random-mtzig:init))

;;  distribution parameters
(define mu     (f64vector-urange (random-mtzig:f64vector-randu! S rng) 
				 (make-f64vector S -5) (make-f64vector S 5)))

(define sigma  (let* ((r  (random-mtzig:f64vector-randn! (* S S) rng))
		      (x  (f64vector-urange r (make-f64vector (* S S) 0) (make-f64vector (* S S) 5))))
		 (blas:dgemm! blas:RowMajor blas:NoTrans blas:Trans S S S
			      1.0 x x 0.0 (make-f64vector (* S S)))))

(define gpdf  (begin
		(make-normal-pdf S mu sigma)))

(define var    (diag S S sigma))
(define sdev   (f64vector-map sqrt var))

(define samples+weights
  (let ((hi (f64vector-scale 3 sdev))
	(lo (f64vector-scale -3 sdev)))
    (let loop ((i 0) (samples (list)))
      (if (< i N)
	  (let ((x (f64vector-urange (random-mtzig:f64vector-randu! S rng) lo hi)))
	    (let* ((sample   (normal-pdf:sample gpdf  x))
		   (density  (normal-pdf:density gpdf sample)))
	      (loop (fx+ 1 i) (cons (cons density sample) samples))))
	  samples))))

(define spdf (make-sampled-pdf S make-rb-tree samples+weights))  

(print "given mean: " mu)
(print "given covariance: " sigma)
(print "sampled mean: " (sampled-pdf:expectation spdf))
(print "sampled covariance: " (sampled-pdf:covariance spdf))

      
(test-group  "multivariate sampled PDF test"

	     ;; test for mean and covariance equality to one decimal place
	     (test "sample-mean"
		   (f64vector-map (lambda (x) (truncate (* 10 x))) mu)
		   (f64vector-map (lambda (x) (truncate (* 10 x))) (sampled-pdf:expectation spdf)))


	     (test "sample-covariance"
		   (f64vector-map (lambda (x) (truncate (* 10 x))) sigma)
		   (f64vector-map (lambda (x) (truncate (* 10 x))) (sampled-pdf:covariance spdf))) )