;;;;
;;;; srfi-25.stk         -- SRFI-25: Multi-dimensional Array Primitives
;;;;
;;;; Copyright © 2021 Jeronimo Pellegrini <j_p@aleph0.info>
;;;;
;;;;
;;;; This program is free software; you can redistribute it and/or modify
;;;; it under the terms of the GNU General Public License as published by
;;;; the Free Software Foundation; either version 2 of the License, or
;;;; (at your option) any later version.
;;;;
;;;; This program is distributed in the hope that it will be useful,
;;;; but WITHOUT ANY WARRANTY; without even the implied warranty of
;;;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
;;;; GNU General Public License for more details.
;;;;
;;;; You should have received a copy of the GNU General Public License
;;;; along with this program; if not, write to the Free Software
;;;; Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307,
;;;; USA.
;;;;
;;;; This file is a derivative work from the  implementation of
;;;; this SRFI by  Shiro Kawai, adapted by John Cowan, Kevin Wortman (2015)
;;;; it is copyrighted as:
;;;;
;;;;;; Copyright (C) Jussi Piitulainen (2001). All Rights Reserved.
;;;;;;
;;;;;; Permission is hereby granted, free of charge, to any person obtaining
;;;;;; a copy of this software and associated documentation files (the
;;;;;; "Software"), to deal in the Software without restriction, including
;;;;;; without limitation the rights to use, copy, modify, merge, publish,
;;;;;; distribute, sublicense, and/or sell copies of the Software, and to
;;;;;; permit persons to whom the Software is furnished to do so, subject to
;;;;;; the following conditions:
;;;;;;
;;;;;; The above copyright notice and this permission notice shall be
;;;;;; included in all copies or substantial portions of the Software.
;;;;;;
;;;;;; THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
;;;;;; EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
;;;;;; MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
;;;;;; NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
;;;;;; LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
;;;;;; OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
;;;;;; WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
;;;;
;;;;           Author: Jeronimo Pellegrini [j_p@aleph0.info]
;;;;    Creation date: 27-Mar-2021 02:33
;;;;

(select-module srfi/25)

(export array-for-each-index
        array->vector
        array->list
        array-map!
        array-map
        array-copy
        tabulate-array
        tabulate-array!

        array-retabulate!
        share-row
        share-column
        share-array/origin
        array-append
        transpose
        share-nths)

;;
;;  Finish array type constuction:
;;    - plug a function to describe arrays
;;    - define a new class for arrays and plug it for class-of
;;    - provide a way to read arrays
;;

;; Customize "describe"
(%user-type-proc-set! 'array 'describe
                      (lambda (x port)
                        (format port "an array of rank ~A and size ~A"
                                (array-rank x)
                                (array-size x))))
;; Customize "class-of"
(define-class <array> (<top>) ())
(export <array>)

(%user-type-proc-set! 'array 'class-of <array>)



;; Reader for arrays
(define-reader-ctor '<array>
  (lambda args
    (apply array (apply shape (car args)) (cdr args))))


(%compile-time-define  shape          array-shape    array
                       array-rank     array-set!     make-array
                       array-ref      array-end      array-start
                       shape-for-each array?         shape?
                       share-array    array-size)

;; And why not have a setter for arrays?
(set! (setter array-ref) array-set!)

;; EXTRA: almost all procdures in arlib.scm, not part of the
;;        SRFI spec but included with the reference implementation.
;;
;;; (array-shape arr)
;;; (array-length arr dim)
;;; (array-size arr)
;;; (shape-for-each shp proc [ind])
;;; (array-for-each-index arr proc [ind])
;;; (tabulate-array shp proc)
;;; (tabulate-array! shp proc ind)
;;; (array-map [shp] proc arr0 arr1 ...)
;;; (array-map! arr [shp] proc arr0 arr1 ...)
;;; (array->vector arr)
;;; (array->list arr)
;;; (share-row arr k)
;;; (share-column arr k)
;;; (share-array/origin arr k ...)
;;; (share-array/origin arr ind)
;;; (share-nths arr dim n)
;;; (transpose arr dim ...)
;;; (array-append dim arr0 arr1 ...)
;;; (array-retabulate! arr shp proc [ind])

;; NOT implemented:
;;; (share-array/prefix arr k ...)

#|
<doc EXT array-length
 * (array-length array dim)
 *
 * Returns the length of dimension |dim| in array |array|.
doc>
|#
(define (array-length arr dim)
  (when (not (array? arr))
    (error "bad array ~S" arr))
  (- (array-end arr dim)
     (array-start arr dim)))

#|
<doc EXT array-for-each-index
 * (array-for-each-index arr proc [index-object])
 *
 * Will loop through all valid indices of |array|, applying |proc|
 * to those indices.
 *
 * If |index-object| is not provided, then |proc| must accept
 * as many arguments as the number of dimensions that the shape
 * describes.
 *
 * If |index-object| is provided, it is used as a place to store the
 * indices, so |proc| must accept a vector or an array (this is to avoid
 * pushing and popping too many values when calling proc).
 * |index-object|, when present, must be aither a vector or array.
 *
 * See the documentation of |shape-for-each| for more information
 * on |index-object|.
doc>
|#
(define (array-for-each-index arr proc :optional (ind #f))
  (when (not (array? arr))
    (error "bad array ~S" arr))
  (when (not (procedure? proc))
    (error "bad procedure ~S" proc))

  (if ind
      (shape-for-each (array-shape arr) proc ind)
      (shape-for-each (array-shape arr) proc)))


#|
<doc EXT array->vector
 * (array->vector array)
 *
 * Returns a vector that contains a copy of the elements of |array|,
 * in row-major order. The new vector does not share elements with
 * the original array (it is a fresh copy).
 * This is not recursive, and will not flatten the array.
doc>
|#
(define (array->vector arr)
  (when (not (array? arr))
    (error "bad array ~S" arr))
  (let ((shape (array-shape arr))
        (rank  (array-rank arr))
        (len   (array-size arr)))
    (let ((vec (make-vector len))
          (arr-ind (make-vector rank))
          (vec-ind 0))
      (array-for-each-index arr
                            (lambda (arr-ind)
                              (vector-set! vec vec-ind (array-ref arr arr-ind))
                              (set! vec-ind (+ 1 vec-ind)))
                            arr-ind)
      vec)))

#|
<doc EXT array->list
 * (array->list array)
 *
 * Returns a list that contains a copy of the elements of |array|,
 * in row-major order.
 * This is not recursive, and will not flatten the array.
doc>
|#
(define (array->list arr)
  (vector->list (array->vector arr)))

#|
<doc EXT array-map!
 * (array-map! array [shape] proc arr~0~ arr~1~ ...)
 *
 * For each valid index |idx|, applies proc to the corresponding
 * position in |arr~0~|, |arr~1~|, ...  and then sets the same
 * place in |array| to the result.
 *
 * If |shape| is specified, it should specify a subarray of
 * |array|, and only that section will be mapped.
doc>
|#
(define (array-map! target shape-or-proc . more)
  (when (null? more) (error "at least three arguments required, ~S provided"
                            (+ 2 (length more))))

  (when (not (array? target)) (error "bad array ~S" target))

    (let ((shape (if (shape? shape-or-proc) shape (array-shape target)))
          (proc  (if (shape? shape-or-proc) (car more) shape-or-proc))
          (args  (if (shape? shape-or-proc) (cdr more) more))
          (idx-vec (make-vector (array-rank target))))

      (when (not (shape? shape))    (error "bad shape ~S" shape))
      (when (not (procedure? proc)) (error "bad procedure ~S" proc))
      (when (not (every array? args)) (error "non-array in array list"))
      (let ((list-shape (array->list shape)))
        (when (not (every (lambda (arr) (equal? list-shape arr))
                          (map array->list (map array-shape (cons target args)))))
          (error "shapes are not compatible: ~S"
                 (map array->list (map array-shape (cons shape args))))))

      (shape-for-each shape
                      (lambda (iv) ; idx-vec goes here
                        (array-set! target
                                    iv
                                    (apply proc
                                           (map (lambda (source)
                                                  (array-ref source iv))
                                                args))))
                      idx-vec)))


#|
<doc EXT array-map
 * (array-map [shape] proc arr~0~ arr~1~ ...)
 *
 * This procedure is similar to |map| for lists:
 * it will run |proc| on an element of each of the
 * |arr~0~|, |arr~1~|, ... arguments, storing the result in
 * the equivalent position of a newly created array.
 *
 * The shapes of the arrays must be the same.
 *
 * The procedure will create a new array with shape |shape|
 * (or `|arr~0~|`'s shape, if |shape| was not specified).
doc>
|#
(define (array-map . args)
  (when (< (length args) 2) (error "at least two arguments required, ~S provided"
                                   (length args)))

  (let ((shape (cond ((shape? (car args))  ;; (array-map shape proc arr0 ...)
                      (car args))

                     ((array? (cadr args)) ;; (array-map proc arr0 ...)
                      (array-shape (cadr args)))

                     (else (error "first argument ~S is not a shape, but second ~S is not an array"
                                  (car args) (cdr args))))))
    (let ((new (make-array shape #void)))
      (apply array-map! (cons new args))
      new)))

#|
<doc EXT array-copy
 * (array-copy array)
 *
 * Returns a copy of |array|.
 * The new copy will have no data shared with any other array, even if the
 * argument |array| did.
doc>
|#
(define (array-copy arr)
  (array-map (lambda (x) x) arr))


#|
<doc EXT tabulate-array
 * (tabulate-array shape proc)
 * (tabulate-array shape proc idx)
 *
 * Returns a new array of shape |shape|, populated according
 * to |proc|. Each valid index in |shape| is passed to |proc|,
 * and the result is place in the according array position.
 *
 * |idx| is an object that may be used to store the indices, and
 * it may be either a vector or an array. If it is not present, or
 * if it is `#f`, then an index vector will be created internally.
 *
doc>
|#

(define (tabulate-array shape proc :optional (ind #f))
  ;; we use vector index objects always, regardless of what the
  ;; user gave us.
  (when (not (shape? shape)) (error "bad shape ~S" shape))
  (when (not (procedure? proc)) (error "bad procedure ~S" proc))
  (let* ((arr (make-array shape)))
    (cond ((or (vector? ind) (array? ind))
           (lambda (idx) (array-set! arr idx (proc idx)))
           (shape-for-each shape (lambda (idx) (array-set! arr
                                                      idx
                                                      (proc idx)))
                           ind))

          ((not ind)
           (shape-for-each shape (lambda idx (array-set! arr
                                                    (list->vector idx)
                                                    (apply proc idx)))))

          (else
           (error "index object ~S is none of vector, array or #f" ind)))
    arr))


;; The following is here just for compatibility, requiring the third argument.
;; * Racket only implements tabulate-array, and does not accept the index argument
;; * Gauche implements tabulate-array, and accepts an optional ind argumet. This
;;   seems consistent with the other procedures.
(define (tabulate-array! shape proc ind)
  (tabulate-array shape proc ind))







#|
<doc EXT array-retabulate!
 * (array-retabulate! arr shp proc [index-object])
 *
 * Sets the elements of |arr| in |shape| to the value of |proc| at that
 * index, using |index-object| if provided. This is similar to
 * |tabulate-array!|, except that the array is given by the user.
 *
 * @lisp
 * (define arr (array (shape 0 2 0 2) 'a 'b 'c 'd))
 * (array-retabulate! arr (shape 0 2 0 2) (lambda (x y) (+ 1 x y)))
 * arr => #,(<array> (0 2 0 2) 1 2 2 3)
 * @end lisp
doc>
|#
(define (array-retabulate! arr shape proc :optional (ind #f))
  ;; we use vector index objects always, regardless of what the
  ;; user gave us.
  (when (not (shape? shape)) (error "bad shape ~S" shape))
  (when (not (procedure? proc)) (error "bad procedure ~S" proc))
  (cond ((or (vector? ind) (array? ind))
         (lambda (idx) (array-set! arr idx (proc idx)))
         (shape-for-each shape (lambda (idx) (array-set! arr
                                                    idx
                                                    (proc idx)))
                         ind))

        ((not ind)
         (shape-for-each shape (lambda idx (array-set! arr
                                                  (list->vector idx)
                                                  (apply proc idx)))))

        (else
         (error "index object ~S is none of vector, array or #f" ind)))
  arr)

#|
<doc EXT share-row
 * (share-row arr k)
 *
 * Shares whatever the first index is about. The result has one dimension less.
 *
 * @lisp
 * (define a (array (shape 0 2 0 2 0 2) -1 -2 -3 -4 -5 -6 -7 -8))
 *
 * (share-row a 0) => #,(<array> (0 2 0 2) -1 -2 -3 -4)
 * (share-row a 1) => #,(<array> (0 2 0 2) -5 -6 -7 -8)
 * @end lisp
doc>
|#
(define (share-row arr k)
  (share-array
   arr
   (let ((bounds (array->list (array-shape arr))))
     (apply shape (cddr bounds)))
   (lambda ks
     (apply values k ks))))

#|
<doc EXT share-column
 * (share-column arr k)
 *
 * Shares whatever the second index is about. The result has one dimension less.
 *
 * @lisp
 * (define a (array (shape 0 2 0 2 0 2) -1 -2 -3 -4 -5 -6 -7 -8))
 *
 * (share-column a 1) => #,(<array> (0 2 0 2) -3 -4 -7 -8)
 * (share-column a 0) => #,(<array> (0 2 0 2) -1 -2 -5 -6)
 * @end lisp
doc>
|#
(define (share-column arr k)
  (share-array
   arr
   (let ((bounds (array->list (array-shape arr))))
     (apply shape
            (car bounds) (cadr bounds)
            (cddddr bounds)))
   (lambda ks
     (apply values (car ks) k (cdr ks)))))

#|
<doc EXT share-array/origin
 * (share-array/origin arr k ...)
 * (share-array/origin arr index)
 *
 * change the origin of |arr| to |k| ..., with |index| a vector or zero-based
 * one-dimensional array that contains k ...
 *
 * @lisp
 * (define a (array (shape 0 2 0 2 ) -1 -2 -3 -4))
 *
 * (share-array/origin  a 1 1) => #,(<array> (1 3 1 3) -1 -2 -3 -4)
 * @end lisp
doc>
|#
(define (share-array/origin arr . xs)
  (let ((new (if (or (null? xs)
                     (integer? (car xs)))
                 xs
                 (apply (lambda (x)
                          (if (vector? x)
                              (vector->list x)
                              (if (array? x)
                                  (array->list x)
                                  (error "share-array/origin: bad thing"))))
                        xs))))
    (do ((k (array-rank arr) (- k 1))
         (old '() (cons (array-start arr (- k 1)) old)))
      ((= k 0)
       (let ((ds (map - new old)))
         (share-array
          arr
          (tabulate-array
           (shape 0 (array-rank arr) 0 2)
           (lambda (r k)
             (case k
               ((0) (+ (array-start arr r) (list-ref ds r)))
               ((1) (+ (array-end arr r) (list-ref ds r))))))
          (lambda ks
            (apply values (map - ks ds)))))))))


#|
<doc EXT array-append
 * (array-append dim arr~1~ arr~2~ ...)
 *
 * Appends arrays |arr~1~|, |arr~2~|, ... along the specified dimension |dim|.
 * The arrays must have equally many dimensions and all other dimensions
 * equally long.
 *
 * @lisp
 * (define a (array (shape 0 2 0 3) 11 22 33 44 55 66))
 * (define b (array (shape 0 3 0 3) -11 -22 -33 -44 -55 -66 -77 -88 -99))
 * (define c (array (shape 0 1 0 3) 'a 'b 'c))
 *
 * (array-append 0 a b c) =>  #,(<array> (0 6 0 3)
 *                                       11  22  33
 *                                       44  55  66
 *                                      -11 -22 -33
 *                                      -44 -55 -66
 *                                      -77 -88 -99
 *                                        a   b   c)
 * @end lisp
doc>
|#
(define (array-append dim arr . ars)
  (let* ((total (do ((m (array-length arr dim)
                        (+ m (array-length (car r) dim)))
                     (r ars (cdr r)))
                  ((null? r) m)))
         (common (array-shape arr))
         (origin (array->vector (share-column common 0)))
         (index (make-vector (array-rank arr))))
    (array-set! common dim 1 (+ (array-start arr dim) total))
    (let ((result (make-array common)))
      (array-set! common dim 1 (array-start arr dim))
      (let wok ((arr arr)
                (ars ars))
        (vector-set! origin dim (array-ref common dim 1))
        (let ((arr1 (share-array/origin arr origin)))
          (array-set! common dim 0 (array-start arr1 dim))
          (array-set! common dim 1 (array-end arr1 dim))
          (shape-for-each
           common
           (lambda (index)
             (array-set! result index (array-ref arr1 index)))
           index))
        (if (pair? ars)
            (wok (car ars) (cdr ars))))
      result)))


(define (matrix-times a b)
  (or (and (= (array-rank a) 2)
           (= (array-rank b) 2))
      (error "matrix-times: arrays are not matrices"))
  (let ((r0 (array-start a 0))  (rn (array-end a 0))
        (t0 (array-start a 1))  (tn (array-end a 1))
        (u0 (array-start b 0))  (un (array-end b 0))
        (k0 (array-start b 1))  (kn (array-end b 1)))
    (or (= (- tn t0) (- un u0))
        (error "matrix-times: matrices are not compatible"))
    (let ((ab (make-array (shape r0 rn k0 kn))))
      (do ((r r0 (+ r 1)))
        ((= r rn))
        (do ((k k0 (+ k 1)))
          ((= k kn))
          (do ((t t0 (+ t 1))
               (u u0 (+ u 1))
               (s 0 (+ s (* (array-ref a r t)
                            (array-ref b u k)))))
            ((and (= t tn)
                  (= u un))
             (array-set! ab r k s)))))
      ab)))

#|
<doc EXT transpose
 * (transpose arr k ...)
 *
 * Shares |arr| with permuted dimensions. Each dimension from 0
 * inclusive to rank exclusive must appear once in |k| ...
 *
 * This is a generalized transpose. It can permute the dimensions any which
 * way. The permutation is provided by a permutation matrix: a square matrix
 * of zeros and ones, with exactly one one in each row and column, or a
 * permutation of the rows of an identity matrix; the size of the matrix
 * must match the number of dimensions of the array.
 *
 * The default permutation is |[ 0 1 , 1 0 ]| of course, but any permutation
 * array can be specified, and the shape array of the original array is then
 * multiplied with it, and index column vectors of the new array with its
 * inverse, from left, to permute the rows appropriately.
 * @lisp
 * (transpose (array (shape 0 4 0 4)
 *                   -1  -2   -3  -4
 *                   -5  -6   -7  -8
 *                   -9  -10 -11 -12
 *                   -13 -14 -15 -16))
 *  => #,(<array> (0 4 0 4)
 *               -1 -5  -9 -13
 *               -2 -6 -10 -14
 *               -3 -7 -11 -15
 *               -4 -8 -12 -16)
 *
 * (transpose (array (shape 0 3 0 3 0 2)
 *                   -1 -2
 *                   -3 -4
 *                   -5 -6
 *
 *                   -7 -8
 *                   -9 -10
 *                   -11 -12
 *
 *                   -13 -14
 *                   -15 -16
 *                   -17 -18))
 *  => #,(<array> (0 2 0 3 0 3)
 *                -1  -7 -13
 *                -3  -9 -15
 *                -5 -11 -17
 *
 *                -2  -8 -14
 *                -4 -10 -16
 *                -6 -12 -18)
 * @end lisp
doc>
|#
(define (permutation-matrix . ds)
  (let* ((n (length ds))
         (arr (make-array (shape 0 n 0 n) 0)))
    (do ((k 0 (+ k 1))
         (ds ds (cdr ds)))
      ((= k n))
      (array-set! arr k (car ds) 1))
    arr))

(define (transpose a . p0)
  (let* ((r (array-rank a))
         (permutation (apply permutation-matrix
                             (if (pair? p0)
                                 p0
                                 (do ((ds '() (cons d ds))
                                      (d 0 (+ d 1)))
                                   ((= d r)
                                    ;; reverse dimensions
                                    ds)))))
         (inverse-permutation (share-array permutation
                                           (array-shape permutation)
                                           (lambda (r k)
                                             ;; transpose
                                             (values k r)))))
    (share-array
     a
     (matrix-times permutation (array-shape a))
     (lambda ks0
       (apply values
              (array->list
               (matrix-times
                inverse-permutation
                (apply array (shape 0 r 0 1) ks0))))))))



#|
<doc EXT share-nths
 * (share-nths a d n)
 *
 * |Share-nths| takes every |n|th slice along dimension |d| into a shared array.
 * This preserves the origin.
 *
 * @lisp
 * (define a (array (shape 0 4 0 4)
 *                  -1 -2 -3 -4
 *                  -5 -6 -7 -8
 *                  -9 -10 -11 -12
 *                  -13 -14 -15 -16))
 *
 * (share-nths a 0 2)
 *  => #,(<array> (0 2 0 4) -1  -2  -3  -4
 *                          -9 -10 -11 -12)
 *
 * (share-nths a 1 2)
 *  => #,(<array> (0 4 0 2) -1  -3  -5  -7
 *                          -9 -11 -13 -15)
 * @end lisp
doc>
|#
(define (share-nths arr d n)
  (let* ((bounds (array->vector (array-shape arr)))
         (b (vector-ref bounds (* 2 d)))
         (e (vector-ref bounds (+ (* 2 d) 1))))
    (vector-set! bounds (+ (* 2 d) 1) (+ b (quotient (+ n (- e b 1)) n)))
    (share-array
     arr
     (apply shape (vector->list bounds))
     (lambda ks
       (apply values
              (let d/nk ((u 0) (ks ks))
                (if (= u d)
                    (cons (+ b (* n (- (car ks) b))) (cdr ks))
                    (cons (car ks) (d/nk (+ u 1) (cdr ks))))))))))


(provide "srfi/25")
