;;; Copyright (c) 2023 Jeremy Steward
;;; All rights reserved.
;;;
;;; Redistribution and use in source and binary forms, with or without
;;; modification, are permitted provided that the following conditions are met:
;;;
;;; 1. Redistributions of source code must retain the above copyright notice,
;;; this list of conditions and the following disclaimer.
;;;
;;; 2. Redistributions in binary form must reproduce the above copyright notice,
;;; this list of conditions and the following disclaimer in the documentation
;;; and/or other materials provided with the distribution.
;;;
;;; 3. Neither the name of the copyright holder nor the names of its
;;; contributors may be used to endorse or promote products derived from this
;;; software without specific prior written permission.
;;;
;;; THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
;;; AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
;;; IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
;;; ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
;;; LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
;;; CONSEQUENTIAL DAMAGES  INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
;;; SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
;;; INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
;;; CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
;;; ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
;;; POSSIBILITY OF SUCH DAMAGE.

;; Comparator for vectors
(define vector-index-comparator
  (make-vector-comparator (make-comparator fixnum? fx=? fx<? number-hash)
                          vector?
                          vector-length
                          vector-ref))

;; Record type representing ranges of indices
;;
;; An interval is a record type that holds a start and end index representing a
;; multi-dimensional, lexicographically ordered set of half-open intervals over
;; a lexicographically ordered set of indices.
;;
;; By using an interval, one can iterate over a range of indices in multiple
;; dimensions without having to necessarily implement proper bounds checks
;; every time.
(define-record-type <interval>
  ;; Constructor for intervals
  (make-interval-internal start end)

  ;; Predicate to check if an object is an interval
  interval?

  ;; The start of the interval
  (start interval-start)

  ;; The end of the interval
  (end interval-end))

;; Checks if the interval arg is an interval. If it is not, signals an error
;; that describes that arg is not an interval in the function with the name loc
;; (a symbol) and name arg-name (also a symbol).
;;
;; Unused internally, exported for backwards compatibility.
(define (check-interval loc arg arg-name)
  (check-arg interval? arg (list loc arg-name)))

(define-checked (make-interval (start vector?) (end vector?))
  (unless (fx=? (vector-length start) (vector-length end))
    (error 'make-interval
           "Start and end indices in interval cannot be of differing ranks."
           start
           end))
  (unless (fx>? (vector-length start) 0)
    (error 'make-interval
           "Interval must have a rank > 0"
           start
           end))
  (for-each fixnum-range-fold
            (inspect
              (lambda (i)
                (let ((si (vector-ref start i))
                      (ei (vector-ref end i)))
                  (unless (and (fixnum? si) (fx>=? si 0))
                    (error 'make-interval
                           "Start index in interval must only contain non-negative fixnums."
                           start
                           i))
                  (unless (and (fixnum? ei) (fx>=? ei 0))
                    (error 'make-interval
                           "End index in interval must only contain non-negative fixnums."
                           end
                           i)))))
            (range 0 (vector-length start)))
  (make-interval-internal start end))

(define-checked (make-default-interval (dimension vector?))
  (make-interval (make-vector (vector-length dimension) 0)
                 dimension))

;; Gets the length of the interval (i.e. number of indices inside of it)
(define-checked (interval-length (interval interval?))
  (let* ((start (interval-start interval))
         (end (interval-end interval))
         (rank (vector-length end)))
    (if (>? vector-index-comparator start end)
      0
      (max 0
           (transduce fixnum-range-fold
                      (map
                        (lambda (i) (fx- (vector-ref end i) (vector-ref start i))))
                      (collect-product)
                      (range 0 rank))))))

;; Predicate which returns true iff the interval length is zero.
(define (interval-empty? interval)
  (fx=? (interval-length interval) 0))

;; Predicate that returns true iff the provided index is within the half-open
;; interval described by `interval`.
(define (interval-contains? interval index)
  (let ((start (interval-start interval))
        (end (interval-end interval)))
    (and (>=? vector-index-comparator index start)
         (transduce vector-fold
                    (zip-vector end)
                    (collect-all
                      (lambda (pair)
                        (fx<? (car pair) (cdr pair))))
                    index))))

;; Transducer aware folding operation over an interval. Starts at the `start`
;; index of an interval and folds over increments of the index
;; lexicographically until reaching the `end`.
(define-checked (interval-fold f sentinel (xs interval?))
  (let* ((start (interval-start xs))
         (end (interval-end xs))
         (rank (vector-length end)))
    (if (<? vector-index-comparator start end)
      (case rank
        ((0) sentinel)
        ((1)
         (let ((si (vector-ref start 0))
               (ei (vector-ref end 0)))
           (let loop ((i si)
                      (sentinel sentinel))
             (if (fx<? i ei)
               (let ((x (f sentinel (vector i))))
                 (if (reduced? x)
                   (unwrap x)
                   (loop (fx+ i 1) x)))
               sentinel))))
        ((2)
         (let ((si (vector-ref start 0))
               (sj (vector-ref start 1))
               (ei (vector-ref end 0))
               (ej (vector-ref end 1)))
           (let loop-i ((i si)
                        (sentinel sentinel))
             (if (fx<? i ei)
               (let loop-j ((j sj)
                            (sentinel sentinel))
                 (if (fx<? j ej)
                   (let ((x (f sentinel (vector i j))))
                     (if (reduced? x)
                       (unwrap x)
                       (loop-j (fx+ j 1) x)))
                   (loop-i (fx+ i 1) sentinel)))
               sentinel))))
        ((3)
         (let ((si (vector-ref start 0))
               (sj (vector-ref start 1))
               (sk (vector-ref start 2))
               (ei (vector-ref end 0))
               (ej (vector-ref end 1))
               (ek (vector-ref end 2)))
           (let loop-i ((i si) (sentinel sentinel))
             (if (fx<? i ei)
               (let loop-j ((j sj) (sentinel sentinel))
                 (if (fx<? j ej)
                   (let loop-k ((k sk) (sentinel sentinel))
                     (if (fx<? k ek)
                       (let ((x (f sentinel (vector i j k))))
                         (if (reduced? x)
                           (unwrap x)
                           (loop-k (fx+ k 1) x)))
                       (loop-j (fx+ j 1) sentinel)))
                   (loop-i (fx+ i 1) sentinel)))
               sentinel))))
        ((4)
         (let ((si (vector-ref start 0))
               (sj (vector-ref start 1))
               (sk (vector-ref start 2))
               (sl (vector-ref start 3))
               (ei (vector-ref end 0))
               (ej (vector-ref end 1))
               (ek (vector-ref end 2))
               (el (vector-ref end 3)))
           (let loop-i ((i si) (sentinel sentinel))
             (if (fx<? i ei)
               (let loop-j ((j sj) (sentinel sentinel))
                 (if (fx<? j ej)
                   (let loop-k ((k sk) (sentinel sentinel))
                     (if (fx<? k ek)
                       (let loop-l ((l sl) (sentinel sentinel))
                         (if (fx<? l el)
                           (let ((x (f sentinel (vector i j k l))))
                             (if (reduced? x)
                               (unwrap x)
                               (loop-l (fx+ l 1) x)))
                           (loop-k (fx+ k 1) sentinel)))
                       (loop-j (fx+ j 1) sentinel)))
                   (loop-i (fx+ i 1) sentinel)))
               sentinel))))
        (else
          (let loop-index ((sentinel sentinel)
                           (current (vector-copy start)))
            (let ((x (f sentinel (vector-copy current))))
              (if (reduced? x)
                (unwrap x)
                (let loop-rank ((axis (fx- rank 1)))
                  (let ((increment-at-axis (fx+ (vector-ref current axis) 1)))
                    (cond
                      ((fx<? increment-at-axis (vector-ref end axis))
                       (vector-set! current axis increment-at-axis)
                       (loop-index x current))
                      ((fx=? axis 0) x)
                      (else
                        (vector-set! current axis (vector-ref start axis))
                        (loop-rank (fx- axis 1)))))))))))
      sentinel)))

;; Transducer aware folding operation over an interval. Starts at the `end`
;; index of an interval and folds over decrements of the index
;; lexicographically until reaching the `start`.
(define-checked (reverse-interval-fold f sentinel (xs interval?))
  (let* ((start (interval-start xs))
         (end (interval-end xs))
         (rank (vector-length end)))
    (if (<? vector-index-comparator start end)
      (case rank
        ((0) sentinel)
        ((1)
         (let ((si (vector-ref start 0))
               (ei (vector-ref end 0)))
           (let loop ((i (fx- ei 1))
                      (sentinel sentinel))
             (if (fx>=? i si)
               (let ((x (f sentinel (vector i))))
                 (if (reduced? x)
                   (unwrap x)
                   (loop (fx- i 1) x)))
               sentinel))))
        ((2)
         (let ((si (vector-ref start 0))
               (sj (vector-ref start 1))
               (ei (vector-ref end 0))
               (ej (vector-ref end 1)))
           (let loop-i ((i (fx- ei 1))
                        (sentinel sentinel))
             (if (fx>=? i si)
               (let loop-j ((j (fx- ej 1))
                            (sentinel sentinel))
                 (if (fx>=? j sj)
                   (let ((x (f sentinel (vector i j))))
                     (if (reduced? x)
                       (unwrap x)
                       (loop-j (fx- j 1) x)))
                   (loop-i (fx- i 1) sentinel)))
               sentinel))))
        ((3)
         (let ((si (vector-ref start 0))
               (sj (vector-ref start 1))
               (sk (vector-ref start 2))
               (ei (vector-ref end 0))
               (ej (vector-ref end 1))
               (ek (vector-ref end 2)))
           (let loop-i ((i (fx- ei 1)) (sentinel sentinel))
             (if (fx>=? i si)
               (let loop-j ((j (fx- ej 1)) (sentinel sentinel))
                 (if (fx>=? j sj)
                   (let loop-k ((k (fx- ek 1)) (sentinel sentinel))
                     (if (fx>=? k sk)
                       (let ((x (f sentinel (vector i j k))))
                         (if (reduced? x)
                           (unwrap x)
                           (loop-k (fx- k 1) x)))
                       (loop-j (fx- j 1) sentinel)))
                   (loop-i (fx- i 1) sentinel)))
               sentinel))))
        ((4)
         (let ((si (vector-ref start 0))
               (sj (vector-ref start 1))
               (sk (vector-ref start 2))
               (sl (vector-ref start 3))
               (ei (vector-ref end 0))
               (ej (vector-ref end 1))
               (ek (vector-ref end 2))
               (el (vector-ref end 3)))
           (let loop-i ((i (fx- ei 1)) (sentinel sentinel))
             (if (fx>=? i si)
               (let loop-j ((j (fx- ej 1)) (sentinel sentinel))
                 (if (fx>=? j sj)
                   (let loop-k ((k (fx- ek 1)) (sentinel sentinel))
                     (if (fx>=? k sk)
                       (let loop-l ((l (fx- el 1)) (sentinel sentinel))
                         (if (fx>=? l sl)
                           (let ((x (f sentinel (vector i j k l))))
                             (if (reduced? x)
                               (unwrap x)
                               (loop-l (fx- l 1) x)))
                           (loop-k (fx- k 1) sentinel)))
                       (loop-j (fx- j 1) sentinel)))
                   (loop-i (fx- i 1) sentinel)))
               sentinel))))
        (else
          (let loop-index ((sentinel sentinel)
                           (current (apply vector
                                           (transduce fixnum-range-fold
                                                      (map
                                                        (lambda (i)
                                                          (fx- (vector-ref end i) 1)))
                                                      (collect-list)
                                                      (range 0 rank)))))
            (let ((x (f sentinel (vector-copy current))))
              (if (reduced? x)
                (unwrap x)
                (let loop-rank ((axis (fx- rank 1)))
                  (let ((decrement-at-axis (fx- (vector-ref current axis) 1)))
                    (cond
                      ((fx>=? decrement-at-axis (vector-ref start axis))
                       (vector-set! current axis decrement-at-axis)
                       (loop-index x current))
                      ((fx=? axis 0) x)
                      (else
                        (vector-set! current axis (fx- (vector-ref end axis) 1))
                        (loop-rank (fx- axis 1)))))))))))
      sentinel)))

(define-flatten-transducer flatten-interval interval? interval-fold)
(define-chain-transducer chain-interval interval-fold)

(define-flatten-transducer reverse-flatten-interval interval? reverse-interval-fold)
(define-chain-transducer reverse-chain-interval reverse-interval-fold)

(define-checked (zip-interval (interval interval?))
  (let* ((start (interval-start interval))
         (end (interval-end interval))
         (rank (vector-length end))
         (current (vector-copy start)))
    (lambda (reducer)
      (if (or (fx=? rank 0) (<? vector-index-comparator current end))
        (case rank
          ((1)
           (case-lambda
             (() (reducer))
             ((result) (reducer result))
             ((result item)
              (let ((i (vector-ref current 0)))
                (cond
                  ((fx<? i (vector-ref end 0))
                   (vector-set! current 0 (fx+ i 1))
                   (reducer result (cons item (vector i))))
                  (else
                    (make-reduced result)))))))
          ((2)
           (case-lambda
             (() (reducer))
             ((result) (reducer result))
             ((result item)
              (let ((i (vector-ref current 0))
                    (j (vector-ref current 1)))
                (cond
                  ((fx<? j (vector-ref end 1))
                   (vector-set! current 1 (fx+ j 1))
                   (reducer result (cons item (vector i j))))
                  (else
                    (let ((i (fx+ i 1))
                          (j (vector-ref start 1)))
                      (vector-set! current 0 i)
                      (vector-set! current 1 (fx+ j 1))
                      (if (fx<? i (vector-ref end 0))
                        (reducer result (cons item (vector i j)))
                        (make-reduced result)))))))))
          ((3)
           (case-lambda
             (() (reducer))
             ((result) (reducer result))
             ((result item)
              (let ((i (vector-ref current 0))
                    (j (vector-ref current 1))
                    (k (vector-ref current 2)))
                (cond
                  ((fx<? k (vector-ref end 2))
                   (vector-set! current 2 (fx+ k 1))
                   (reducer result (cons item (vector i j k))))
                  (else
                    (let ((j (fx+ j 1))
                          (k (vector-ref start 2)))
                      (vector-set! current 2 (fx+ k 1))
                      (cond
                        ((fx<? j (vector-ref end 1))
                         (vector-set! current 1 j)
                         (reducer result (cons item (vector i j k))))
                        (else
                          (let ((i (fx+ i 1))
                                (j (vector-ref start 1)))
                            (vector-set! current 0 i)
                            (vector-set! current 1 j)
                            (if (fx<? i (vector-ref end 0))
                              (reducer result (cons item (vector i j k)))
                              (make-reduced result))))))))))))
          ((4)
           (case-lambda
             (() (reducer))
             ((result) (reducer result))
             ((result item)
              (let ((i (vector-ref current 0))
                    (j (vector-ref current 1))
                    (k (vector-ref current 2))
                    (l (vector-ref current 3)))
                (cond
                  ((fx<? l (vector-ref end 3))
                   (vector-set! current 3 (fx+ l 1))
                   (reducer result (cons item (vector i j k l))))
                  (else
                    (let ((k (fx+ k 1))
                          (l (vector-ref start 3)))
                      (vector-set! current 3 (fx+ l 1))
                      (cond
                        ((fx<? k (vector-ref end 2))
                         (vector-set! current 2 k)
                         (reducer result (cons item (vector i j k l))))
                        (else
                          (let ((j (fx+ j 1))
                                (k (vector-ref start 2)))
                            (vector-set! current 2 k)
                            (cond
                              ((fx<? j (vector-ref end 1))
                               (vector-set! current 1 j)
                               (reducer result (cons item (vector i j k l))))
                              (else
                                (let ((i (fx+ i 1))
                                      (j (vector-ref start 1)))
                                  (vector-set! current 0 i)
                                  (vector-set! current 1 j)
                                  (if (fx<? i (vector-ref end 0))
                                    (reducer result (cons item (vector i j k l)))
                                    (make-reduced result)))))))))))))))
          (else
            (case-lambda
              (() reducer)
              ((result) (reducer result))
              ((result item)
               (let ((zip-item (vector-copy current)))
                 (let loop-rank ((axis (fx- rank 1)))
                   (let ((increment-at-axis (fx+ (vector-ref current axis) 1)))
                     (cond
                       ((fx<? increment-at-axis (vector-ref end axis))
                        (vector-set! current axis increment-at-axis)
                        (reducer result (cons item zip-item)))
                       ((fx=? axis 0)
                        (make-reduced
                          (reducer result (cons item zip-item))))
                       (else
                         (vector-set! current axis (vector-ref start axis))
                         (loop-rank (fx- axis 1)))))))))))
        (case-lambda
          (() reducer)
          ((result) (reducer result))
          ((result item)
           (make-reduced result)))))))

(define-checked (reverse-zip-interval (interval interval?))
  (let* ((start (interval-start interval))
         (end (interval-end interval))
         (rank (vector-length end))
         (current (apply vector
                         (transduce fixnum-range-fold
                                    (map
                                      (lambda (i) (fx- (vector-ref end i) 1)))
                                    (collect-list)
                                    (range 0 rank)))))
    (lambda (reducer)
      (if (or (fx=? rank 0) (<? vector-index-comparator start end))
        (case rank
          ((1)
           (case-lambda
             (() (reducer))
             ((result) (reducer result))
             ((result item)
              (let ((i (vector-ref current 0)))
                (cond
                  ((fx>=? i (vector-ref start 0))
                   (vector-set! current 0 (fx- i 1))
                   (reducer result (cons item (vector i))))
                  (else
                    (make-reduced result)))))))
          ((2)
           (case-lambda
             (() (reducer))
             ((result) (reducer result))
             ((result item)
              (let ((i (vector-ref current 0))
                    (j (vector-ref current 1)))
                (cond
                  ((fx>=? j (vector-ref start 1))
                   (vector-set! current 1 (fx- j 1))
                   (reducer result (cons item (vector i j))))
                  (else
                    (let ((i (fx- i 1))
                          (j (fx- (vector-ref end 1) 1)))
                      (vector-set! current 0 i)
                      (vector-set! current 1 (fx- j 1))
                      (if (fx>=? i (vector-ref start 0))
                        (reducer result (cons item (vector i j)))
                        (make-reduced result)))))))))
          ((3)
           (case-lambda
             (() (reducer))
             ((result) (reducer result))
             ((result item)
              (let ((i (vector-ref current 0))
                    (j (vector-ref current 1))
                    (k (vector-ref current 2)))
                (cond
                  ((fx>=? k (vector-ref start 2))
                   (vector-set! current 2 (fx- k 1))
                   (reducer result (cons item (vector i j k))))
                  (else
                    (let ((j (fx- j 1))
                          (k (fx- (vector-ref end 2) 1)))
                      (vector-set! current 2 (fx- k 1))
                      (cond
                        ((fx>=? j (vector-ref start 1))
                         (vector-set! current 1 j)
                         (reducer result (cons item (vector i j k))))
                        (else
                          (let ((i (fx- i 1))
                                (j (fx- (vector-ref end 1) 1)))
                            (vector-set! current 0 i)
                            (vector-set! current 1 j)
                            (if (fx>=? i (vector-ref start 0))
                              (reducer result (cons item (vector i j k)))
                              (make-reduced result))))))))))))
          ((4)
           (case-lambda
             (() (reducer))
             ((result) (reducer result))
             ((result item)
              (let ((i (vector-ref current 0))
                    (j (vector-ref current 1))
                    (k (vector-ref current 2))
                    (l (vector-ref current 3)))
                (cond
                  ((fx>=? l (vector-ref start 3))
                   (vector-set! current 3 (fx- l 1))
                   (reducer result (cons item (vector i j k l))))
                  (else
                    (let ((k (fx- k 1))
                          (l (fx- (vector-ref end 3) 1)))
                      (vector-set! current 3 (fx- l 1))
                      (cond
                        ((fx>=? k (vector-ref start 2))
                         (vector-set! current 2 k)
                         (reducer result (cons item (vector i j k l))))
                        (else
                          (let ((j (fx- j 1))
                                (k (fx- (vector-ref end 2) 1)))
                            (vector-set! current 2 k)
                            (cond
                              ((fx>=? j (vector-ref start 1))
                               (vector-set! current 1 j)
                               (reducer result (cons item (vector i j k l))))
                              (else
                                (let ((i (fx- i 1))
                                      (j (fx- (vector-ref end 1) 1)))
                                  (vector-set! current 0 i)
                                  (vector-set! current 1 j)
                                  (if (fx>=? i (vector-ref start 0))
                                    (reducer result (cons item (vector i j k l)))
                                    (make-reduced result)))))))))))))))
          (else
            (case-lambda
              (() reducer)
              ((result) (reducer result))
              ((result item)
               (let ((zip-item (vector-copy current)))
                 (let loop-rank ((axis (fx- rank 1)))
                   (let ((decrement-at-axis (fx- (vector-ref current axis) 1)))
                     (cond
                       ((fx>=? decrement-at-axis (vector-ref start axis))
                        (vector-set! current axis decrement-at-axis)
                        (reducer result (cons item zip-item)))
                       ((fx=? axis 0)
                        (make-reduced
                          (reducer result (cons item zip-item))))
                       (else
                         (vector-set! current axis (fx- (vector-ref end axis) 1))
                         (loop-rank (fx- axis 1)))))))))))
        (case-lambda
          (() reducer)
          ((result) (reducer result))
          ((result item)
           (make-reduced result)))))))

(define flatten-pair
  (lambda (reducer)
    (case-lambda
      (() (reducer))
      ((result) (reducer result))
      ((result item)
       (if (pair? item)
         (let ((result (reducer result (car item))))
           (if (reduced? result)
             result
             (reducer result (cdr item))))
         (reducer result item))))))

(define-checked (interleave-interval (interval interval?))
  (compose (zip-interval interval) flatten-pair))

(define-checked (reverse-interleave-interval (interval interval?))
  (compose (reverse-zip-interval interval) flatten-pair))

(cond-expand
  (chicken-5
    (import (only (chicken base) define-record-printer))
    (define-record-printer (<interval> interval output-port)
      (format output-port
              "#<interval ~S ~S>"
              (interval-start interval)
              (interval-end interval)))))