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

;; Slices an array starting at the `start` index and ending at the optional
;; `end` index.
(define array-slice
  (case-lambda-checked
    ((array start)
     (array-slice array start (interval-end (array-shape array))))
    (((array array?) (start vector?) (end vector?))
     (let* ((original-interval (array-interval array))
            (original-start (interval-start original-interval))
            (original-end (interval-end original-interval))
            (original-rank (array-rank array))
            (original-len (interval-length original-interval))
            (rank (vector-length end)))
       (unless (fx=? original-rank rank)
         (error 'array-slice
                "The rank of the provided interval is not equal to the original rank."
                original-rank
                rank))
       (let* ((offset-index-by
                (lambda (i j rank)
                  (transduce fixnum-range-fold
                             (map (lambda (r) (fx+ (vector-ref i r) (vector-ref j r))))
                             (collect-vector rank)
                             (range 0 rank))))
              (adjusted-start (offset-index-by start original-start rank))
              (adjusted-end (offset-index-by end original-start rank))
              (adjusted-interval (make-interval adjusted-start adjusted-end))
              (adjusted-len (interval-length adjusted-interval)))
         (unless (fx>=? original-len adjusted-len)
           (error 'array-slice
                  "The length of the provided interval is not equal to the original shape."
                  original-len
                  adjusted-len))
         (if (and (>=? vector-index-comparator adjusted-start original-start)
                  (<=? vector-index-comparator adjusted-end original-end))
           (make-array-internal (array-storage-class array)
                                adjusted-interval
                                (array-stride array)
                                (array-storage-object array))
           (error 'array-slice
                  "Interval exceeds array dimensions"
                  original-interval
                  start
                  end)))))))

;; Monadic transposition of an array. This is done by reversing the indices in
;; the interval, as well as the stride.
(define-checked (array-transpose (array array?))
  (let* ((interval (array-interval array))
         (reverse-index
           (lambda (idx)
             (apply vector (transduce vector-fold values (collect-reverse-list) idx))))
         (transposed-interval
           (make-interval (reverse-index (interval-start interval))
                          (reverse-index (interval-end interval))))
         (transposed-stride (reverse-index (array-stride array))))
    (make-array-internal (array-storage-class array)
                         transposed-interval
                         transposed-stride
                         (array-storage-object array))))

;; Singular triadic transposition of an array. This swaps the ordering of two
;; axes in an interval. For example, if you have an array with an interval of
;; #(2 1 0) through #(5 3 3), you can then swap the first and second axes by
;; calling (array-swap-axes array 0 1), which would then leave you with the
;; same array, but with the interval #(1 2 0) through #(3 5 3).
;;
;; This is similar to how array-transpose works, but merely swaps two axes
;; instead of reversing the entire interval and stride.
(define-checked (array-swap-axes (array array?)
                                 (to-axis non-negative-fixnum?)
                                 (from-axis non-negative-fixnum?))
  (let* ((interval (array-interval array))
         (rank (vector-length (interval-end interval))))
    (unless (fx<? to-axis rank)
      (error 'array-swap-axes
             "to-axis is greater than the array rank."
             to-axis))
    (unless (fx<? from-axis rank)
      (error 'array-swap-axes
             "from-axis is greater than the array rank."
             from-axis))
    (let ((swapped-start (vector-copy (interval-start interval)))
          (swapped-end (vector-copy (interval-end interval)))
          (swapped-stride (vector-copy (array-stride array))))
      (vector-swap! swapped-start to-axis from-axis)
      (vector-swap! swapped-end to-axis from-axis)
      (vector-swap! swapped-stride to-axis from-axis)
      (make-array-internal (array-storage-class array)
                           (make-interval swapped-start swapped-end)
                           swapped-stride
                           (array-storage-object array)))))

;; Removes the `axis` from the array's interval / stride iff the length of that
;; axis is 1, otherwise signals an error.
(define-checked (array-squeeze-axis (array array?) (axis non-negative-fixnum?))
  (let* ((shape (array-shape array))
         (rank (vector-length (interval-end shape))))
    (unless (fx<? axis rank)
      (error 'array-squeeze-axis "Axis is greater or equals array rank" axis rank))
    (let ((axis-length (vector-ref (interval-end shape) axis)))
      (cond
        ((and (fx=? rank 1)
              (fx=? axis 0)
              (fx=? axis-length 1))
         (array-ref array (interval-start shape)))
        ((and (fx<? axis rank)
              (fx=? axis-length 1))
         (let* ((interval (array-interval array))
                (squeeze-idx
                  (lambda (idx axis)
                    (transduce fixnum-range-fold
                               (compose
                                 (remove (lambda (i) (fx=? i axis)))
                                 (map (lambda (i) (vector-ref idx i))))
                               (collect-vector)
                               (range 0 rank))))
                (new-start (squeeze-idx (interval-start interval) axis))
                (new-end (squeeze-idx (interval-end interval) axis))
                (new-stride (squeeze-idx (array-stride array) axis)))
           (make-array-internal (array-storage-class array)
                                (make-interval new-start new-end)
                                new-stride
                                (array-storage-object array))))
        ((not (fx=? axis-length 1))
         (error 'array-squeeze-axis "Axis is not of length 1" axis axis-length))
        (else
          (error 'array-squeeze-axis "unreachable!"))))))

;; Removes all axes in the array whose length is 1.
;;
;; If every axis would be squeezed out by this operation, returns the singular
;; value in the array at that point.
(define-checked (array-squeeze (array array?))
  (let* ((shape (array-shape array))
         (dim (interval-end shape))
         (rank (vector-length dim))
         (squeeze-index
           (lambda (dim index rank)
             (transduce fixnum-range-fold
                        (compose
                          (remove
                            (lambda (axis) (fx=? (vector-ref dim axis) 1)))
                          (map
                            (lambda (axis) (vector-ref index axis))))
                        (collect-vector rank)
                        (range 0 rank))))
         (new-start (squeeze-index dim (interval-start shape) rank))
         (new-end (squeeze-index dim (interval-end shape) rank))
         (new-stride (squeeze-index dim (array-stride array) rank)))
    (make-array-internal (array-storage-class array)
                         (make-interval new-start new-end)
                         new-stride
                         (array-storage-object array))))

(define (increment-after val)
  (lambda (reducer)
    (case-lambda
      (() reducer)
      ((result) (reducer result))
      ((result item)
       (if (fx>=? item val)
         (reducer result (fx+ item 1))
         (reducer result item))))))

;; Adds an `axis` of length 1 to the internal array's interval / stride.
;;
;; Can be thought of as the opposite of `array-squeeze-axis`.
(define-checked (array-expand-axis (array array?) (axis non-negative-fixnum?))
  (let* ((shape (array-shape array))
         (rank (vector-length (interval-end shape))))
    (unless (fx<? axis rank)
      (error 'array-expand-axis "Axis is greater or equals array rank" axis rank))
    (let* ((interval (array-interval array))
           (expand-idx
             (lambda (idx axis val)
               (let ((new-idx (make-vector (fx+ rank 1) #f)))
                 (vector-set! new-idx axis val)
                 (for-each fixnum-range-fold
                           (compose
                             (increment-after axis)
                             enumerate
                             (inspect
                               (lambda (pair)
                                 (let ((idx-i (car pair))
                                       (new-i (cdr pair)))
                                   (vector-set! new-idx new-i (vector-ref idx idx-i))))))
                           (range 0 rank))
                 new-idx)))
           (new-start (expand-idx (interval-start interval) axis 0))
           (new-end (expand-idx (interval-end interval) axis 1))
           (new-stride (expand-idx (array-stride array) axis 1)))
      (make-array-internal (array-storage-class array)
                           (make-interval new-start new-end)
                           new-stride
                           (array-storage-object array)))))