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

;;; Input/output

;; The external representation of an array consists of a # character, followed by
;; the letter a, followed by a sequence of digits indicating the rank of the
;; array, followed by a coded representation of the storage class, all with no
;; intervening whitespace. This prefix is followed, after optional whitespace, by
;; the representation of a nested list produced as if by array->nested-list. The
;; prefix is interpreted case-insensitively.

(define (storage-class-from-identifier id)
  (let ((storage-class
          (transduce list-fold
                     (filter (lambda (s) (eq? id (storage-class-short-id s))))
                     (collect-first #f)
                     (list vector-storage-class
                           u8vector-storage-class
                           s8vector-storage-class
                           u16vector-storage-class
                           s16vector-storage-class
                           u32vector-storage-class
                           s32vector-storage-class
                           u64vector-storage-class
                           s64vector-storage-class
                           f32vector-storage-class
                           f64vector-storage-class
                           c64vector-storage-class
                           c128vector-storage-class))))
    (if storage-class
      storage-class
      vector-storage-class)))

(define-checked (array->nested-list (array array?))
  (let* ((dimension (interval-end (array-shape array)))
         (rank (array-rank array))
         (upper (fx- rank 1))
         (chunker
           (apply compose
                  (transduce fixnum-range-fold
                             (map
                               (lambda (i)
                                 (let ((n (vector-ref dimension (fx- upper i))))
                                   (chunks n (collect-list)))))
                             (collect-reverse-list)
                             (range 0 upper)))))
    (transduce array-fold chunker (collect-list) array)))

(define (rcar lst n)
  (let loop ((lst lst)
             (i 0))
    (if (fx<? i n)
      (loop (car lst) (fx+ i 1))
      lst)))

(define-checked (nested-list->array (rank non-negative-fixnum?)
                                    (storage-class storage-class?)
                                    (nested-list list?))
  (let ((dimension
          (transduce fixnum-range-fold
                     (map
                       (lambda (i)
                         (transduce list-fold
                                    values
                                    (collect-count)
                                    (rcar nested-list i))))
                     (collect-vector)
                     (range 0 rank))))
    (make-array-from-storage
      storage-class
      dimension
      (transduce list-fold
                 flatten-list
                 ((transducible-collector
                    (storage-class-transducible storage-class)))
                 nested-list))))

(define (char->digit c)
  (case c
    ((#\0) 0)
    ((#\1) 1)
    ((#\2) 2)
    ((#\3) 3)
    ((#\4) 4)
    ((#\5) 5)
    ((#\6) 6)
    ((#\7) 7)
    ((#\8) 8)
    ((#\9) 9)
    (else #f)))

(define read-number
  (case-lambda
    (() (read-number (current-input-port)))
    ((port)
     (let loop ((next-char (peek-char port))
                (number 0))
       (let ((digit (char->digit next-char)))
         (cond
           (digit
             (read-char)
             (loop (peek-char port)
               (+ (* 10 number) digit)))
           (else number)))))))

;; Reads the external representation of an array from input-port (the current
;; input port if input-port is not specified) and returns the corresponding
;; array. If the coded representation of the storage class is not recognized,
;; vector-storage-class is used; this permits the addition of new coded storage
;; classes in a backward compatible way.
(define array-read
  (case-lambda-checked
    (()
     (array-read (current-input-port)))
    (((port port?))
     (read-char port) ; character #\#
     (read-char port) ; character #\a
     (let ((rank (read-number port))
           (storage-class (storage-class-from-identifier (read port)))
           (nested-list (read port)))
       (nested-list->array rank storage-class nested-list)))))

;; Writes the external representation of array from output-port (the current
;; output port if output-port is not specified) and returns an unspecified
;; value.
(define array-write
  (case-lambda-checked
    ((array)
     (array-write array (current-output-port)))
    (((array array?) (port port?))
     (format port
             "#a~s~s~s"
             (array-rank array)
             (storage-class-short-id (array-storage-class array))
             (array->nested-list array)))))

(cond-expand
  (chicken-5
    (import (only (chicken base) define-record-printer)
            (only (chicken read-syntax) set-sharp-read-syntax!))

    (define-record-printer (<array> array output-port)
      (array-write array output-port))

    (set-sharp-read-syntax!
      #\a
      (lambda (port)
        (let ((rank (read-number port))
              (storage-class (storage-class-from-identifier (read port)))
              (nested-list (read port)))
          `(quote ,(nested-list->array rank storage-class nested-list)))))))