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

(test-group "Construct arrays with make-array"
  (test-assert "make-array returns something that passes array?"
               (array? (make-array vector-storage-class #(480 640) 0)))

  (test-assert "make-array without fill returns something that passes array?"
               (array? (make-array vector-storage-class #(480 640))))

  (test-assert "make-array fills array with 'fill' value"
               (transduce array-fold
                          values
                          (collect-all (lambda (val) (eq? val 999)))
                          (make-array vector-storage-class #(1000) 999)))

  (test-error "make-array with a fill that can't be used with storage class is an error"
              (make-array u8vector-storage-class #(1000 10) "a-string")))

(test-group "Construct arrays with make-array-from-storage"
  (test-assert "make-array-from-storage returns something that passes array?"
               (array? (make-array-from-storage
                         vector-storage-class
                         #(3 3)
                         #(0 1 2 3 4 5 6 7 8)) ))

  (test-error "make-array-from-storage errors if dimension and length mismatch."
              (make-array-from-storage
                         vector-storage-class
                         #(3 3 3)
                         #('a 'b 'c))))

(test-group "Construct arrays with array-tabulate"
  (test-assert "array-tabulate returns something that passes array?"
               (array? (array-tabulate vector-storage-class #(100 100) (lambda (i) 'a)))))

(test-group "Construct arrays with array-broadcast"
  (define array (make-array vector-storage-class #(100 100) 999))
  (define bcast (array-broadcast array -30))

  (test-assert "array-broadcast returns something that passes array?"
               (array? bcast))

  (test-assert "array-broadcast does not modify initial array"
               (transduce array-fold
                          values
                          (collect-all (lambda (i) (eq? i 999)))
                          array))

  (test "Shape of broadcasted array matches original array"
        (array-shape array)
        (array-shape bcast))

  (test-assert "All values in broadcasted array are set to the fill obj"
               (transduce array-fold
                          values
                          (collect-all (lambda (i) (eq? i -30)))
                          bcast))

  (test-error "array-broadcast with bad-fill is an error"
              (array-broadcast (make-array u8vector-storage-class #(10 10) 2)
                               "a-string")))

(test-group "Array shape"
  (define dimension #(480 640))
  (define array (make-array vector-storage-class dimension))

  (test-assert "array-shape returns an interval"
               (interval? (array-shape array)))

  (test "array-shape starts at default index"
        (make-vector 2 0)
        (interval-start (array-shape array)))

  (test "array-shape ends at dimension"
        dimension
        (interval-end (array-shape array))))

(test-group "Array rank"
  (define array-1 (make-array vector-storage-class #(10)))
  (define array-2 (make-array vector-storage-class #(10 10)))
  (define array-3 (make-array vector-storage-class #(10 10 10)))
  (define array-4 (make-array vector-storage-class #(10 10 10 10)))

  (test "array-1 has rank 1"
        1
        (array-rank array-1))

  (test "array-2 has rank 2"
        2
        (array-rank array-2))

  (test "array-3 has rank 3"
        3
        (array-rank array-3))

  (test "array-4 has rank 4"
        4
        (array-rank array-4)))

(test-group "Array ref, set! and update!"
  (define array (make-array-from-storage vector-storage-class
                                         #(3 3)
                                         (vector 'a 'b 'c
                                                 'd 'e 'f
                                                 'g 'h 'i)))

  (test "array-ref at (0, 0) is 'a"
        'a
        (array-ref array #(0 0)))

  (test "array-ref at (2, 2) is 'i"
        'i
        (array-ref array #(2 2)))

  (array-set! array #(1 1) 999)
  (test "array-set! correctly sets value"
        999
        (array-ref array #(1 1)))

  (test-assert "array-set! does not set any other value"
               (transduce array-interval-fold
                          values
                          (collect-all
                            (lambda (idx)
                              (if (equal? idx #(1 1))
                                (eq? (array-ref array idx) 999)
                                (symbol? (array-ref array idx)))))
                          array))

  (array-update! array #(1 1) (lambda (val) (+ val 1)))
  (test "array-update! correctly updates value"
        1000
        (array-ref array #(1 1)))

  (test-assert "array-set! does not set any other value"
               (transduce array-interval-fold
                          values
                          (collect-all
                            (lambda (idx)
                              (if (equal? idx #(1 1))
                                (eq? (array-ref array idx) 1000)
                                (symbol? (array-ref array idx)))))
                          array))

  (test-error "out-of-bounds ref is an error"
              (array-ref array #(10 10)))

  (test-error "out-of-bounds set! is an error"
              (array-set! array #(10 10) 999))

  (test-error "out-of-bounds update! is an error"
              (array-update! array #(10 10) (lambda (val) 999))))

(test-group "Transducible functions"
  (define array
    (make-array-from-storage vector-storage-class
                             #(3 3)
                             (vector 'a 'b 'c
                                     'd 'e 'f
                                     'g 'h 'i)))

  (define reverse-array
    (make-array-from-storage vector-storage-class
                             #(3 3)
                             (vector 'i 'h 'g
                                     'f 'e 'd
                                     'c 'b 'a)))

  (test-assert "folding / collecting same array is equal"
               (array=? array
                        (transduce array-fold
                                   values
                                   (collect-array (array-storage-class array)
                                                  (interval-end (array-shape array)))
                                   array)))

  (test-assert "reverse folding / collecting same array is equal"
               (array=? reverse-array
                        (transduce reverse-array-fold
                                   values
                                   (collect-array (array-storage-class array)
                                                  (interval-end (array-shape array)))
                                   array)))

  (test "zipping array produces expected list"
        (transduce array-fold values (collect-list) array)
        (transduce range-fold
                   (compose
                     (zip-array array)
                     (map cdr))
                   (collect-list)
                   (counter 0)))

  (test "reverse zipping array produces expected list"
        (transduce array-fold values (collect-reverse-list) array)
        (transduce range-fold
                   (compose
                     (reverse-zip-array array)
                     (map cdr))
                   (collect-list)
                   (counter 0))))

  (test-group "Array view tests"
    (define array
      (make-array-from-storage vector-storage-class
                               #(3 3)
                               (vector 'a 'b 'c
                                       'd 'e 'f
                                       'g 'h 'i)))
    (test-assert "slice is equivalent to sub-array"
                 (array=? (make-array-from-storage vector-storage-class
                                                   #(2 2)
                                                   (vector 'e 'f
                                                           'h 'i))
                          (array-slice array #(1 1) #(3 3))))

    (test-assert "transpose is equivalent to transposed data"
                 (array=? (make-array-from-storage vector-storage-class
                                                   #(3 3)
                                                   (vector 'a 'd 'g
                                                           'b 'e 'h
                                                           'c 'f 'i))
                          (array-transpose array)))

    (test-assert "swap-axis on 2D array is equivalent to transpose"
                 (array=? (array-transpose array)
                          (array-swap-axes array 0 1)))

    (define squeezable-array
      (make-array-from-storage vector-storage-class
                               #(1 2 1 3 1)
                               (vector 'a 'b 'c
                                       'd 'e 'f)))

    (test "squeezing axis 0 produces correct shape"
          (vector 2 1 3 1)
          (interval-end
            (array-shape
              (array-squeeze-axis squeezable-array 0))))

    (test "squeezing axis 2 produces correct shape"
          (vector 1 2 3 1)
          (interval-end
            (array-shape
              (array-squeeze-axis squeezable-array 2))))

    (test "squeezing axis 4 produces correct shape"
          (vector 1 2 1 3)
          (interval-end
            (array-shape
              (array-squeeze-axis squeezable-array 4))))

    (test "squeezing all axes produces correct shape"
          (vector 2 3)
          (interval-end
            (array-shape
              (array-squeeze squeezable-array))))

    (define expandable-array
      (make-array-from-storage vector-storage-class
                               #(3)
                               (vector 'a 'b 'c)))

    (test "expanding array at axis 0 produces correct shape"
          (vector 1 3)
          (interval-end
            (array-shape
              (array-expand-axis expandable-array 0))))

    (test-error "expanding array at axis 1 is an error"
                (interval-end
                  (array-shape
                    (array-expand-axis expandable-array 1)))))

(test-group "Copy and convert arrays"
  (define array
    (make-array-from-storage vector-storage-class
                             #(3 3)
                             (vector 'a 'b 'c
                                     'd 'e 'f
                                     'g 'h 'i)))

  (test-assert "array-copy produces equal array"
               (array=? array (array-copy array)))

  (test-assert "array-copy with start bound produces sub-array"
               (array=? (array-slice array #(1 1))
                        (array-copy array #(1 1))))

  (let ((from (make-array vector-storage-class #(3 3) #t))
        (to (make-array vector-storage-class #(3 3) #f)))
    (array-copy! to #(0 0) from)
    (test-assert "array-copy! overwrites all previous values in array"
                 (transduce array-fold values (collect-all values) to)))

  (let ((from (make-array vector-storage-class #(3 3) #t))
        (to (make-array vector-storage-class #(4 3) #f))
        (true-interval (make-interval #(1 0) #(4 3))))
    (array-copy! to #(1 0) from)
    (test-assert "array-copy! with offset only copies offset values"
                 (and (transduce array-fold
                                 (compose
                                   (zip-interval (array-shape array))
                                   (filter
                                     (lambda (pair)
                                       (interval-contains? true-interval (cdr pair))))
                                   (map car))
                                 (collect-all values)
                                 to)
                      (transduce array-fold
                                 (compose
                                   (zip-interval (array-shape array))
                                   (filter
                                     (lambda (pair)
                                       (not (interval-contains? true-interval (cdr pair)))))
                                   (map car))
                                 (collect-all not)
                                 to))))

  (let ((from (make-array vector-storage-class #(4 4) #t))
        (to (make-array vector-storage-class #(4 3) #f))
        (true-interval (make-interval #(0 0) #(3 3))))
    (array-copy! to #(0 0) from #(1 1))
    (test-assert "array-copy! from start only copies from-range of values"
                 (and (transduce array-fold
                                 (compose
                                   (zip-interval (array-shape array))
                                   (filter
                                     (lambda (pair)
                                       (interval-contains? true-interval (cdr pair))))
                                   (map car))
                                 (collect-all values)
                                 to)
                      (transduce array-fold
                                 (compose
                                   (zip-interval (array-shape array))
                                   (filter
                                     (lambda (pair)
                                       (not (interval-contains? true-interval (cdr pair)))))
                                   (map car))
                                 (collect-all not)
                                 to))))

  (let ((from (make-array vector-storage-class #(4 4) #t))
        (to (make-array vector-storage-class #(4 3) #f))
        (true-interval (make-interval #(0 0) #(2 2))))
    (array-copy! to #(0 0) from #(1 1) #(3 3))
    (test-assert "array-copy! from bounds only copies from-range of values"
                 (and (transduce array-fold
                                 (compose
                                   (zip-interval (array-shape array))
                                   (filter
                                     (lambda (pair)
                                       (interval-contains? true-interval (cdr pair))))
                                   (map car))
                                 (collect-all values)
                                 to)
                      (transduce array-fold
                                 (compose
                                   (zip-interval (array-shape array))
                                   (filter
                                     (lambda (pair)
                                       (not (interval-contains? true-interval (cdr pair)))))
                                   (map car))
                                 (collect-all not)
                                 to))))

  (test "array-reshape produces shape of #(9)"
        (vector 9)
        (interval-end
          (array-shape
            (array-reshape array #(9)))))

  (test "array-reshape produces shape of #(9 1)"
        (vector 9 1)
        (interval-end
          (array-shape
            (array-reshape array #(9 1)))))

  (let* ((array (make-array u32vector-storage-class #(3 3) 2))
         (reclassified (array-reclassify array f64vector-storage-class)))
    (test "array-reclassify converts array to correct storage class"
          f64vector-storage-class
          (array-storage-class reclassified))

    (test-assert "array-reclassify converts values without modifying them"
                 (transduce array-fold
                            values
                            (collect-all
                              (lambda (v)
                                (< (- v 2.0) 1e-12)))
                            reclassified)))

  (test-assert "array-append produces correct array"
               (array=? (make-array-from-storage vector-storage-class
                                                 #(5 2)
                                                 (vector 'a 'b
                                                         'c 'd
                                                         'e 'f
                                                         'g 'h
                                                         'i 'j))
                        (array-append 0
                                      (make-array-from-storage
                                        vector-storage-class
                                        #(2 2)
                                        (vector 'a 'b
                                                'c 'd))
                                      (make-array-from-storage
                                        vector-storage-class
                                        #(3 2)
                                        (vector 'e 'f
                                                'g 'h
                                                'i 'j))))))