;;;
;;; Matrix operations on arrays
;;;
;;;
;;;  Copyright (C) 2004 Alex Shinn (foof@synthcode.com)
;;;
;;;  Permission to use, copy, modify, distribute this software and
;;;  accompanying documentation for any purpose is hereby granted,
;;;  provided that existing copyright notices are retained in all
;;;  copies and that this notice is included verbatim in all
;;;  distributions.
;;;  This software is provided as is, without express or implied
;;;  warranty.  In no circumstances the author(s) shall be liable
;;;  for any damages arising out of the use of this software.
;;;
;;; $Id: matrix.scm,v 1.4 2006/04/06 04:09:31 shirok Exp $
;;;

(select-module gauche.array)

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; general array manipulation

(define (array-concatenate a b . opt)
  (let* ((dimension (get-optional opt 0))
         (a-start (start-vector-of a))
         (a-end (end-vector-of a))
         (b-start (start-vector-of b))
         (b-end (end-vector-of b))
         (diff (s32vector->vector (s32vector-sub a-start b-start)))
         (a-rank (s32vector-length a-start))
         (b-rank (s32vector-length b-start)))
    (unless (= a-rank b-rank)
      (error "rank mismatch"))
    (unless (every (lambda (d)
                     (or (= d dimension)
                         (= (- (s32vector-ref a-end d) (s32vector-ref a-start d))
                            (- (s32vector-ref b-end d) (s32vector-ref b-start d)))))
                   (iota a-rank))
      (error "shape mismatch"))
    (let ((c-start (s32vector-copy a-start))
          (c-end (s32vector-copy a-end)))
      (s32vector-set! c-end dimension (+ (s32vector-ref c-end dimension)
                                         (- (s32vector-ref b-end dimension)
                                            (s32vector-ref b-start dimension))))
      (let ((c (make-minimal-backend-array
                (list a b) (start/end-vector->shape c-start c-end)))
            (off (- (s32vector-ref a-end dimension) (s32vector-ref a-start dimension))))
        (array-for-each-index a
          (lambda (i)
            (array-set! c i (array-ref a i)))
          (make-vector a-rank))
        (array-for-each-index b
          (lambda (i)
            (let ((j (vector-copy i)))
              (dotimes (dim a-rank)
                (vector-set! j dim (+ (vector-ref j dim) (vector-ref diff dim))))
              (vector-set! j dimension (+ off (vector-ref j dimension)))
              (array-set! c j (array-ref b i))))
          (make-vector b-rank))
        c))))

(define (array-transpose a . opt)
  (let-optionals* opt ((dim1 0) (dim2 1))
    (let* ((sh (copy-object (array-shape a)))
           (rank (array-rank a))
           (tmp0 (array-ref sh dim1 0))
           (tmp1 (array-ref sh dim1 1)))
      (array-set! sh dim1 0 (array-ref sh dim2 0))
      (array-set! sh dim1 1 (array-ref sh dim2 1))
      (array-set! sh dim2 0 tmp0)
      (array-set! sh dim2 1 tmp1)
      (let ((res (make-array sh)))
        (array-for-each-index a
          (lambda (vec1)
            (let* ((vec2 (vector-copy vec1))
                   (tmp (vector-ref vec2 dim1)))
              (vector-set! vec2 dim1 (vector-ref vec2 dim2))
              (vector-set! vec2 dim2 tmp)
              (array-set! res vec2 (array-ref a vec1))))
          (make-vector rank))
        res))))

(define (array-rotate-90 a . opt)
  (let-optionals* opt ((dim1 0) (dim2 1))
    (let* ((sh (copy-object (array-shape a)))
           (rank (array-rank a))
           (tmp0 (array-ref sh dim1 0))
           (tmp1 (array-ref sh dim1 1))
           (off (- (array-ref sh dim1 1) 1)))
      (array-set! sh dim1 0 (array-ref sh dim2 0))
      (array-set! sh dim1 1 (array-ref sh dim2 1))
      (array-set! sh dim2 0 tmp0)
      (array-set! sh dim2 1 tmp1)
      (let ((res (make-array sh)))
        (array-for-each-index a
          (lambda (vec1)
            (let* ((vec2 (vector-copy vec1))
                   (tmp (vector-ref vec2 dim1)))
              (vector-set! vec2 dim1 (vector-ref vec2 dim2))
              (vector-set! vec2 dim2 (- off tmp))
              (array-set! res vec2 (array-ref a vec1))))
          (make-vector rank))
        res))))

(define (array-flip! a . opt)
  (let* ((dimension (get-optional opt 0))
         (lo (s32vector-ref (start-vector-of a) dimension))
         (end (end-vector-of a))
         (hi (s32vector-ref end dimension))
         (half (quotient hi 2))
         (rank (s32vector-length end)))
    (array-for-each-index-by-dimension a (delete! dimension (iota rank))
      (lambda (vec1)
        (do ((i lo (+ i 1)))
            ((= i half))
          (vector-set! vec1 dimension i)
          (let* ((j (- hi i 1))
                 (vec2 (vector-copy vec1))
                 (tmp (array-ref a vec1)))
            (vector-set! vec2 dimension j)
            (array-set! a vec1 (array-ref a vec2))
            (array-set! a vec2 tmp))))
      (make-vector rank))))

(define (array-flip a . args)
  (let ((res (copy-object a)))
    (apply array-flip! res args)
    res))


;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; linear algebra

(define (identity-array n . opt)
  (let ((res (make-array-internal (get-optional opt <array>)
                                  (shape 0 n 0 n) 0)))
    (do ((i 0 (+ i 1)))
        ((= i n) res)
      (array-set! res i i 1))))

;; Gaussian elimination, returns factor applied to determinant
(define (array-row-echelon! a)
  (let* ((start (start-vector-of a))
         (row-start (s32vector-ref start 0))
         (col-start (s32vector-ref start 1))
         (end (end-vector-of a))
         (row-end (s32vector-ref end 0))
         (col-end (s32vector-ref end 1))
         (row-col-offset (- row-start col-start)))
    (define (row-swap! i j)
      (do ((k col-start (+ k 1)))
          ((= k col-end))
        (let ((temp (array-ref a i k)))
          (array-set! a i k (array-ref a j k))
          (array-set! a j k temp))))
    (define (row-sub! i j factor)
      (do ((k col-start (+ k 1)))
          ((= k col-end))
        (array-set! a i k (- (array-ref a i k) (* factor (array-ref a j k))))))
    (let loop ((i row-start) (factor 1))
      (let ((col (- i row-col-offset)))
        (cond
          ((= i row-end) factor)
          ((zero? (array-ref a i col))
           ;; pivot non-zero row to top
           (let loop2 ((j (+ i 1)))
             (cond ((= j row-end) 0)
                   ((zero? (array-ref a j col))
                    (loop2 (+ j 1)))
                   (else
                    (row-swap! j i)
                    (loop i (* factor -1))))))
          (else
           ;; eliminate other non-zero rows
           (let loop2 ((j (+ i 1)))
             (cond ((= j row-end) (loop (+ i 1) factor))
                   ((zero? (array-ref a j col))
                    (loop2 (+ j 1)))
                   (else
                    (let ((factor (/ (array-ref a j col)
                                     (array-ref a i col))))
                      (row-sub! j i factor)
                      (loop2 (+ j 1))))))))))))

(define (array-solve-left-identity! a)
  (array-row-echelon! a)
  (let* ((start (start-vector-of a))
         (row-start (s32vector-ref start 0))
         (col-start (s32vector-ref start 1))
         (end (end-vector-of a))
         (row-end (s32vector-ref end 0))
         (col-end (s32vector-ref end 1))
         (row-col-offset (- row-start col-start)))
    ;; zero-out
    (do ((i (- row-end 1) (- i 1)))   ; for-each row in reverse
        ((< i row-start))
      (let ((divisor (array-ref a i (- i row-col-offset))))
        (unless (zero? divisor)
          (do ((j (- i 1) (- j 1)))   ; for-each lower row
              ((< j row-start))
            (let ((factor (/ (array-ref a j (- i row-col-offset)) divisor)))
              (array-set! a i (- j row-col-offset) 0)
              (do ((k (- j row-col-offset) (+ k 1))) ; for-each column
                  ((= k col-end))
                (array-set! a j k (- (array-ref a j k)
                                     (* factor (array-ref a i k))))))))))
    ;; reduce
    (do ((i row-start (+ i 1)))       ; for-each row
        ((= i row-end))
      (let ((divisor (array-ref a i (- i row-col-offset))))
        (unless (or (zero? divisor) (= divisor 1))
          (array-set! a i (- i row-col-offset) 1)
          (do ((j (- i row-col-offset -1) (+ j 1))) ; for-each higher column
              ((= j col-end))
            (array-set! a i j (/ (array-ref a i j) divisor))))))))

(define (array-inverse a)
  (let* ((start (start-vector-of a))
         (end (end-vector-of a))
         (rank (s32vector-length start))
         (n (- (s32vector-ref end 0) (s32vector-ref start 0)))
         (m (- (s32vector-ref end 1) (s32vector-ref start 1))))
    (unless (= 2 rank)
      (error "can only compute inverses of 2D arrays"))
    (unless (= n m)
      (error "can only compute inverses of square matrices"))
    (let* ((class (class-of a))
           (id (identity-array n (if (or (eq? class <f32array>)
                                         (eq? class <f64array>))
                                   class <array>)))
           (tmp (array-concatenate a id 1)))
      (array-solve-left-identity! tmp)
      (and (= 1 (array-ref tmp (- (s32vector-ref end 0) 1)
                           (- (s32vector-ref end 1) 1)))
           (subarray tmp (shape (s32vector-ref start 0) (s32vector-ref end 0)
                                (s32vector-ref end 1) (+ (s32vector-ref end 1) n)))))))


(define (determinant! a)
  (let* ((start (s32vector->list (start-vector-of a)))
         (end (s32vector->list (end-vector-of a)))
         (row-col-offset (- (car start) (cadr start)))
         (factor (array-row-echelon! a)))
    (unless (= 2 (length start)) ; add determinant for the 2x2x2 case?
      (error "can't compute hyperdeterminants in the general case"))
    (unless (apply = (map - end start))
      (error "can't compute determinants of non-square matrices"))
    (apply * factor (map (lambda (i) (array-ref a i (- i row-col-offset)))
                         (map (cute + <> (car start))
                              (iota (- (car end) (car start))))))))

(define (determinant a)
  (let ((class (class-of a)))
    (if (or (eq? class <f32array>)
            (eq? class <f64array>)
            (eq? class <array>))
      (determinant! (copy-object a))
      (let* ((rank (s32vector-length (start-vector-of a)))
             (b (tabulate-array (array-shape a)
                                (lambda (ind) (array-ref a ind))
                                (make-vector rank))))
        (determinant! b)))))


;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; matrix arithmetic

(define (array-mul a b) ; NxM * MxP => NxP
  (let ((a-start (start-vector-of a))
        (a-end (end-vector-of a))
        (b-start (start-vector-of b))
        (b-end (end-vector-of b)))
    (unless (= 2 (s32vector-length a-start) (s32vector-length b-start))
      (error "array-mul matrices must be of rank 2"))
    (let* ((a-start-row (s32vector-ref a-start 0))
           (a-end-row (s32vector-ref a-end 0))
           (a-start-col (s32vector-ref a-start 1))
           (a-end-col (s32vector-ref a-end 1))
           (b-start-col (s32vector-ref b-start 1))
           (b-end-col (s32vector-ref b-end 1))
           (n (- a-end-row a-start-row))
           (m (- a-end-col a-start-col))
           (p (- b-end-col b-start-col))
           (a-col-b-row-off (- a-start-col (s32vector-ref b-start 0)))
           (res (make-minimal-backend-array (list a b) (shape 0 n 0 p))))
      (unless (= m (- (s32vector-ref b-end 0) (s32vector-ref b-start 0)))
        (errorf "dimension mismatch: can't mul shapes ~S and ~S"
                (array-shape a) (array-shape b)))
      (do ((i a-start-row (+ i 1)))       ; for-each row of a
          ((= i a-end-row) res)
        (do ((k b-start-col (+ k 1)))     ; for-each col of b
            ((= k b-end-col))
          (let ((tmp 0))
            (do ((j a-start-col (+ j 1))) ; for-each col of a & row of b
                ((= j a-end-col))
              (inc! tmp (* (array-ref a i j) (array-ref b (- j a-col-b-row-off) k))))
            (array-set! res (- i a-start-row) (- k b-start-col) tmp)))))))

(define (array-div-left a b)
  (array-mul (array-inverse b) a))

(define (array-div-right a b)
  (array-mul a (array-inverse b)))

(define (array-expt ar pow)
  (let loop ((a ar) (n pow))
    (case n
      ((0) (identity-array (s32vector-length (start-vector-of a))))
      ((1) a)
      ((2) (array-mul a a))
      ((3) (array-mul (array-mul a a) a))
      (else
       (let* ((res1 (loop a (ash n -1)))
              (res (array-mul res1 res1)))
         (if (odd? n)
           (array-mul res a)
           res))))))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; element-wise operations (advantage over a direct array-map! is
;; ability to intermingle scalars)

;; add

(define-method array-add-elements! ((a <array-base>) b c . rest)
  (array-add-elements! a b)
  (apply array-add-elements! a c rest))

(define-method array-add-elements! ((a <array-base>))
  a)

(define-method array-add-elements! ((a <array-base>) (b <number>))
  (array-map! a (lambda (i) (+ b i)) a))

(define-method array-add-elements! ((a <number>) (b <array-base>))
  (array-map! (copy-object b) (lambda (i) (+ a i)) b))

(define-method array-add-elements! ((a <array-base>) (b <array-base>))
  (array-map! a (lambda (i j) (+ i j)) a b))

(define (array-add-elements a . rest)
  (let ((res (copy-object a)))
    (apply array-add-elements! res rest)
    res))

;; sub

(define-method array-sub-elements! ((a <array-base>))
  (array-sub-elements! 1 a))

(define-method array-sub-elements! ((a <array-base>) b c . rest)
  (array-sub-elements! a b)
  (apply array-sub-elements! a c rest))

(define-method array-sub-elements! ((a <array-base>) (b <number>))
  (array-map! a (lambda (i) (- i b)) a))

(define-method array-sub-elements! ((a <number>) (b <array-base>))
  (array-map! (copy-object b) (lambda (i) (- a i)) b))

(define-method array-sub-elements! ((a <array-base>) (b <array-base>))
  (array-map! a (lambda (i j) (- i j)) a b))

(define (array-sub-elements a . rest)
  (let ((res (copy-object a)))
    (apply array-sub-elements! res rest)
    res))

;; mul

(define-method array-mul-elements! ((a <array-base>) b c . rest)
  (array-mul-elements! a b)
  (apply array-mul-elements! a c rest))

(define-method array-mul-elements! ((a <array-base>))
  a)

(define-method array-mul-elements! ((a <array-base>) (b <number>))
  (array-map! a (lambda (i) (* b i)) a))

(define-method array-mul-elements! ((a <number>) (b <array-base>))
  (array-map! (copy-object b) (lambda (i) (* a i)) b))

(define-method array-mul-elements! ((a <array-base>) (b <array-base>))
  (array-map! a (lambda (i j) (* i j)) a b))

(define (array-mul-elements a . rest)
  (let ((res (copy-object a)))
    (apply array-mul-elements! res rest)
    res))

;; div

(define-method array-div-elements! ((a <array-base>) b c . rest)
  (array-div-elements! a b)
  (apply array-div-elements! a c rest))

(define-method array-div-elements! ((a <array-base>))
  a)

(define-method array-div-elements! ((a <array-base>) (b <number>))
  (array-map! a (lambda (i) (/ i b)) a))

(define-method array-div-elements! ((a <number>) (b <array-base>))
  (array-map! (copy-object b) (lambda (i) (/ a i)) b))

(define-method array-div-elements! ((a <array-base>) (b <array-base>))
  (array-map! a (lambda (i j) (/ i j)) a b))

(define (array-div-elements a . rest)
  (let ((res (copy-object a)))
    (apply array-div-elements! res rest)
    res))


;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; utilities

;; Nicely formatted printing for arrays of any rank, by default both
;; readable (using #,() syntax) and filled so columns line up.  Ranks
;; higher than 2 are represented as successive 2D drawings.
;;
;; Can generate (generally non-readable) box drawings with :left,
;; :right, etc. keywords.  For examples, Jacal-style matrix output can
;; be done with:
;;
;;  (pretty-print-array a #t :readable? #f :left "[" :right "]")
;;
;; and a tic-tac-toe board can be printed with
;;
;;  (pretty-print-array a #t :readable? #f :line-char "|"
;;     :top "-" :middle-col "|" :center-row "-")
;;
;; The API is subject to change.

(define (pretty-print-array a . opt)
  (let* ((p1 (get-optional opt (current-output-port)))
         (p (cond ((port? p1) p1)
                  (p1 (current-output-port))
                  (else (open-output-string))))
         (rank (array-rank a))
         (row-dim (- rank 2))
         (col-dim (- rank 1))
         (shape (array-shape a))
         (shape-list (array->list shape))
         (inner-shape (take-right shape-list 4))
         (nrows (cadr inner-shape))
         (ncols (cadddr inner-shape)))
    (let-keywords* (if (pair? opt) (cdr opt) '())
        ((readable? #t)
         (fill? #t)
         (line-char #f)
         (left line-char)
         (right left)
         (middle-col #\space)
         (top line-char)
         (bottom top)
         (center-row #f)
         (pad-char #\space))
      (define (replicate x width)
        (if (char? x)
          (make-string width x)
          (let ((len (string-length x)))
            (if (eq? len 1)
              (make-string width (string-ref x 0))
              (receive (quot rem) (quotient&remainder width len)
                (string-append
                 (apply string-append (map (lambda (n) x) (iota quot)))
                 (substring x 0 rem)))))))
      (define (pad x width)
        (let* ((res (x->string x))
               (len (string-length res)))
          (if (< len width)
            (string-append (make-string (- width len) pad-char) res )
            res)))
      (define (print-width x)
        (cond ((char? x) 1)
              ((string? x) (string-length x))
              (else (string-length (x->string x)))))
      (let ((fixed-width (+ (if left (print-width left) 0)
                            (if right (print-width right) 0)
                            (if middle-col (* (- ncols 1) (print-width middle-col)) 0))))
        (define (column-widths vec)
          (let loop1 ((col (caddr inner-shape))
                      (res '()))
            (cond
              ((< col ncols)
               (vector-set! vec col-dim col)
               (let loop2 ((row (car inner-shape))
                           (x 0))
                 (cond
                   ((< row nrows)
                    (vector-set! vec row-dim row)
                    (loop2 (+ row 1) (max x (print-width (array-ref a vec)))))
                   (else
                    (loop1 (+ col 1) (cons x res))))))
              (else (reverse res)))))
        (define (print2d vec)
          (newline p)
          (let* ((widths (if fill? (column-widths (vector-copy vec)) '()))
                 (total-width (apply + fixed-width widths)))
            (when top
              (display (if fill? (replicate top total-width) top) p)
              (newline p))
            (let loop1 ((row (car inner-shape)))
              (when (< row nrows)
                (when left (display left p))
                (vector-set! vec row-dim row)
                (let loop2 ((col (caddr inner-shape))
                            (w widths))
                  (when (< col ncols)
                    (vector-set! vec col-dim col)
                    (display (if (pair? w)
                               (pad (array-ref a vec) (car w))
                               (array-ref a vec))
                             p)
                    (when (and middle-col (< (+ col 1) ncols))
                      (display middle-col p))
                    (loop2 (+ col 1) (if (pair? w) (cdr w) w))))
                (when right (display right p))
                (newline p)
                (when (and center-row (< (+ row 1) nrows))
                  (display (if fill? (replicate bottom total-width) bottom))
                  (newline p))
                (loop1 (+ row 1))))
            (when bottom
              (display (if fill? (replicate bottom total-width) bottom) p)
              (newline p))))
        (cond
          ((< rank 2)
           (if readable? (write a p) (display a p)))
          (else
           (when readable?
             (format p "#,(~A ~S" (class-name (class-of a)) shape-list))
           (if (= rank 2)
             (print2d (make-vector rank))
             (array-for-each-index-by-dimension
              a (drop-right (iota rank) 2)
              print2d
              (make-vector rank)))
           (when readable?
             (format p "  )"))
           (newline p)))
        (if (not p1) (get-output-string p))))))


;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; internal utility to keep arrays uniform when possible

(define *signed-arrays*
  (list <s64array> <s32array> <s16array> <s8array>))

(define (make-minimal-backend-array ls sh . args)
  (define (return class)
    (apply make-array-internal class sh args))
  (define (signed? x)
    (member x *signed-arrays*))
  (define (size x)
    (cond ((or (eq? x <s64array>) (eq? x <u64array>)) 64)
          ((or (eq? x <s32array>) (eq? x <u32array>)) 32)
          ((or (eq? x <s16array>) (eq? x <u16array>)) 16)
          ((or (eq? x <s8array>) (eq? x <u8array>)) 8)
          (else 0)))
  (define (backend signed? size)
    (if signed?
      (case size
        ((64) <s64array>) ((32) <s32array>) ((16) <s16array>) ((8) <s8array>) (else <array>))
      (case size
        ((64) <u64array>) ((32) <u32array>) ((16) <u16array>) ((8) <u8array>) (else <array>))))
  (define (join a b)
    (cond
      ((or (eq? a <array>) (eq? b <array>))
       <array>)
      ((or (eq? a <f32array>) (eq? a <f64array>))
       (if (or (eq? b <f32array>) (eq? b <f64array>))
         (if (eq? a b) a <f64array>)
         <array>))
      (else
       (let ((a-s? (signed? a))
             (a-size (size a))
             (b-s? (signed? b))
             (b-size (size b)))
         (if (or (zero? a-size) (zero? b-size))
           <array>
           (let ((res-size (max a-size b-size)))
             (backend (or a-s? b-s?)
                      (if (if a-s? (not b-s?) b-s?)
                        (* 2 res-size) ; signs differ, double size
                        res-size))))))))
  (let loop ((class (class-of (car ls))) (l ls))
    (if (null? l)
      (return class)
      (let ((next (join class (class-of (car l)))))
        (if (eq? next <array>)
          (return <array>) ; already most general array
          (loop next (cdr l)))))))

(provide "gauche/matrix")

