;;; Chickadee Game Toolkit
;;; Copyright © 2021 David Thompson <davet@gnu.org>
;;;
;;; Chickadee is free software: you can redistribute it and/or modify
;;; it under the terms of the GNU General Public License as published
;;; by the Free Software Foundation, either version 3 of the License,
;;; or (at your option) any later version.
;;;
;;; Chickadee is distributed in the hope that it will be useful, but
;;; WITHOUT ANY WARRANTY; without even the implied warranty of
;;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
;;; General Public License for more details.
;;;
;;; You should have received a copy of the GNU General Public License
;;; along with this program.  If not, see
;;; <http://www.gnu.org/licenses/>.

(define-module (chickadee data quadtree)
  #:use-module (chickadee math rect)
  #:use-module (ice-9 format)
  #:use-module (srfi srfi-9)
  #:use-module (srfi srfi-9 gnu)
  #:export (make-quadtree
            quadtree?
            quadtree-leaf?
            quadtree-bounds
            quadtree-max-depth
            quadtree-max-size
            quadtree-depth
            quadtree-size
            quadtree-q1
            quadtree-q2
            quadtree-q3
            quadtree-q4
            quadtree-clear!
            quadtree-insert!
            quadtree-delete!
            quadtree-find
            quadtree-fold
            quadtree-for-each))

;; The quadrants:
;;
;; *------*------*
;; |      |      |
;; |  Q2  |  Q1  |
;; |      |      |
;; *------*------*
;; |      |      |
;; |  Q3  |  Q4  |
;; |      |      |
;; *------*------*
(define-record-type <quadtree>
  (%make-quadtree bounds max-depth max-size depth size objects leaf?)
  quadtree?
  (bounds quadtree-bounds)
  (max-depth quadtree-max-depth)
  (max-size quadtree-max-size)
  (depth quadtree-depth)
  (size quadtree-size set-quadtree-size!)
  (objects quadtree-objects set-quadtree-objects!)
  (leaf? quadtree-leaf? set-quadtree-leaf!)
  (q1 %quadtree-q1 set-quadtree-q1!)
  (q2 %quadtree-q2 set-quadtree-q2!)
  (q3 %quadtree-q3 set-quadtree-q3!)
  (q4 %quadtree-q4 set-quadtree-q4!))

(define (display-quadtree quadtree port)
  (format port "#<quadtree bounds: ~s depth: ~d size: ~d>"
          (quadtree-bounds quadtree)
          (quadtree-depth quadtree)
          (quadtree-size quadtree)))

(set-record-type-printer! <quadtree> display-quadtree)

(define* (make-quadtree bounds #:key (max-size 5) (max-depth 4))
  "Return a new quadtree that covers the area BOUNDS.  Each node will
try to hold at maximum MAX-SIZE objects and the tree depth will be
restricted to MAX-DEPTH."
  (%make-quadtree bounds max-depth max-size 0 0 (make-vector max-size #f) #t))

(define (quadtree-empty? quadtree)
  "Return #t if QUADTREE has no objects."
  (= (quadtree-size quadtree) 0))

(define (quadtree-empty-leaf? quadtree)
  "Return #t if QUADTREE is an empty leaf node."
  (and (quadtree-empty? quadtree) (quadtree-leaf? quadtree)))

(define (quadtree-full? quadtree)
  "Return #t if QUADTREE is at or over desired maximum capacity."
  (>= (quadtree-size quadtree) (quadtree-max-size quadtree)))

(define (quadtree-max-depth? quadtree)
  "Return #t if QUADTREE is at the maximum allowed depth."
  (= (quadtree-depth quadtree) (quadtree-max-depth quadtree)))

(define (quadtree-q1 quadtree)
  (if (quadtree-leaf? quadtree)
      #f
      (%quadtree-q1 quadtree)))

(define (quadtree-q2 quadtree)
  (if (quadtree-leaf? quadtree)
      #f
      (%quadtree-q2 quadtree)))

(define (quadtree-q3 quadtree)
  (if (quadtree-leaf? quadtree)
      #f
      (%quadtree-q3 quadtree)))

(define (quadtree-q4 quadtree)
  (if (quadtree-leaf? quadtree)
      #f
      (%quadtree-q4 quadtree)))

(define (quadtree-clear! quadtree)
  "Clear QUADTREE."
  (vector-fill! (quadtree-objects quadtree) #f)
  (set-quadtree-size! quadtree 0)
  (unless (quadtree-leaf? quadtree)
    (set-quadtree-leaf! quadtree #t)
    (quadtree-clear! (%quadtree-q1 quadtree))
    (quadtree-clear! (%quadtree-q2 quadtree))
    (quadtree-clear! (%quadtree-q3 quadtree))
    (quadtree-clear! (%quadtree-q4 quadtree))))

(define (quadtree-split! quadtree)
  "Split QUADTREE region into 4 smaller child nodes."
  (let* ((max-depth (quadtree-max-depth quadtree))
         (max-size (quadtree-max-size quadtree))
         (depth (+ (quadtree-depth quadtree) 1))
         (length (vector-length (quadtree-objects quadtree)))
         (bounds (quadtree-bounds quadtree))
         (x (rect-x bounds))
         (y (rect-y bounds))
         (hw (/ (rect-width bounds) 2.0))
         (hh (/ (rect-height bounds) 2.0)))
    (define (make-node x y)
      (%make-quadtree (make-rect x y hw hh) max-depth max-size depth 0
                      (make-vector length #f) #t))
    (set-quadtree-leaf! quadtree #f)
    (unless (%quadtree-q1 quadtree)
      (set-quadtree-q1! quadtree (make-node (+ x hw) (+ y hh)))
      (set-quadtree-q2! quadtree (make-node x (+ y hh)))
      (set-quadtree-q3! quadtree (make-node x y))
      (set-quadtree-q4! quadtree (make-node (+ x hw) y)))))

(define (pick-node quadtree rect)
  "Return the child node of QUADTREE that fully contains RECT, or
QUADTREE if RECT overlaps multiple child nodes."
  (let* ((bounds (quadtree-bounds quadtree))
         (bx (rect-x bounds))
         (by (rect-y bounds))
         (bw (rect-width bounds))
         (bh (rect-height bounds))
         (mid-x (+ bx (/ bw 2.0)))
         (mid-y (+ by (/ bh 2.0)))
         (x (rect-x rect))
         (y (rect-y rect))
         (w (rect-width rect))
         (h (rect-height rect)))
    (cond
     ;; Rect does not overlap this node.
     ((or (>= x (+ bx bw))
          (<= (+ x w) bx)
          (>= y (+ by bh))
          (<= (+ y h) by))
      #f)
     ((quadtree-leaf? quadtree)
      quadtree)
     ;; Rect is within the left two quadrants: Q2 and Q3.
     ((and (< x mid-x) (< (+ x w) mid-x))
      (cond
       ;; Rect is within Q3.
       ((and (< y mid-y) (< (+ y h) mid-y))
        (%quadtree-q3 quadtree))
       ;; Rect is within Q2.
       ((> y mid-y)
        (%quadtree-q2 quadtree))
       ;; Rect spans both Q2 and Q3.
       (else quadtree)))
     ;; Rect is within the right two quadrants: Q1 and Q4.
     ((and (> x mid-x))
      (cond
       ;; Rect is within Q4.
       ((and (< y mid-y) (< (+ y h) mid-y))
        (%quadtree-q4 quadtree))
       ;; Rect is within Q1.
       ((> y mid-y)
        (%quadtree-q1 quadtree))
       ;; Rect spans both Q1 and Q4.
       (else quadtree)))
     (else quadtree))))

(define (quadtree-add! quadtree rect object)
  "Add OBJECT to the list of objets in QUADTREE."
  (let ((objects (quadtree-objects quadtree))
        (size (quadtree-size quadtree)))
    (when (= size (vector-length objects))
      (let* ((new-size (* size 2))
             (new-objects (make-vector new-size #f)))
        (let loop ((i 0))
          (when (< i size)
            (vector-set! new-objects i (vector-ref objects i))
            (loop (+ i 1))))
        (set-quadtree-objects! quadtree new-objects)))
    (vector-set! (quadtree-objects quadtree) size (cons rect object))
    (set-quadtree-size! quadtree (+ size 1))))

(define (quadtree-merge-maybe! quadtree)
  "Remove child nodes if they are all empty leaf nodes."
  (when (and (quadtree-empty-leaf? (%quadtree-q1 quadtree))
             (quadtree-empty-leaf? (%quadtree-q2 quadtree))
             (quadtree-empty-leaf? (%quadtree-q3 quadtree))
             (quadtree-empty-leaf? (%quadtree-q4 quadtree)))
    ;; We don't actually get rid of the child nodes.  This means that
    ;; the quadtree can take more memory than it has to, but it also
    ;; means that the quadtree doesn't allocate needlessly when
    ;; objects are constantly being added/removed.
    (set-quadtree-leaf! quadtree #t)))

(define (quadtree-insert! quadtree rect object)
  "Insert OBJECT with bounding box RECT into QUADTREE."
  (let ((node (pick-node quadtree rect)))
    (cond
     ;; The rect doesn't fit into the parent node.
     ((not node)
      #f)
     ;; The rect fits completely within one of the child nodes, so
     ;; descend into that node and repeat the process.
     ((not (eq? quadtree node))
      (quadtree-insert! node rect object))
     ;; The node is a leaf node that is at or over the desired
     ;; capacity, so we need to split it and redistribute the objects.
     ;; Nodes that have reached the maximum allowed depth cannot be
     ;; split.
     ((and (quadtree-full? quadtree)
           (quadtree-leaf? quadtree)
           (not (quadtree-max-depth? quadtree)))
      (let ((objects (quadtree-objects quadtree))
            (size (quadtree-size quadtree)))
        (quadtree-split! quadtree)
        (let loop ((i 0))
          (if (< i size)
              (let ((obj (vector-ref objects i)))
                (loop (+ i 1))
                (quadtree-insert! quadtree (car obj) (cdr obj)))
              (begin
                (set-quadtree-size! quadtree 0)
                (vector-fill! objects #f))))
        (quadtree-insert! quadtree rect object)
        (quadtree-merge-maybe! quadtree)))
     ;; The node is either under the desired maximum objects threshold
     ;; or has no choice but to accept another object because there is
     ;; no child node that fully contains the rect or we have reached
     ;; the maximum allowed tree depth.
     (else
      (quadtree-add! quadtree rect object)))))

(define (quadtree-delete! quadtree rect object)
  "Delete OBJECT, who occupies the space RECT, from QUADTREE."
  (let ((node (pick-node quadtree rect)))
    (cond
     ((not node)
      #f)
     ((eq? quadtree node)
      (let ((objects (quadtree-objects quadtree))
            (size (quadtree-size quadtree)))
        (let loop ((i 0))
          (cond
           ((= i size)
            #f)
           ((eq? (cdr (vector-ref objects i)) object)
            (let ((new-size (- size 1)))
              (vector-set! objects i (vector-ref objects new-size))
              (vector-set! objects new-size #f)
              (set-quadtree-size! quadtree new-size)
              #t))
           (else
            (loop (+ i 1)))))))
     (else
      (and (quadtree-delete! node rect object)
           (begin
             (quadtree-merge-maybe! quadtree)
             #t))))))

(define (quadtree-find quadtree rect pred)
  "Return the first object in QUADTREE in the vicinity of RECT that
satisfies PRED."
  (let ((objects (quadtree-objects quadtree))
        (size (quadtree-size quadtree))
        (next-node (pick-node quadtree rect)))
    (if next-node
        (let loop ((i 0))
          (if (< i size)
              (let ((object (cdr (vector-ref objects i))))
                (if (pred object)
                    object
                    (loop (+ i 1))))
              (cond
               ((and (eq? next-node quadtree) (not (quadtree-leaf? quadtree)))
                (or (quadtree-find (%quadtree-q1 quadtree) rect pred)
                    (quadtree-find (%quadtree-q2 quadtree) rect pred)
                    (quadtree-find (%quadtree-q3 quadtree) rect pred)
                    (quadtree-find (%quadtree-q4 quadtree) rect pred)))
               ((eq? next-node quadtree)
                #f)
               (else
                (quadtree-find next-node rect pred)))))
        #f)))

(define (quadtree-fold quadtree rect init proc)
  "Apply PROC to all objects in the vicinity of RECT in QUADTREE to
build a result and return that result.  INIT is the initial result.
If there are no objects in the vicinity of RECT, just INIT is
returned."
  (let ((objects (quadtree-objects quadtree))
        (size (quadtree-size quadtree))
        (next-node (pick-node quadtree rect)))
    (if next-node
        (let loop ((i 0)
                   (memo init))
          (if (< i size)
              (loop (+ i 1) (proc (cdr (vector-ref objects i)) memo))
              (cond
               ((and (eq? next-node quadtree) (not (quadtree-leaf? quadtree)))
                (let* ((q1 (quadtree-fold (%quadtree-q1 quadtree) rect memo proc))
                       (q2 (quadtree-fold (%quadtree-q2 quadtree) rect q1 proc))
                       (q3 (quadtree-fold (%quadtree-q3 quadtree) rect q2 proc)))
                  (quadtree-fold (%quadtree-q4 quadtree) rect q3 proc)))
               ((eq? next-node quadtree)
                memo)
               (else
                (quadtree-fold next-node rect memo proc)))))
        init)))

(define (quadtree-for-each quadtree rect proc)
  "Call PROC for all objects in the vicinity of RECT in QUADTREE."
  (let ((objects (quadtree-objects quadtree))
        (size (quadtree-size quadtree))
        (next-node (pick-node quadtree rect)))
    (when next-node
      (let loop ((i 0))
        (when (< i size)
          (proc (cdr (vector-ref objects i)))
          (loop (+ i 1))))
      (cond
       ((and (eq? next-node quadtree) (not (quadtree-leaf? quadtree)))
        (quadtree-for-each (%quadtree-q1 quadtree) rect proc)
        (quadtree-for-each (%quadtree-q2 quadtree) rect proc)
        (quadtree-for-each (%quadtree-q3 quadtree) rect proc)
        (quadtree-for-each (%quadtree-q4 quadtree) rect proc))
       ((eq? next-node quadtree)
        *unspecified*)
       (else
        (quadtree-for-each next-node rect proc))))))