;; Unit tests for kd-tree


(use typeclass kd-tree test srfi-1 srfi-4-utils random-mtzig)



(define=> (make-consistency-tests <Point> <KdTree>)
  (lambda (pts t x k r)

    (define (sort-points pts) 
      (sort pts (lambda (a b) (negative? (compare-distance x a b)))))
    (define sorted-points (sort-points pts))
    (define sliced-points (sort-points (filter (lambda (p) (and (<= 0. (coord 0 p)) (<= (coord 0 p) 0.5))) pts)))
    (define tree-points (sort-points (kd-tree->list t)))
    (define tree-points-indices (kd-tree->list* t))

    (define nn (car sorted-points))
    (define nn1 (kd-tree-nearest-neighbor t x))

    (define knn (if (> (length sorted-points) k)
		    (take sorted-points k) 
		    sorted-points))

    (define rnn (let ((r2 (* r r)))
		  (sort-points
		   (filter
		    (lambda (y) (<= (dist2 x y) r2))
		    sorted-points)
		  )))

#|
    (define scale-factors '(1.0 0.5 1.0))

    (define (compare-scaled-distance p a b . reltol)
      (let ((dist2 (sdist2 scale-factors)))
	(let ((delta (- (dist2 p a) (dist2 p b))))
	  (if (null? reltol) 
	      delta 
	      (if (<= delta (car reltol)) 0 delta)))))

    (define (sort-scaled-points pts) 
      (sort pts (lambda (a b) (negative? (compare-scaled-distance x a b)))))
    
    (define srnn (let ((r2 (* r r))  (dist2 (sdist2 scale-factors)))
		   (sort-scaled-points 
		    (filter
		     (lambda (y) (<= (dist2 x y) r2))
		     (sort-scaled-points sorted-points))
		    )))
|#
    
    (test-group
     "KD tree consistency"
     
     (test-assert "monotonically increasing indices"
	   (sorted? tree-points-indices (lambda (x y) (< (car x) (car y)))))

     (test-assert "tree-is-valid?"
	   (kd-tree-is-valid? t))

     (test-assert "tree-all-subtrees-are-valid?"
	   (kd-tree-all-subtrees-are-valid? t))

     (test "kd-tree->list"
	  sorted-points tree-points)

     (test "nearest-neighbor"
	   nn
	   (kd-tree-nearest-neighbor t x))

     (test "k-nearest-neighbors"
	   knn
	   (sort-points (kd-tree-k-nearest-neighbors t k x)))
     
     (test "near-neighbors"
	   rnn
	   (sort-points (kd-tree-near-neighbors t r x)))

#|
     (test "near-neighbors scaled"
	   srnn
	   (sort-scaled-points (kd-tree-near-neighbors t r x factors: scale-factors)))
|#
     (test "slice"
	   sliced-points
	   (sort-points (kd-tree-slice 0 0. 0.5 t)))
     )))

(define consistency-tests/3d (make-consistency-tests Point3d KdTree3d))
(define consistency-tests/2d (make-consistency-tests Point2d KdTree2d))


(define (test1)
  (let ((n 100000) (k 40) (r 0.2) (randst (random-mtzig:init)))
  
  (let recur ((ntrials 1))
    (let* ((xs (random-mtzig:f64vector-randn! n randst))
	   (ys (random-mtzig:f64vector-randn! n randst))
	   (zs (random-mtzig:f64vector-randn! n randst))
	   (pts (let precur ((i (- n 1)) (ax '()))
		  (let ((p (make-point (f64vector-ref xs i)
				       (f64vector-ref ys i)
				       (f64vector-ref zs i))))
		    (if (positive? i)
			(precur (- i 1) (cons p ax))
			(cons p ax)))))
	   (t     (with-instance ((<KdTree> KdTree3d)) (list->kd-tree pts leaf-factor: 300)))
	   (dd    (print "tree constructed!"))
	   (xi    (inexact->exact (modulo (random-mtzig:random! randst) n)))
	   (x     (list-ref pts xi))
	   (xx    (with-instance ((<Point> Point3d)) 
				 (make-point (+ 0.1 (coord 0 x)) 
					     (- (coord 1 x) 0.1) 
					     (+ 0.1 (coord 2 x)))))
	   )

      (consistency-tests/3d pts t xx k r)
;;      (consistency-tests/2d pts t xx k r)

      (if (positive? ntrials)
	  (recur (- ntrials 1)))
      ))
  ))


(define (test2)
  (let ((n 100000) (k 40) (r 0.2) (randst (random-mtzig:init)))
  
  (let recur ((ntrials 1))
    (let ((xs (random-mtzig:f64vector-randn! n randst))
	  (ys (random-mtzig:f64vector-randn! n randst))
	  (zs (random-mtzig:f64vector-randn! n randst))
	  (cmpfn (lambda (i1 v1 i2 v2) (< v1 v2))))

      (f64vector-quick-sort! xs cmpfn)
      (f64vector-quick-sort! ys cmpfn)
      (f64vector-quick-sort! zs cmpfn)

      (print "arrays sorted!")

      (let* (
	     (axv   (list xs ys zs))
	     (t     (with-instance ((<KdTree> KdTree3d)) (f64vector->kd-tree axv leaf-factor: 300)))
	     (dd    (print "tree constructed!"))
	     (pts (let precur ((i (- n 1)) (ax '()))
		  (let ((p (make-point (f64vector-ref xs i)
				       (f64vector-ref ys i)
				       (f64vector-ref zs i))))
		    (if (positive? i)
			(precur (- i 1) (cons p ax))
			(cons p ax)))))
	     (xi    (inexact->exact (modulo (random-mtzig:random! randst) n)))
	     (x     (list-ref pts xi))
	     (xx    (with-instance ((<Point> Point3d)) 
				   (make-point (+ 0.1 (coord 0 x)) 
					       (- (coord 1 x) 0.1) 
					       (+ 0.1 (coord 2 x)))))
	     )

;;	(print (with-instance ((<KdTree> KdTree3d)) (kd-tree->list* t)))
	(consistency-tests/3d pts t xx k r)
	;;      (consistency-tests/2d pts t xx k r)
	
	(if (positive? ntrials)
	    (recur (- ntrials 1)))
	)))
  ))


(test1)
(test2)