;;
;; Multivariate test of normal-pdf.
;;
;; This test creates a 3-dimensional normal distribution with a
;; random mean and variance. It then samples from this distribution
;; and calculates the sample mean and variance for comparison.
;;

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

(define-utility-matrices f64)

;; Number of samples 
(define N 100000)

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


(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 (upper M N A)
  (f64matrix-map blas:RowMajor M N A (lambda (i j v) (if (>= j i) v 0.0))))

;;
;;  Computes the arithmetic mean of a list of f64 vectors using the
;;  recurrence relation mean_(n) = mean(n-1) + (v[n] -
;;  mean(n-1))/(n+1)
;;
(define (mean s vs)
    (let loop ((i 0) (vs vs) (mean (make-f64vector s 0.0)))
      (if (null? vs)  mean
	  (loop (fx+ 1 i) (cdr vs)
		(let ((d (f64vector-sub (car vs) mean)))
		  (f64vector-sum mean (f64vector-scale (/ 1 (+ 1 i)) d)))))))

(define (covariance s vs mean)
   (let* ((cov  (matrix-zeros S S)))
     (let ((ds (map (lambda (x) (f64vector-sub x mean)) vs)))
       (let ((n (fold (lambda (d n) 
			(blas:dsyr! blas:RowMajor blas:Upper s 1.0 d cov)
			(fx+ 1 n)) 0 ds)))
	 (f64vector-scale (/ 1 (- n 1)) cov)))))
			 

   ;; 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 (< lo hi) (let ((d (- hi lo)))
		     (f64vector-map (lambda (x) (+ lo (* d x))) x))))
   
   (define rng  (random-mtzig:init))

   ;;  distribution parameters
   (define mu     (f64vector-urange (random-mtzig:f64vector-randu! S rng) -5 5))
   (define sigma  (let ((x (f64vector-urange (random-mtzig:f64vector-randu! (* S S) rng) 0 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
		   (print "mu = " mu)
		   (print "sigma = " sigma)
		   (make-normal-pdf S mu sigma)))

   (define input (list-tabulate N (lambda (i) (random-mtzig:f64vector-randn! S rng))))

   (define samples
     (let loop ((i 0) (input input) (samples (list)))
       (if (null? input)  samples
	   (let ((samples (cons (normal-pdf:sample gpdf  (car input))  samples)))
	     (loop (fx+ 1 i) (cdr input) samples)))))

   (if (<= N 100) (print "input = " input))
   (if (<= N 100) (print "samples = " samples))
;   (print "density = " (cadr samples+density))
      
(test-group "multivariate normal 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))) 
		       (mean S samples)))

  (test "sample-covariance"
	(f64vector-map (lambda (x) (truncate (* 10 x))) (upper S S sigma))
	(f64vector-map (lambda (x) (truncate (* 10 x)))
		       (covariance S samples (mean S samples))))
  )