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

;; A record for representing a multi-dimensional array of a given index and
;; storage class.
;;
;; Arrays can be treated as roughly analagous to Scheme vectors, except that
;; they may span several dimensions dependent upon the interval used to index
;; the array. Anything can be used to index an array so long as the appropriate
;; index-class is provided (see the index-classes module for more information).
;; Index-classes define how data is conceptually "laid out" in the array, or
;; how one might reference a unit or element of data within the array.
;;
;; The storage-class of an array defines the underlying datatype used to store
;; data in an array. Storage-classes are inherently a linear, random-access
;; data type. It is assumed that data contained within a storage class is
;; contiguous, or can transparently be represented as though it were. See the
;; storage-classes module for more information.
(define-record-type <array>
  ;; Constructs a new array from a provided index class, storage-class, shape,
  ;; stride, offset, and storage-object.
  ;;
  ;; This constructor is meant to be internal so as to hide explicit shape /
  ;; stride / offset generation through the public API, since this is largely
  ;; something that can be derived and does not require user effort to
  ;; generate.
  (make-array-internal storage-class
                       interval
                       stride
                       storage-object)

  ;; Predicate function for checking if an object is an array.
  array?

  ;; The storage-class of this array.
  (storage-class array-storage-class)

  ;; An interval describing the internal shape (i.e. the start index is not
  ;; normalized to default / zero) of the array.
  (interval array-interval)

  ;; The stride of the array. The stride is used to determine how to convert an
  ;; index into a linear offset into the underlying storage-object.
  (stride array-stride)

  ;; The storage object (e.g. vector, f64vector, etc.) that contains the
  ;; array's data.
  (storage-object array-storage-object))

;; Calculates the stride of an array based on the shape. Implicitly assumes
;; row-major ordering (as does most of this API), but can optionally be used to
;; specify column major ordering if necessary.
(define-checked (dimension->stride (dimension vector?))
  (let ((xs (vector-cumulate
              fx*
              1
              (list->vector
                (cons 1 (transduce fixnum-range-fold
                                   (map (lambda (i) (vector-ref dimension i)))
                                   (collect-reverse-list)
                                   (range 1 (vector-length dimension))))))))
    (apply vector (transduce vector-fold values (collect-reverse-list) xs))))

;;; Constructors

;; Constructs an array filled with `fill` values, with `index-class` and
;; `storage-class` defining the rank/indexing and storage of the array
;; respectively. `dim` is an index of `index-class` defining the total
;; dimensions of the array.
(define make-array
  (case-lambda-checked
    ((storage-class dimension)
     (make-array storage-class
                 dimension
                 (storage-class-default-element storage-class)))
    (((storage-class storage-class?) (dimension vector?) fill)
     (let ((stride (dimension->stride dimension))
           (len (transduce vector-fold values (collect-product) dimension))
           (interval (make-default-interval dimension)))
       (make-array-internal storage-class
                            interval
                            stride
                            (make-storage-object storage-class len fill))))))

;; Constructs an array with the provided `storage-object`, with `index-class`
;; and `storage-class` defining the rank/indexing and storage-kind of the array
;; respectively. `dim` is an index of `index-class` defining the total
;; dimensions of the array.
(define-checked (make-array-from-storage (storage-class storage-class?) (dimension vector?) storage-object)
  (comparator-check-type (storage-class-comparator storage-class) storage-object)
  (let ((len (transduce vector-fold values (collect-product) dimension)))
    (if (fx=? len (storage-object-length storage-class storage-object))
      (make-array-internal storage-class
                           (make-default-interval dimension)
                           (dimension->stride dimension)
                           storage-object)
      (error 'make-array-from-storage
             (format #f
                     "Storage length mismatch. Expected dimension ~s with length ~s but got ~s"
                     dimension
                     len
                     (storage-object-length storage-class storage-object))
             storage-object))))

;; Gets the rank (number of dimensions) of the array.
(define-checked (array-rank (array array?))
  (vector-length (interval-end (array-interval array))))

;; Gets the shape of the array, which is an index range starting at the default
;; index and ending at the upper dimension of the array.
(define-checked (array-shape (array array?))
  (let* ((interval (array-interval array))
         (start (interval-start interval))
         (end (interval-end interval))
         (rank (vector-length end)))
    (make-default-interval
      (transduce fixnum-range-fold
                 (map
                   (lambda (i) (fx- (vector-ref end i) (vector-ref start i))))
                 (collect-vector)
                 (range 0 rank)))))

;; Helper collector for using fixnum arithmetic to compute a fixnum sum
(define collect-fx-sum
  (case-lambda
    (() (collect-fx-sum 0))
    ((sentinel)
     (case-lambda
       (() sentinel)
       ((result) result)
       ((result item) (fx+ result item))))))

;; Converts an index object and stride into a linear offset, assuming that the
;; lexicographic ordering of the index is the major order.
;;
;; Both the index and the stride must be indices of the provided index class.
(define-checked (index->offset (index vector?) (stride vector?) (rank non-negative-fixnum?))
  (case rank
    ((1)
     (fx* (vector-ref index 0) (vector-ref stride 0)))
    ((2)
     (fx+ (fx* (vector-ref index 0) (vector-ref stride 0))
          (fx* (vector-ref index 1) (vector-ref stride 1))))
    ((3)
     (fx+ (fx+ (fx* (vector-ref index 0) (vector-ref stride 0))
               (fx* (vector-ref index 1) (vector-ref stride 1)))
          (fx* (vector-ref index 2) (vector-ref stride 2))))
    ((4)
     (fx+ (fx+ (fx+ (fx* (vector-ref index 0) (vector-ref stride 0))
                    (fx* (vector-ref index 1) (vector-ref stride 1)))
               (fx* (vector-ref index 2) (vector-ref stride 2)))
          (fx* (vector-ref index 3) (vector-ref stride 3))))
    (else
      (transduce fixnum-range-fold
                 (map (lambda (axis)
                        (fx* (vector-ref index axis) (vector-ref stride axis))))
                 (collect-fx-sum)
                 (range 0 rank)))))

;; Gets an individual element at `index` from the `array`
(define-checked (array-ref (array array?) (index vector?))
  (let* ((interval (array-interval array))
         (start (interval-start interval))
         (rank (vector-length start))
         (offset-index
           (transduce fixnum-range-fold
                      (map
                        (lambda (i) (fx+ (vector-ref start i) (vector-ref index i))))
                      (collect-vector)
                      (range 0 rank)))
         (storage-object (array-storage-object array))
         (storage-class (array-storage-class array))
         (storage-ref (storage-class-ref storage-class))
         (stride (array-stride array)))
    (if (interval-contains? interval offset-index)
      (storage-ref storage-object (index->offset offset-index stride rank))
      (error 'array-ref "Index out of bounds" array index))))

;; Sets an individual element at `index` in the `array` to the provided
;; `value`.
(define-checked (array-set! (array array?) (index vector?) value)
  (let* ((interval (array-interval array))
         (start (interval-start interval))
         (rank (vector-length start))
         (offset-index
           (transduce fixnum-range-fold
                      (map
                        (lambda (i) (fx+ (vector-ref start i) (vector-ref index i))))
                      (collect-vector)
                      (range 0 rank)))
         (storage-object (array-storage-object array))
         (storage-class (array-storage-class array))
         (storage-set! (storage-class-set storage-class))
         (stride (array-stride array)))
    (if (interval-contains? interval offset-index)
      (storage-set! storage-object (index->offset offset-index stride rank) value)
      (error 'array-set! "Index out of bounds" array index))))

;; Updates the value at index as if applying `proc` to it.
(define-checked (array-update! (array array?) (index vector?) proc)
  (let* ((interval (array-interval array))
         (start (interval-start interval))
         (rank (vector-length start))
         (offset-index
           (transduce fixnum-range-fold
                      (map
                        (lambda (i) (fx+ (vector-ref start i) (vector-ref index i))))
                      (collect-vector)
                      (range 0 rank)))
         (storage-object (array-storage-object array))
         (storage-class (array-storage-class array))
         (storage-set! (storage-class-set storage-class))
         (storage-ref (storage-class-ref storage-class))
         (stride (array-stride array))
         (offset (index->offset index stride rank)))
    (if (interval-contains? interval offset-index)
      (storage-set! storage-object offset (proc (storage-ref storage-object offset)))
      (error 'array-update! "Index out of bounds" array index))))

;; Predicate which evaluates to `#t` iff the array is some kind of array view,
;; otherwise `#f`.
;;
;; An array view is any array that has been sliced or has had its axes
;; re-arranged in some way (transposition, swap-axes, squeeze, expand), while
;; maintaining the same overall interval rank.
(define-checked (array-view? (array array?))
  (not
    (and (fx=? (interval-length (array-interval array))
               (storage-object-length (array-storage-class array)
                                      (array-storage-object array)))
         (=? vector-index-comparator
             (array-stride array)
             (dimension->stride (interval-end (array-shape array)))))))

;; Predicate that evaluates if the `left` and `right` arrays are equal in the
;; sense that they:
;;
;; 0. Are both arrays. An error is signalled if at least one is not an array.
;; 1. Share the same shape and rank
;; 2. Are equivalent (in the sense of equal?) for all values contained within
;; the array.
;;
;; This is a very expensive predicate! One should take care not to use it in a
;; tight loop. If anything, this is mostly only intended for comparisons within
;; tests.
(define (array=? left right)
  (and (array? left)
       (array? right)
       (equal? (interval-end (array-shape left))
               (interval-end (array-shape right)))
       (transduce array-fold
                  (zip-array right)
                  (collect-all (lambda (pair) (equal? (car pair) (cdr pair))))
                  left)))

;; Returns a newly allocated array whose shape, index-class, and storage class
;; are the same as array and all of whose elements are `value`.
(define-checked (array-broadcast (array array?) value)
  (make-array (array-storage-class array) (interval-end (array-shape array)) value))