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

;;; Operations over arrays

;; It is at times helpful to be able to perform various array-specific
;; operations, or to perform operations against arrays similarly to how scalar
;; procedures operate. Examples include adding a scalar to an array, or a
;; Hadamard product.
;;
;; Writing the common forms for these is possible with the existing transducers
;; API, but leaves something to be desired when you have to do it many times
;; within a project. To that end, these act as convenience operations over
;; arrays.

;; Compares dimensions for two arrays and returns a partial ordering. The
;; ordering is:
;;
;; - GREATER: iff the rank of A is greater than the rank of B, and the B-rank
;;            of dimension A is equivalent to dimension B
;; - EQUAL:   iff dim-a and dim-b are exactly equal
;; - LESS:    iff the rank of A is less than the rank of B, and the A-rank of
;;            dimension B is equivalent to dimension A
;; - #f:      if the two dimensions do not satisfy any of the above criteria
;;            and are incompatible.
(define-checked (partial-compare-dimensions (dim-a vector?)
                                            (dim-b vector?))
  (let* ((rank-a (vector-length dim-a))
         (rank-b (vector-length dim-b))
         (subdim-is-equal
           (transduce fixnum-range-fold
                      values
                      (collect-all
                        (lambda (i) (eq? (vector-ref dim-a (fx- rank-a i))
                                         (vector-ref dim-b (fx- rank-b i)))))
                      (range 1 (fx+ 1 (min rank-a rank-b))))))
    (cond
      ((and (fx>? rank-a rank-b) subdim-is-equal)
       'GREATER)
      ((and (fx=? rank-a rank-b) subdim-is-equal)
       'EQUAL)
      ((and (fx<? rank-a rank-b) subdim-is-equal)
       'LESS)
      (else #f))))

;; The implementation for pointwise-extension of a scalar proc when two arrays
;; are provided.
(define (pointwise-extend-arrays-impl scalar-proc
                                      a
                                      b
                                      storage-class)
  ;; The assumption is that this is already called from pointwise-extend and so
  ;; does not make attempts to check if a and b are arrays
  (let* ((shape-a (array-shape a))
         (dimension-a (interval-end shape-a))
         (shape-b (array-shape b))
         (dimension-b (interval-end shape-b)))
    (case (partial-compare-dimensions dimension-a dimension-b)
      ('GREATER
       (let ((length-b (interval-length shape-b)))
         (transduce
           array-fold
           (compose
             (chunks-exact length-b (collect-vector length-b))
             (map
               (lambda (sub-a)
                 (transduce vector-fold
                            (compose
                              (zip-array b)
                              (map (lambda (pair) (scalar-proc (car pair) (cdr pair)))))
                            (collect-vector length-b)
                            sub-a)))
             flatten-vector)
           (collect-array storage-class dimension-a)
           a)))
      ('EQUAL
       (transduce array-fold
                  (compose
                    (zip-array b)
                    (map (lambda (pair) (scalar-proc (car pair) (cdr pair)))))
                  (collect-array storage-class dimension-a)
                  a))
      ('LESS
       (let ((length-a (interval-length shape-a)))
         (transduce
           array-fold
           (compose
             (chunks-exact length-a (collect-vector length-a))
             (map
               (lambda (sub-b)
                 (transduce vector-fold
                            (compose
                              (zip-array a)
                              ;; NOTE: The order of car / cdr is kept distinctly so that
                              ;; we preserve the a<>b ordering of the initial arguments
                              (map (lambda (pair) (scalar-proc (cdr pair) (car pair)))))
                            (collect-vector length-a)
                            sub-b)))
             flatten-vector)
           (collect-array storage-class dimension-b)
           b)))
      ;; If partial-compare-dimensions returns #f we know that the dimensions
      ;; are incompatible.
      (else
        (error `(pointwise-extend ,scalar-proc)
               "Array dimensions are incompatible"
               shape-a
               shape-b)))))

;; Produces a dyadic (arity 2) procedure that is the pointwise extension of a
;; dyadic (arity 2) scalar procedure.
;;
;; The resulting procedure will accept any two scalars or arrays that have
;; compatible dimensions and will perform the scalar procedure across those
;; dimensions. Two arrays have compatible dimensions iff:
;;
;; 1. The shape of the arrays is the exact same
;; 2. The shape of one of the arrays is a lower rank than the other, and the
;;    sub-dimension of the first array is compatible with the second. e.g.
;;    #(5 1) would be compatible with the sub-dimension of #(3 5 1).
;; 3. One or two of the arguments are scalars, and not arrays.
;;
;; The `scalar-proc` must be a procedure with at least arity of 2. It is
;; applied with the ordering such that all operations will respect similar
;; ordering to the caller. For example, if your scalar-proc was string-append,
;; you would not see arguments re-ordered, but instead the resulting procedure
;; produced by pointwise-extend (which takes two arguments, lets say A and B)
;; would apply the scalar procedure such that individual array elements would
;; be ordered as though the procedure was called as `(scalar-proc a b)`.
;;
;; The returned procedure takes an optional 3rd argument, which is the storage
;; class that the final array (as a result of the extended operation) should
;; use. This defaults to the storage class of the first argument (if it is an
;; array), and then the storage-class of the second argument, if the first is
;; not an array.
(define-checked (pointwise-extend (scalar-proc procedure?))
  (letrec ((extended-proc
             (case-lambda
               ((a b)
                (extended-proc a b (or (and (array? a)
                                            (array-storage-class a))
                                       (and (array? b)
                                            (array-storage-class b)))))
               ((a b storage-class)
                (let ((a-is-array (array? a))
                      (b-is-array (array? b)))
                  (cond
                    ;; CASE: both 'A' and 'B' are arrays
                    ((and a-is-array b-is-array)
                     (pointwise-extend-arrays-impl scalar-proc a b storage-class))
                    ;; CASE: 'A' is an array and 'b' is a scalar
                    ((and a-is-array (not b-is-array))
                     (transduce array-fold
                                (map
                                  (lambda (element)
                                    (scalar-proc element b)))
                                (collect-array storage-class
                                               (interval-end (array-shape a)))
                                a))
                    ;; CASE: 'a' is a scalar and 'B' is an array
                    ((and b-is-array (not a-is-array))
                     (transduce array-fold
                                (map
                                  (lambda (element)
                                    (scalar-proc a element)))
                                (collect-array storage-class
                                               (interval-end (array-shape b)))
                                b))
                    ;; CASE: both 'a' and 'b' are scalars
                    (else
                      (scalar-proc a b))))))))
    extended-proc))