;;;; Copyright (c) 2017, 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.

(define-record-type array
  (%make-array storage-class
               shape
               stride
               offset
               mutable?
               storage-object)
  array?
  (storage-class %array-storage-class : (struct storage-class))
  (shape %array-shape (setter %array-shape) : (vector-of fixnum))
  (stride %array-stride (setter %array-stride) : (vector-of fixnum))
  (offset %array-offset : fixnum)
  (mutable? %array-mutable? (setter %array-mutable?) : boolean)
  (storage-object %array-storage-object))

(define-check+error-type array array?)

(define-check+error-type
  array-index
  (lambda (obj)
    (and (vector? obj)
         (vector-every fixnum? obj))))

(define *UNSPECIFIED* (list 'unspecified))

;; 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 (shape->stride shape #!optional (make-row-major? #t))
  (if (and (fx> (vector-length shape) 2)
           (fx= 1 (vector-ref shape
                              (sub1 (vector-length shape)))))
    (let* ((total (vector-fold fx* 1 shape))
           (stride (vector-map (cute fx/ total <>)
                               (vector-cumulate fx* 1 shape))))
      (unless make-row-major? (vector-reverse! stride))
      stride)
    (let ((f (vector-ref shape 0)))
     (let ((stride (vector-map (cute fx/ <> f)
                               (vector-cumulate fx* 1 shape))))
       (when make-row-major? (vector-reverse! stride))
       stride))))

(define (row-major? array)
  (apply >= (vector->list (%array-stride array))))


;; Quick and dirty way to get the "rank" of an array. Rank here refers to the
;; same notation as in Fortran or C, not actual mathematical rank.
(define (%array-rank array)
  (vector-length (%array-stride array)))

(define (%array-lower-bound array)
  (make-vector (%array-rank array) 0))

(define (%array-upper-bound array)
  (vector-copy (%array-shape array)))

;; Converts an array index to a storage index.
;; This API is unsafe, see the public version without the % prefix for more
;; information.
(define-inline (%array-index->storage-index array index)
  (fx+ (%array-offset array)
       (match (list (%array-stride array) index)
         ;; Rank-1 Case
         ((#(si) #(i)) (fx* si i))
         ;; Rank-2 Case
         ((#(si sj) #(i j))
          (fx+ (fx* si i)
               (fx* sj j)))
         ;; Rank-3 Case
         ((#(si sj sk) #(i j k))
          (fx+ (fx* si i)
               (fx+ (fx* sj j)
                    (fx* sk k))))
         ;; Rank-4 Case
         ((#(si sj sk sv) #(i j k v))
          (fx+ (fx* si i)
               (fx+ (fx* sj j)
                    (fx+ (fx* sk k)
                         (fx* sv v)))))
         ;; Rank-N Case
         (_ (vector-fold fx+ 0 (vector-map fx* (%array-stride array) index))))))



;; Increments an index up to the upper-bound.
;; Some examples assume upper bound of #(3 3 3):
;;
;; #(0 0 0) -> #(0 0 1)
;; #(0 0 1) -> #(0 0 2)
;; #(0 0 2) -> #(0 1 0)
;;
;; Note how in every case it increments up to the upper bound of a given
;; dimension, but never to it. If the index cannot be incremented any further
;; beyond its max value, then the upper bound is returned.
(define (%increment-index index upper-bound)
  (let loop ((dim (sub1 (vector-length index)))
             (index index)
             (upper-bound (vector-copy upper-bound)))
    (cond
      ((or (fx< dim 0)
           (fx>= (vector-ref index 0)
                 (vector-ref upper-bound 0)))
       upper-bound)
      ((fx= (add1 (vector-ref index dim))
            (vector-ref upper-bound dim))
       (vector-set! index dim 0)
       (loop (sub1 dim) index upper-bound))
      (else
        (vector-set! index dim (add1 (vector-ref index dim)))
        index))))

;; Executes proc on every possible index value of array, in lexicographic
;; order. e.g. (%array-for-each-index print <some-array>) prints all the
;; indices in lexicographic order.
(define-inline (%array-for-each-index proc start end)
   (match (list start end)
    ;; Rank-1 Case
    ((#(li) #(ui))
     (let i-loop ((i li))
      (when (< i ui)
        (proc (vector i))
        (i-loop (add1 i)))))
    ;; Rank-2 Case
    ((#(li lj) #(ui uj))
     (let i-loop ((i li))
      (when (< i ui)
        (let j-loop ((j lj))
         (if (< j uj)
           (begin
             (proc (vector i j))
             (j-loop (add1 j)))
           (i-loop (add1 i)))))))
    ;; Rank-3 Case
    ((#(li lj lk) #(ui uj uk))
     (let i-loop ((i li))
      (when (< i ui)
        (let j-loop ((j lj))
         (if (< j uj)
           (let k-loop ((k lk))
            (if (< k uk)
              (begin
                (proc (vector i j k))
                (k-loop (add1 k)))
              (j-loop (add1 j))))
           (i-loop (add1 i)))))))
    ;; Rank-4 Case
    ((#(li lj lk lv) #(ui uj uk uv))
     (let i-loop ((i li))
      (when (< i ui)
        (let j-loop ((j lj))
         (if (< j uj)
           (let k-loop ((k lk))
            (if (< k uk)
              (let v-loop ((v lv))
               (if (< v uv)
                 (begin
                   (proc (vector i j k v))
                   (v-loop (add1 v)))
                 (k-loop (add1 k))))
              (j-loop (add1 j))))
           (i-loop (add1 i)))))))
    ;; Rank-N Case
    (((and start (? vector?) (and end (? vector?))))
     (let loop ((index (vector-copy start))
                (upper-bound (vector-copy end)))
       (cond
         ((vector= fx= index upper-bound) (void))
         (else
           (proc index)
           (loop (%increment-index index upper-bound) upper-bound)))))
    ;; Default rule
    (_ (error "start and end must be vectors of the same rank." start end))))


(define-inline (%array-ref array index)
  (let ((storage-index (%array-index->storage-index array index))
        (ref (storage-class-accessor (%array-storage-class array)))
        (storage-object (%array-storage-object array)))
    (ref storage-object storage-index)))

;;; Constructors

;; Returns a newly allocated mutable array (with a newly allocated storage
;; object of the specified storage-class) that has the specified bounds. The
;; values of the elements default to the default fill value of the storage
;; class if the fill argument is omitted.
(define (make-array storage-class shape #!optional (fill *UNSPECIFIED*))
  (check-storage-class 'make-array storage-class)
  (let* ((nelems (vector-fold * 1 shape))
         (stride (shape->stride shape))
         (storage-object
           (make-storage-object storage-class
                                nelems
                                (if (eq? fill *UNSPECIFIED*)
                                  (storage-class-default-fill storage-class)
                                  fill))))
    (%make-array storage-class
                 shape
                 stride
                 0
                 #t
                 storage-object)))


;; Iterates over the indexes of array in lexicographical order starting at the
;; index start and ending at the index end (start inclusive, end exclusive),
;; and calling proc on each index. Whatever proc returns becomes the value of
;; the array at the index.
(define (array-tabulate! proc array)
  (unless (array-mutable? array)
    (error "(array-tabulate!): Cannot mutate array if array is not mutable." array))
  (let* ((array array)
         (storage-set! (storage-class-mutator
                         (%array-storage-class array)))
         (data (%array-storage-object array)))
    (%array-for-each-index
      (lambda (i)
        (storage-set!
          data
          (%array-index->storage-index array i)
          (proc i)))
      (%array-lower-bound array)
      (%array-upper-bound array))
    array))


;; Returns a newly allocated array (with a newly allocated storage object of
;; the specified storage-class) that has the specified bounds. The values of
;; the elements are computed by calling proc on every possible index of the
;; array in lexicographic order. If mutable? is true, the array can mutate its
;; storage object.
(define (array-tabulate storage-class proc shape mutable?)
  (let* ((array (make-array storage-class shape)))
    (array-tabulate! proc array)
    (set! (%array-mutable? array) mutable?)
    array))

;; Returns a newly allocated array whose bounds and storage class are the same
;; as array and all of whose elements are obj.
(define (array-broadcast array obj)
  (check-array 'array-broadcast array)
  (make-array (%array-storage-class array) (%array-shape array) obj))

;;; Predicates

;; Returns #t if obj is an array, and #f otherwise.
;; Defined in the record definition at the top.
;; (define (array? obj))

;; Returns #t if array can mutate its storage object, and #f otherwise.
(define (array-mutable? array)
  (check-array 'array-mutable? array)
  (%array-mutable? array))

;;; Metadata accessors

;; Returns the storage class with which array was created.
(define (array-storage-class array)
  (check-array 'array-storage-class array)
  (%array-storage-class array))

;; Returns the storage object underlying this array. Note that this can be #f
;; in the case of storage classes without actual storage.
(define (array-storage-object array)
  (check-array 'array-storage-object array)
  (%array-storage-object array))

;; Return the rank (number of dimensions) of array.
(define (array-rank array)
  (check-array 'array-rank array)
  (%array-rank array))

(define (array-shape array)
  (check-array 'array-shape array)
  (vector-copy (%array-shape array)))

;; Return the stride of array.
(define (array-stride array)
  (check-array 'array-stride array)
  (vector-copy (%array-stride array)))

;; Returns the offset of array.
(define (array-offset array)
  (check-array 'array-offset array)
  (%array-offset array))

;; Converts the index of array to the equivalent storage index.
(define (array-index->storage-index array index)
  (check-array 'array-index->storage-index array)
  (check-array-index 'array-index->storage-index index)
  (let ((rank (%array-rank array)))
   (unless (= (vector-length index) rank)
     (error "Incorrect number of dimensions in index for array rank."
            rank
            index))
   (%array-index->storage-index array index)))

;;; Accessors

;; Returns the value of the element of array specified by index. Note that this
;; is different from the array-ref of most Lisp systems, which specifies the
;; index as a sequence of arguments.
(define (array-ref array index)
  (check-array 'array-ref array)
  (check-array-index 'array-ref index)
  (unless (eq? (vector-length index)
               (%array-rank array))
    (error "(array-ref): index must have as many dimensions as array rank" index))
  (unless (vector= < index (%array-shape array))
    (error "(array-ref): index must be less than the max shape along all axes"
           index
           (%array-shape array)))
  (%array-ref array index))

;; Applies array-ref to the array using the first index. If there are more
;; arguments, apply array-ref to the result using the next index. Repeat until
;; there are no more indexes, returning the last result. It is an error if any
;; intermediate result is not an array. (APL enlist.)
(define (array-recursive-ref array index #!rest indices)
  (fold (flip array-ref) array (cons index indices)))

;; Iterates over the elements of array in lexicographic order, starting at the
;; index start and ending just before the index end, and calling proc on each
;; element. Each invocation of proc receives the value of the element at that
;; index. The value returned by proc is discarded.
(define (array-for-each proc array)
  (check-array 'array-for-each array)
  (%array-for-each-index
    (lambda (i) (proc (%array-ref array i)))
    (%array-lower-bound array)
    (%array-upper-bound array)))

;; Iterates over the indexes of array in lexicographic order, starting at the
;; index start and ending at the index end, and calling proc on each index. The
;; value returned by proc is discarded.
(define (array-for-each-index proc array)
  (check-array 'array-for-each-index array)
  (%array-for-each-index
    proc
    (%array-lower-bound array)
    (%array-upper-bound array)))

;;; Mutators

;; These procedures return an unspecified value.

(define-inline (%array-set! array index value)
  (let ((storage-index (%array-index->storage-index array index))
        (storage-set! (storage-class-mutator (%array-storage-class array)))
        (storage-object (%array-storage-object array)))
    (storage-set! storage-object storage-index value)))

;; Sets the value of the element of array specified by index to value. Note
;; that this is different from the array-set! of most Lisp systems, which
;; specifies the index as a sequence of arguments.
(define (array-set! array index value)
  (check-array 'array-set! array)
  (check-array-index 'array-set! index)
  (unless (array-mutable? array)
    (error "(array-set!): Cannot mutate array if array is not mutable." array))
  (unless (vector= < index (%array-shape array))
    (error "(array-set!): index must be less than the max shape along all axes"
           index
           (%array-shape array)))
  (%array-set! array index value))