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

(import r7rs)

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

(import (mcgoron weight-balanced-trees internal)
        (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))

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

;;; ;;;;;;;;;;;;;;;;;;;
;;; 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 node number-comparator middle))))
  (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 node number-comparator middle)))))
  (test-property search-does-not-find (list (split-vector-generator))))

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

(define split-finds
  (call/inserted
   (lambda (left middle right node)
     (let-values (((new-left found? new-right)
                   (split number-comparator node middle)))
       found?))))

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

(define split-does-not-find
  (call-w/o-inserted
   (lambda (left middle right node)
     (let-values (((new-left found? new-right)
                   (split number-comparator node middle)))
       (not found?)))))

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