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

;; NOTE: Because of the way I include the files in the top-level module, some
;; definitions for helper procedures won't be used here. Check
;; `generalized-arrays/generalized-arrays-internal.scm` for these procedures /
;; record-types, as they should be located there.

;;; The whole array

;; Returns a newly allocated array with the same bounds as the arrays; it is an
;; error if they do not all have the same bounds. For each valid index value,
;; proc is invoked, passing each corresponding element of the arrays. Whatever
;; proc returns becomes the value of the storage element corresponding to that
;; index in the result array. The order of invocations of proc is not
;; specified.
(define (array-map storage-class proc array #!rest arrays)
  (check-array 'array-map array)
  (check-storage-class 'array-map storage-class)
  (for-each (lambda (a) (check-array 'array-map a))
            arrays)
  (%array-map storage-class proc array arrays))


;; It is an error if the arrays do not all have the same bounds. For each valid
;; index value, proc is invoked, passing each corresponding element of the
;; arrays.  Whatever proc returns becomes the value of the storage element
;; corresponding to that index in the first specified array. The order of
;; invocations of proc is not specified. Returns an undefined value.
(define (array-map! proc array #!rest arrays)
  (check-array 'array-map! array)
  (for-each (lambda (x) (check-array 'array-map! x))
            arrays)
  (unless (array-mutable? array)
    (error "(array-map!): immutable arrays cannot be mapped using array-map!"))
  (%array-map! proc array arrays))

;; Returns a newly allocated array with the same bounds as the arrays; it is an
;; error if they do not all have the same bounds. For each valid index value,
;; proc is invoked in lexicographic order, passing each corresponding element
;; of the arrays and a seed, whose initial value is seed. Proc returns two
;; values, the value of the storage element corresponding to that index in the
;; result array, and the new value of the seed.
(define (array-fold proc seed array #!rest arrays)
  (check-array 'array-map! array)
  (for-each (lambda (x) (check-array 'array-map! x))
            arrays)
  (%array-fold proc seed array arrays))


;; Returns a newly allocated array whose rank is one less than the rank of
;; array, by combining all the groups of elements of length n (where the
;; default is the extent of axis) along axis using proc, a two-argument
;; procedure. The order and number of invocations of proc is unspecified. If
;; there is only one such element, it is unchanged. (APL reduce.)
(define (array-reduce proc seed array axis)
  (check-array 'array-reduce array)
  (check-natural-fixnum 'array-reduce axis 'axis)
  (unless (and (fx>= axis 0)
               (fx< axis (%array-rank array)))
    (error "(array-reduce): axis must be between 0 and the array rank." axis))
  (%array-reduce proc seed array axis))


;; Returns a newly allocated array whose bounds are the same as the bounds of
;; array. Each element along axis is constructed by reducing (as if by
;; array-reduce) successive prefixes of the elements of array along that axis.
;; (APL scan.)
(define (array-cumulate proc seed array axis)
  (check-array 'array-cumulate array)
  (check-natural-fixnum 'array-cumulate axis 'axis)
  (unless (and (fx>= axis 0)
               (fx< axis (%array-rank array)))
    (error "(array-cumulate): axis must be between 0 and the array rank." axis))
  (%array-cumulate proc seed array axis))


;; Returns an exact integer containing the number of elements in array that
;; satisfy pred.
(define (array-count pred array)
  (array-fold (lambda (len x)
                (if (pred x)
                  (add1 len)
                  len))
              0
              array))


;; Returns the index of the first element of array in lexicographic order that
;; satisfies pred, or #f if no index is found.
(define (array-index pred array)
  (check-array 'array-index array)
  (call/cc
    (lambda (k)
      (array-for-each-index
        (lambda (i)
          (when (pred (%array-ref array i))
            (k i)))
        array)
      (k #f))))


;; Returns the first element of array in lexicographic order that satisfies
;; pred, or #f if no element is found.
(define (array-find pred array)
  (check-array 'array-find array)
  (let ((index (array-index pred array)))
    (if index
      (array-ref array index)
      #f)))


;; Returns an array with the same bounds as array except possibly along the
;; axis dimension. The array is sliced along axis and the elements of booleans
;; (a vector of boolean values) are used to decide which slices to incorporate
;; into the result: if the corresponding boolean is #t, the slice is
;; incorporated, otherwise not. (APL compress.)
(define (array-compress array booleans axis)
  (check-array 'array-compress array)
  (check-vector 'array-compress booleans)
  (unless (and (fx>= axis 0)
               (fx< axis (%array-rank array)))
    (error "(array-compress): axis must be between 0 and the array rank." axis))
  (unless (fx= (vector-length booleans)
               (vector-ref (%array-shape array) axis))
    (error "(array-compress): booleans must be same size as axis length"
           (vector-length booleans)
           axis))
  (%array-compress array booleans axis))


;; Returns an array with the same bounds as array except possibly along the
;; axis dimension. Array is sliced along axis and the elements of booleans (a
;; vector of boolean values) are used to decide where, if anywhere, nil is to
;; be interpolated: if the corresponding boolean is #t, nil is interpolated,
;; otherwise the next slice is incorporated. It is an error if the size of
;; booleans is not to the extent of the axis dimension in the result. It is
;; also an error if nil does not have the same bounds as a slice. (APL expand.)
(define (array-expand array booleans nil axis)
  (check-array 'array-expand array)
  (check-array 'array-expand nil)
  (unless (vector= =
                   (%array-shape nil)
                   (remove-axis-from-index (%array-shape array) axis))
    (error "(array-expand): nil must be same shape as array along axis."
           (%array-shape nil)
           (remove-axis-from-index (%array-shape array) axis)))
  (check-vector 'array-expand booleans)
  (unless (and (fx>= axis 0)
               (fx< axis (%array-rank array)))
    (error "(array-expand): axis must be between 0 and the array rank." axis))
  (unless (fx= (vector-length booleans)
               (vector-ref (%array-shape array) axis))
    (error "(array-extend): booleans must be same size as axis length"
           (vector-length booleans)
           axis))
  (%array-expand array booleans nil axis))


;; Returns an array with the same bounds as array. Array is sliced along the
;; axis dimension, and the slices are reassembled in the order given by vector.
;; The slice whose number appears in the first element of vector appears first
;; in the result, and so on. It is an error if vector is not a vector of exact
;; integers.  (Generalized version of APL rotate.)
(define (array-rearrange array order axis)
  (check-array 'array-rearrange array)
  (check-vector 'array-rearrange order)
  (check-natural-fixnum 'array-rearrange axis)
  (unless (fx< axis (%array-rank array))
    (error "(array-rearrange): axis must be between 0 and the array rank." axis))
  (unless (and (fx= (vector-ref (%array-shape array) axis)
                    (vector-length order))
               (vector-every (lambda (x)
                               (fx< x (vector-ref (%array-shape array) axis)))
                             order))
    (error "(array-rearrange): All values in the order vector must be less than the axis length."))
  (%array-rearrange array order axis))


;;; Transformations

;; The procedure proc specifies an affine function that returns an index of
;; array when given an index of the returned array. The array does not retain a
;; dependency on proc. (SRFI 25 share-array.)
(define (array-transform proc shape array)
  (check-array 'array-transform array)
  (check-vector 'array-transform shape)
  (%array-transform proc shape array))

;; Returns an array with the specified bounds whose elements in lexicographic
;; order are the same as the elements of array in lexicographic order. It is an
;; error if there are too many or too few elements. (APL reshape.)
(define (array-reshape array shape)
  (check-array 'array-reshape array)
  (check-vector 'array-reshape shape)
  (unless (= (vector-fold fx* 1 (%array-shape array))
             (vector-fold fx* 1 shape))
    (error "(array-reshape): new shape must have equivalent number of elements as old shape."
           (vector-fold fx* 1 shape)
           (vector-fold fx* 1 (%array-shape array))))
  (let ((reshaped-array (array-copy array (%array-mutable? array))))
   (set! (%array-shape reshaped-array) shape)
   (set! (%array-stride reshaped-array)
         (shape->stride shape (row-major? array)))
   reshaped-array))

;; Returns a one-dimensional array which contains the diagonal elements of
;; array (that is, the elements whose indices are all the same integer).
(define (array-diagonal array)
  (check-array 'array-diagonal array)
  (if (= 1 (%array-rank array))
    array
    (%make-array (%array-storage-class array)
                 (vector (apply min (vector->list (%array-shape array))))
                 (vector (vector-fold fx+ 0 (%array-stride array)))
                 (%array-offset array)
                 (%array-mutable? array)
                 (%array-storage-object array))))

;; Returns an array with the same bounds as array, but whose elements on the
;; specified axis are reversed. (APL reverse.)
(define (array-reverse array axis)
  (check-array 'array-reverse array)
  (check-natural-fixnum 'array-reverse axis)
  (unless (and (fx>= axis 0)
               (fx< axis (%array-rank array)))
    (error "(array-rearrange): axis must be between 0 and the array rank." axis))
  (let ((limit (sub1 (vector-ref (%array-shape array) axis))))
   (array-transform (lambda (index)
                      (let ((j (vector-copy index)))
                       (vector-set! j
                                    axis
                                    (fx- limit (vector-ref j axis)))
                       j))
                    (array-shape array)
                    array)))

;; Returns an array whose axes appear in the reverse order of the axes of
;; array.  This implies that the upper and lower bound are the reverse of the
;; bounds of array. (APL monadic transpose.)
(define (array-transpose array)
  (check-array 'array-transpose array)
  (let ((shape (vector-copy (%array-shape array)))
        (stride (vector-copy (%array-stride array))))
   (vector-reverse! shape)
   (vector-reverse! stride)
   (%make-array (%array-storage-class array)
                shape
                stride
                (%array-offset array)
                (%array-mutable? array)
                (%array-storage-object array))))


;; Returns an array whose axes are an arbitrary permutation of the axes of
;; array.  Vec specifies how to do the permutation: the axis whose number
;; appears in the first element of vector appears as the first axis of the
;; result, and so on. (APL dyadic transpose with integer-valued vector.)
(define (array-rearrange-axes array order)
  (check-array 'array-rearrange-axes array)
  (check-vector 'array-rearrange-axes order)
  (unless (= (vector-length order)
             (%array-rank array))
    (error "(array-rearrange-axes): order must be same length as rank of array."
           (vector-length order)
           (%array-rank array)))
  (unless (vector= fx=
                   (sort order fx<)
                   (vector-unfold values (%array-rank array)))
    (error "(array-rearrange-axes): all axes must be present in order vector."
           order))
  (%array-rearrange-axes array order))


;; Returns a smaller array with the same rank as array whose elements begin at
;; the "lower left" corner specified by start and end at the "upper right"
;; corner specified byend. Unlike array-copy, the result shares its storage
;; object with array. (APL take and drop.)
(define (array-slice array start end #!optional (step *UNSPECIFIED*))
  (check-array 'array-slice array)
  (check-array-index 'array-slice start)
  (check-array-index 'array-slice end)
  (unless (and (vector= fx>= start (%array-lower-bound array))
               (vector= fx< start (%array-upper-bound array))
               (vector= fx< start end))
    (error
      "(array-slice): start index must be greater-than-or-equal to the lower bound and less than the end / upper-bound."
      start
      end
      (%array-lower-bound array)
      (%array-upper-bound array)))
  (unless (vector= fx<= end (%array-upper-bound array))
    (error "(array-slice): end index cannot be greater than the upper-bound of array."
           end
           (%array-shape array)))
  (let ((step (if (eq? step *UNSPECIFIED*)
                (vector-map (lambda _ 1) start)
                step)))
    (unless (vector= >= step (vector-map (lambda _ 0) start))
      (error "(array-slice): step size must be greater than or equal to zero." step))
    (%array-slice array start end step)))


;; Returns an array with the ranks specified by the elements of vector removed
;; from array. It is an error if the extents of the specified ranks are not
;; equal to 1.  (NumPy squeeze)
(define (array-squeeze array #!optional (axes *UNSPECIFIED*))
  (check-array 'array-squeeze array)
  (if (eq? axes *UNSPECIFIED*)
    (let* ((shape (%array-shape array))
           (axes (list->vector
                  (filter (lambda (x) (= 1 (vector-ref shape x)))
                          (iota (vector-length shape))))))
      (%array-squeeze array axes))
    (begin
      (vector-for-each
        (lambda (x)
          (unless (= 1 (vector-ref (%array-shape array) x))
            (error "(array-squeeze): Can only squeeze axes with length of 1.")))
        axes)
      (%array-squeeze array axes))))


;; Returns an array whose rank is one greater than the rank of array. This is
;; accomplished by inserting a new axis numbered 'axis' whose extent is (0:1).
;; (NumPy expand.)
(define (array-unsqueeze array rank)
  (check-array 'array-unsqueeze array)
  (check-natural-fixnum 'array-unsqueeze rank)
  (unless (fx<= rank (%array-rank array))
    (error "(array-unsqueeze): rank must be a value in [0, n] for an n-ranked array."
           rank))
  (%array-unsqueeze array rank))


;;; Copying and conversion

;; These procedures return arrays which do not share their storage objects with
;; any existing arrays.

;; When these procedures return arrays, both the array and the underlying
;; storage object are newly allocated

;; Returns an array containing the elements of array.
;; The stride and offset are implementation-defined. The resulting array is
;; mutable if mutable? is true.  The returned array has the same storage class
;; as array.
(define (array-copy array #!optional (mutable? #t))
  (check-array 'array-copy array)
  (check-boolean 'array-copy mutable?)
  (%array-copy array mutable?))


;; Copies the elements of array from the "to" argument to the "from" argument
;; starting at element "at".  It is an error if there are not enough elements
;; in "to" to make this possible. The storage objects of from and to need not
;; belong to the same storage classes.
(define (array-copy! from to)
  (check-array 'array-copy! from)
  (check-array 'array-copy! to)
  (unless (array-mutable? to)
    (error "(array-copy!): Specified 'to' array is immutable."))
  (unless (vector= >= (%array-shape to) (%array-shape from))
    (error "There is insufficient space available to copy into array."
           (%array-shape to)
           (%array-shape from)))
  (%array-copy! from to))


;; Returns a newly allocated array consisting of the arrays concatenated along
;; axis. It is an error unless the storage classes of the arrays are the same,
;; and the result has the same storage class. It is also an error unless the
;; bounds of all the arrays are the same, with the possible exception of axis.
;; The axisth element of the lower bound of the result is 0; the corresponding
;; element of the upper bound is the sum of the extents of the arrays.
(define (array-append axis array #!rest arrays)
  (check-array 'array-append array)
  (check-natural-fixnum 'array-append axis)
  (for-each
   (lambda (a)
     (check-array 'array-append a)
     (unless (vector= =
                      (remove-axis-from-index (%array-shape array) axis)
                      (remove-axis-from-index (%array-shape a) axis))
       (error "(array-append): shapes must be congruent for all arrays."
              (%array-shape a)
              (%array-shape array))))
   arrays)
  (unless (or (fx>= axis 0)
              (fx< axis (%array-rank array)))
    (error "(array-append): axis must be between 0 and array rank."))
  (%array-append (%array-storage-class array) axis (cons array arrays)))


;; Append repeat copies of array along axis axis, as if by array-append.
;; (Variant of NumPy repeat.)
(define (array-repeat array axis repeat)
  (apply array-append axis (make-list repeat array)))

;; Returns an array with the same bounds and elements as array, but whose
;; storage class is specified by storage-class.
(define (array-reclassify array storage-class)
  (check-array 'array-reclassify array)
  (check-storage-class 'array-reclassify storage-class)
  (let ((new-array (make-array storage-class (%array-shape array))))
    (array-copy array new-array)
    array))

;; Returns a newly allocated vector/list whose elements are also newly
;; constructed vectors/lists, and so on as far down as necessary to cover every
;; axis of the array. Bottom-level Scheme vectors/lists contain the elements of
;; array. Thus, if array has rank 1, the result is a vector/list; if the array
;; has rank 2, the result is a vector/list whose elements are vectors/lists,
;; and so on. As a special case, if array has rank 0, the sole element is
;; returned.
(define (array->nested-vector array)
  (check-array 'array->nested-vector array)
  (let loop ((subarray array))
    (let ((vec (make-vector (vector-ref (%array-shape subarray) 0))))
     (if (= (array-rank subarray) 1)
       (for-each (lambda (i)
                   (vector-set! vec i (%array-ref subarray (vector i))))
                 (iota (vector-length vec)))
       (let ((lower-bound (vector-copy (%array-lower-bound subarray)))
             (upper-bound (vector-copy (%array-shape subarray))))
        (for-each
         (lambda (i)
           (vector-set! lower-bound 0 i)
           (vector-set! upper-bound 0 (add1 i))
           (vector-set! vec
                        i
                        (loop
                         (%array-squeeze (array-slice subarray
                                                      lower-bound
                                                      upper-bound)
                                         (vector 0)))))
         (iota (vector-length vec)))))
     vec)))

(define (array->nested-list array)
  (check-array 'array->nested-list array)
  (let loop ((subarray array))
    (let ((dim-length (vector-ref (%array-shape subarray) 0)))
      (if (= (%array-rank subarray) 1)
        (map (compose (cute %array-ref subarray <>) vector)
             (iota dim-length))
        (let ((lower-bound (vector-copy (%array-lower-bound subarray)))
              (upper-bound (vector-copy (%array-shape subarray))))
          (map (lambda (i)
                 (vector-set! lower-bound 0 i)
                 (vector-set! upper-bound 0 (add1 i))
                 (loop
                  (array-squeeze (array-slice subarray
                                              lower-bound
                                              upper-bound)
                                 (vector 0))))
           (iota dim-length)))))))


;; Returns a newly allocated array of the storage-class with rank rank whose
;; elements are initialized from the vector nested-vector or list nested-list.
;; It is an error if this argument is not rectangular. As a special case, if
;; rank is 0, the sole element is nested-vector or nested-list, which in that
;; case need not be a Scheme vector/list.
(define (nested-vector->array nested-vector storage-class rank)
  (check-vector 'nested-vector->array nested-vector)
  (check-storage-class 'nested-vector->array storage-class)
  (check-positive-fixnum 'nested-vector->array rank)
  (let loop ((vec nested-vector)
             (shape '())
             (i 0))
    (cond
      ((fx>= i rank)
       (array-tabulate storage-class
                       (lambda (index)
                         (vector-fold vector-ref nested-vector index))
                       (list->vector (reverse shape))
                       #t))
      (else
        (loop (vector-ref vec 0)
              (cons (vector-length vec) shape)
              (add1 i))))))


(define-inline (flatten-by-rank rank lst)
  (let loop ((i 1)
             (xs lst))
    (if (< i rank)
      (loop (add1 i) (concatenate xs))
      xs)))

(define (nested-list->array nested-list storage-class rank)
  (check-list 'nested-list->array nested-list)
  (check-storage-class 'nested-list->array storage-class)
  (check-positive-fixnum 'nested-list->array rank)
  (let loop ((lst nested-list)
             (shape '())
             (i 0))
    (cond
      ((fx>= i rank)
       (let* ((flat-list (flatten-by-rank rank nested-list))
              (vshape (list->vector (reverse shape)))
              (storage-object
               (make-storage-object storage-class
                                    (length flat-list))))
         (unless (= (storage-object-length storage-class
                                           storage-object)
                    (vector-fold fx* 1 vshape))
           (error "(nested-list->array): the lengths along each dimension in nested-list must be the same."))
         (for-each (lambda (elem i)
                     ((storage-class-mutator storage-class)
                      storage-object
                      i
                      elem))
                   flat-list
                   (iota (storage-object-length storage-class
                                                storage-object)))
         (%make-array storage-class
                      vshape
                      (shape->stride vshape)
                      0
                      #t
                      storage-object)))
      ((and (fx< (add1 i) rank)
            (not (pair? (car lst))))
       (error "Rank is inconsistent with array shape."
              rank
              (reverse (cons (length lst) shape))))
      (else
        (loop (car lst)
              (cons (length lst) shape)
              (add1 i))))))