;;
;; 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.
;;

(require-extension srfi-1)
(require-extension srfi-4)
(require-extension srfi-4-utils)
(require-extension matrix-utils)
(require-extension blas)
(require-extension normal-pdf)
(require-extension random-mtzig)

(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))

;; test for mean and covariance equality to one decimal place
(equal? (f64vector-map (lambda (x) (truncate (* 10 x))) mu)
	(f64vector-map (lambda (x) (truncate (* 10 x))) 
		       (mean S samples)))

(equal? (f64vector-map (lambda (x) (truncate (* 10 x))) (upper S S sigma))
	(f64vector-map (lambda (x) (truncate (* 10 x)))
		       (covariance S samples (mean S samples))))