#| Copyright 2024 Peter McGoron
 |
 | Licensed under the Apache License, Version 2.0 (the "License");
 | you may not use this file except in compliance with the License.
 | You may obtain a copy of the License at
 |
 |     http://www.apache.org/licenses/LICENSE-2.0
 |
 | Unless required by applicable law or agreed to in writing, software
 | distributed under the License is distributed on an "AS IS" BASIS,
 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 | See the License for the specific language governing permissions and
 | limitations under the License.
 |----------------------------------------------
 | This test suite would be better if property tests (and by association
 | generators) could take multiple values. As it stands now property
 | test generators cannot be correlated without doing a hack like returning
 | a vector (which the current implementation does).
 |#

(import r7rs)

(load "random-number-vector.sld")
(load "util.sld")

(import (mcgoron weight-balanced-trees test util)
        (mcgoron weight-balanced-trees test random-number-vector)
        (prefix (only (mcgoron srfi 64) factory) mcgoron-test-)
        (except (mcgoron srfi 64) factory)
        (srfi 1) (srfi 26) (srfi 64) (srfi 128) (srfi 132) (srfi 133)
        (srfi 158) (srfi 194) (srfi 197) (srfi 252)
        (except (mcgoron weight-balanced-trees internal)
                every)
        (prefix (only (mcgoron weight-balanced-trees internal)
                      every)
                node-))

;;; ;;;;;;;;;;;;;;;;
;;; Parameters

(define verbose (make-parameter #f))
(define max-vector-length (make-parameter 100))

;;; ;;;;;;;;;;;;;;;;;;;;
;;; Setup my test runner
;;; ;;;;;;;;;;;;;;;;;;;;

(test-runner-factory mcgoron-test-factory)
(test-runner-current (test-runner-create))
(set-verbosity! 'fails)

;;; ;;;;;;;;;;;;;;;;;;;
;;; Random in-order vector generator
;;; ;;;;;;;;;;;;;;;;;;;

(define (vector->node-generator)
  (gmap in-order-vector->node (in-order-vector-generator
                               (max-vector-length))))

(define number-comparator
  (make-comparator
   number?
   =
   <
   (lambda (x) x)))

;;; ;;;;;;;;;;;;;;;;;;;;;;
;;; Test node creation
;;; ;;;;;;;;;;;;;;;;;;;;;;

(test-group "vector->node is correct"
  (define (vector->node->list? vec)
    (let* ((node (in-order-vector->node vec))
           (lst (node->in-order-list node)))
      (when (verbose)
        (dynamic-property-set! 'vec vec)
        (dynamic-property-set! 'lst lst))
      (equal? (vector->list vec) lst)))
  (test-property vector->node->list?
                 (list (in-order-vector-generator (max-vector-length)))))

(test-group "vector->node balanced?"
  (test-assert (balanced? '()))
  (test-property balanced? (list (vector->node-generator))))

;;; ;;;;;;;;;
;;; Make a "split" vector. A "split" vector is two sorted vectors and a 
;;; number, where for all indicies i
;;; 
;;;      v1[i] < number
;;;      number < v2[i]
;;;   
;;; ;;;;;;;;;

(define (split-in-order-vector vec)
  ;; Take an in order vector and split it into two in-order vectors and
  ;; an element in between those two vectors.
  (let* ((len (vector-length vec))
         (pivot ((make-random-integer-generator 0 len))))
    (vector (vector-copy-exclusive vec 0 pivot)
            (vector-ref vec pivot)
            (vector-copy-exclusive vec (+ pivot 1) len))))

(define (split-vector-generator)
  (chain (in-order-vector-generator (max-vector-length))
         (gmap split-in-order-vector _)))

(define (call/split-vecs procedure)
  (lambda (vec)
    (let ((left (vector-ref vec 0))
          (middle (vector-ref vec 1))
          (right (vector-ref vec 2)))
      (procedure left middle right))))

(define (join-vectors left middle right)
  (join middle (in-order-vector->node left) (in-order-vector->node right)))

;;; ;;;;;;;;;;;;;;;;;;;;;;;;;;;
;;; join
;;; ;;;;;;;;;;;;;;;;;;;;;;;;;;;

(define join-preserves-order-of
  (call/split-vecs
   (lambda (left middle right)
     ;; Test if joining two in-order vectors with a node in between will
     ;; preserve order in the tree.
     ;; Log the left, middle, and right vectors.
     (when (verbose)
       (dynamic-property-set! 'right right)
       (dynamic-property-set! 'middle middle)
       (dynamic-property-set! 'left left))
     (let ((ret (node->in-order-list
                 (join-vectors left middle right)))
           (orig (append (vector->list left)
                         (list middle)
                         (vector->list right))))
       (when (verbose)
         (dynamic-property-set! 'ret ret))
       (equal? ret orig)))))

(test-group "join of in-order vectors preserves order"
  (parameterize ((verbose #f))
    (test-property join-preserves-order-of
                   (list (split-vector-generator)))))

(test-group "join of in-order vectors is balanced"
  (define (joined-vector-generator)
    (gmap (call/split-vecs join-vectors)
          (split-vector-generator)))
  (test-property balanced? (list (joined-vector-generator))))

;;; ;;;;;;;;;;;;;;;;;;;;;;;;;;
;;; join2
;;; ;;;;;;;;;;;;;;;;;;;;;;;;;;

(define join2-preserves-order-of
  (call/split-vecs
   (lambda (left _ right)
     (let ((ret (node->in-order-list
                 (join2 (in-order-vector->node left)
                        (in-order-vector->node right))))
           (orig (append (vector->list left)
                         (vector->list right))))
       (equal? ret orig)))))

(test-group "join2 of in-order vectors preserves order"
  (test-property join2-preserves-order-of
                 (list (split-vector-generator))))

(test-group "join2 of in-order vectors is balanced"
  (define (join2-split-vector-generator)
    (define (join2-of-split vec)
      (join2 (in-order-vector->node (vector-ref vec 0))
             (in-order-vector->node (vector-ref vec 2))))
    (gmap join2-of-split (split-vector-generator)))
  (test-property balanced?
                 (list (join2-split-vector-generator))))

;;; ;;;;;;
;;; search
;;; ;;;;;;

(define (call/inserted procedure)
  (call/split-vecs
   (lambda (left middle right)
     (let ((node (join-vectors left middle right)))
       (procedure left middle right node)))))

(define (call-w/o-inserted procedure)
  (call/split-vecs
   (lambda (left middle right)
     (let ((node (join2 (in-order-vector->node left)
                        (in-order-vector->node right))))
       (procedure left middle right node)))))

(test-group "search with inserted node"
  (define search-finds
    (call/inserted
     (lambda (left middle right node)
       (when (verbose)
         (dynamic-property-set! 'expected middle)
         (dynamic-property-set! 'right right)
         (dynamic-property-set! 'left left)
         (dynamic-property-set! 'node (node->in-order-list node)))
       (search number-comparator middle node))))
  (parameterize ((verbose #f))
    (test-property search-finds (list (split-vector-generator)))))

(test-group "search without inserted node"
  (define search-does-not-find
    (call-w/o-inserted
     (lambda (left middle right node)
       (not (search number-comparator middle node)))))
  (test-property search-does-not-find (list (split-vector-generator))))

;;; ;;;;;;;;;;;;;;;;;;;;;;;;
;;; split
;;; ;;;;;;;;;;;;;;;;;;;;;;;;

(define (call/split wrapper procedure)
  (wrapper
   (lambda (left middle right node)
     (call-with-values (lambda ()
                         (split number-comparator node middle
                                (lambda () #f)))
       procedure))))

(define split-finds
  (call/split
   call/inserted
   (lambda (new-left found? new-right)
     found?)))

(test-group "split finds"
  (test-property split-finds (list (split-vector-generator))))

(define split-left-balanced
  (call/split
   call/inserted
   (lambda (new-found found? new-right)
     (balanced? new-found))))

(test-group "split left balanced"
  (test-property split-left-balanced (list (split-vector-generator))))

(define split-right-balanced
  (call/split
   call/inserted
   (lambda (new-found found? new-right)
     (balanced? new-right))))

(test-group "split right balanced"
  (test-property split-right-balanced (list (split-vector-generator))))

(define split-does-not-find
  (call/split
   call-w/o-inserted
   (lambda (new-left found? new-right)
     (not found?))))

(test-group "split does not find"
  (test-property split-does-not-find (list (split-vector-generator))))

(define split-left-balanced-w/o-inserted
  (call/split
   call-w/o-inserted
   (lambda (new-left found? new-right)
     (balanced? new-left))))

(test-group "split left balanced without insertion"
  (test-property split-left-balanced-w/o-inserted (list (split-vector-generator))))

(define split-right-balanced-w/o-inserted
  (call/split
   call-w/o-inserted
   (lambda (new-left found? new-right)
     (balanced? new-right))))

(test-group "split right balanced without insertion"
  (test-property split-right-balanced-w/o-inserted (list (split-vector-generator))))

;;; ;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;;; Union
;;; ;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(test-group "union balanced"
  (define (union-balanced node1 node2)
    (balanced? (union number-comparator node1 node2)))
  (test-property union-balanced
                 (list (vector->node-generator)
                       (vector->node-generator))))

(define (union-subset-of-lset node1 node2)
    (let ((node (union number-comparator node1 node2))
          (union (lset-union =
                             (node->in-order-list node1)
                             (node->in-order-list node2))))
      (if (verbose)
          (dynamic-property-set! 'union
                                 (node->in-order-list node)))
      (every (cut search number-comparator <> node) union)))

(test-group "union subset of lset"
  (parameterize ((verbose #t))
    (test-property union-subset-of-lset
                   (list (vector->node-generator)
                         (vector->node-generator)))))

(test-group "union subset of lset, right null"
  (test-property union-subset-of-lset
                 (list (vector->node-generator)
                       (circular-generator '()))))

(test-group "union subset of lset, left null"
  (test-property union-subset-of-lset
                 (list (circular-generator '())
                       (vector->node-generator))))

(define (lset-subset-of-union node1 node2)
  (let ((node (union number-comparator node1 node2))
        (union (lset-union =
                           (node->in-order-list node1)
                           (node->in-order-list node2))))
    (node-every (cut memq <> union) node)))

(test-group "lset subset of union"
  (test-property lset-subset-of-union
                 (list (vector->node-generator)
                       (vector->node-generator))))

(test-group "lset subset of union, right null"
  (test-property lset-subset-of-union
                 (list (vector->node-generator)
                       (circular-generator '()))))

(test-group "lset subset of union, left null"
  (test-property lset-subset-of-union
                 (list (circular-generator '())
                       (vector->node-generator))))

;;; ;;;;;;;;;;;;;;;;;;;;
;;; Intersection
;;; ;;;;;;;;;;;;;;;;;;;;

(test-group "intersection balanced?"
  (define (intersection-balanced? node1 node2)
    (balanced? (intersection number-comparator node1 node2)))
  (test-property intersection-balanced?
                 (list (vector->node-generator)
                       (vector->node-generator))))

(test-group "intersection of itself"
  (define (intersection-of-itself node1)
    (let ((node (intersection number-comparator node1 node1))
          (lst (node->in-order-list node1)))
      (every (cut search number-comparator <> node) lst)))
  (test-property intersection-of-itself
                 (list (vector->node-generator))))

(test-group "intersection of null, left"
  (define (intersection-of-null-left node1)
    (null? (intersection number-comparator node1 '())))
  (test-property intersection-of-null-left
                 (list (vector->node-generator))))

(test-group "intersection of null, right"
  (define (intersection-of-null-right node1)
    (null? (intersection number-comparator '() node1)))
  (test-property intersection-of-null-right 
                 (list (vector->node-generator))))

(test-group "intersection subset of lset"
  (define (intersection-lset node1 node2)
    (let ((node (intersection number-comparator node1 node2))
          (intersection (lset-intersection =
                                           (node->in-order-list node1)
                                           (node->in-order-list node2))))
      (every (cut search number-comparator <> node) intersection)))
  (test-property intersection-lset
                 (list (vector->node-generator)
                       (vector->node-generator))))

(test-group "lset subset of intersection"
  (define (lset-subset-of-intersection node1 node2)
    (let ((node (intersection number-comparator node1 node2))
          (intersection (lset-intersection =
                                           (node->in-order-list node1)
                                           (node->in-order-list node2))))
      (node-every (cut memv <> intersection) node)))
  (test-property lset-subset-of-intersection
                 (list (vector->node-generator)
                       (vector->node-generator))))

;;; ;;;;;;;;;;;;;;;;;;;;;;;;;;
;;; Difference
;;; ;;;;;;;;;;;;;;;;;;;;;;;;;;

(test-group "difference subset of lset"
  (define (difference-subset-of-lset node1 node2)
    (let ((node (difference number-comparator node1 node2))
          (intersection (lset-difference =
                                         (node->in-order-list node1)
                                         (node->in-order-list node2))))
      (every (cut search number-comparator <> node) intersection)))
  (test-property difference-subset-of-lset
                 (list (vector->node-generator)
                       (vector->node-generator))))

(test-group "lset subset of difference"
  (define (lset-subset-of-difference node1 node2)
    (let ((node (difference number-comparator node1 node2))
          (intersection (lset-difference =
                                         (node->in-order-list node1)
                                         (node->in-order-list node2))))
      (node-every (cut memq <> intersection) node)))
  (test-property lset-subset-of-difference
                 (list (vector->node-generator)
                       (vector->node-generator))))

;;; ;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;;; Update
;;; ;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(define (insert elem set)
  (update number-comparator
          set
          elem
          (lambda (value) value)
          (lambda ()
            (wb-tree-node elem '() '()))))

(test-group "update of unordered list is balanced"
  (define (insert-unordered lst)
    (let ((set (fold insert '() lst)))
      (balanced? set)))

  (test-property insert-unordered
                 (list (list-generator-of
                        (exact-integer-generator)))))

(test-group "update of unordered list is subset of list"
  (define (subset-of-list lst)
    (let ((set (fold insert '() lst)))
      (when (verbose)
        (dynamic-property-set! 'lst lst)
        (dynamic-property-set! 'set (node->in-order-list set)))
      (every (cut search number-comparator <> set) lst)))
  (parameterize ((verbose #f))
    (test-property subset-of-list
                   (list (list-generator-of
                          (exact-integer-generator))))))

(test-group "list is subset of update of unordered list"
  (define (list-is-subset-of lst)
    (let ((set (fold insert '() lst)))
      (node-every (cut memq <> lst) set)))
  (test-property list-is-subset-of
                 (list (list-generator-of
                        (exact-integer-generator)))))

;;; ;;;;;;;;;;;;;;;;
;;; Delete
;;; ;;;;;;;;;;;;;;;;

(define (call/difference procedure)
  (lambda (lst1 lst2)
    (let* ((set1 (fold insert '() lst1))
           (set2 (fold insert '() lst2))
           (diff (difference number-comparator set1 set2))
           (deleted (fold (lambda (elem set)
                            (when (verbose)
                              (dynamic-property-set! 'elem elem)
                              (dynamic-property-set! 'set set))
                            (delete number-comparator elem set))
                          set1
                          lst2)))
      (procedure lst1 lst2 diff deleted))))

(test-group "delete of two sets is subset of difference"
  (define (delete-and-difference lst1 lst2 diff deleted)
    (node-every (cut search number-comparator <> diff) deleted))
  (parameterize ((verbose #f))
    (test-property (call/difference delete-and-difference)
                   (list (list-generator-of (exact-integer-generator))
                         (list-generator-of (exact-integer-generator))))))

(test-group "difference is subset of delete of two sets"
  (define (difference-and-delete lst1 lst2 diff deleted)
    (node-every (cut search number-comparator <> deleted) diff))
  (test-property (call/difference difference-and-delete)
                 (list (list-generator-of (exact-integer-generator))
                       (list-generator-of (exact-integer-generator)))))