;; Matrix-vector multiplication and matrix solution by Gaussian elimination
;; with partial pivoting.

(require 'array-map)

;; Return a function multiplying a vector by matrix MATRIX.
;; This is done by taking the matrix one column at a time, so no
;; dot products are directly calculated.

(define (matrix-multiplier matrix)
  (let* ((spec (car (array-dimensions matrix)))
	 (cols (enclose-array matrix 0)))
    (lambda (v)
      (or (equal? (list spec) (array-dimensions v))
	  (error "matrix-multiplier: bad shape" (array-dimensions v)))
      (let ((res (make-uniform-array (array-prototype v) spec))
	    (term (make-uniform-array (array-prototype v) spec)))
	(array-map! res +)
	(array-for-each 
	 (lambda (col valias)
	   (array-map! term * valias col)
	   (array-map! res + res term))
	 cols
	 (enclose-array
	  (make-shared-array v (lambda (i j) j) spec spec) 0))
	res))))

(define (mat:pivot! mat)
  (let* ((spec (car (array-shape mat)))
	 (irow (car spec))
	 (pivcol (array-column mat irow)))
    (let search ((i irow)
		 (ipiv irow)
		 (pivmax (abs (array-ref pivcol irow))))
      (if (array-in-bounds? pivcol i)
	  (if (< pivmax (abs (array-ref pivcol i)))
	      (search (+ i 1) i (abs (array-ref pivcol i)))
	      (search (+ i 1) ipiv pivmax))
	  (begin
	    (or (= ipiv irow)
		(let ((tmprow (make-uniform-array (array-prototype mat) spec))
		      (pivrow (array-row mat ipiv))
		      (toprow (array-row mat irow)))
		  (array-copy! pivrow tmprow)
		  (array-copy! toprow pivrow)
		  (array-copy! tmprow toprow)))
	    (if (= ipiv irow)
		#f 
		ipiv))))))

(define (mat:solver! mat)
  (let* ((spec (car (array-shape mat)))
	 (irow (car spec))
	 (subspec (list (1+ irow) (cadr spec)))
	 (ipivot (mat:pivot! mat)))
    (let ((elim-row
	   (make-shared-array mat
			      (lambda inds (cons irow inds))
			      subspec))
	  (elim-col
	   (make-shared-array mat
			      (lambda (i) (list i irow))
			      subspec)))
      (let ((rows (enclose-array (subarray mat subspec subspec) 1)))
	(array-map! elim-col
		    (lambda (row-head row)
		      (let ((mult (/ row-head
				     (array-ref mat irow irow))))
			(array-map! row - row
				    (array-map (array-prototype mat)
					       *
					       elim-row
					       (scalar->array mult subspec)))
			mult))
		    elim-col rows))
      (let ((subsolve!
	     (and (<= (car subspec) (cadr spec))
		  (mat:solver! (subarray mat subspec subspec)))))
	(lambda (soln)
	  (if ipivot
	      (let ((tmp (array-ref soln ipivot)))
		(array-set! soln (array-ref soln irow) ipivot)
		(array-set! soln tmp irow)))
	  (let ((piv-soln
		 (scalar->array (exact->inexact (array-ref soln irow))
				subspec))
		(subsoln
		 (subarray soln subspec)))
	    (array-map! subsoln - subsoln
			(array-map (array-prototype soln)
				   * piv-soln elim-col))
	    (if subsolve! (subsolve! subsoln))
	    (let ((a (array-ref soln irow)))
	      (array-for-each (lambda (x) (set! a (- a x)))
			      (array-map (array-prototype soln)
					 * subsoln elim-row))
	      (array-set! soln (/ a (array-ref mat irow irow)) irow)))
	  soln)))))

;; Return a function taking a vector rhs, solving the system MATRIX * x = rhs.
(define (matrix-solver matrix)
  (let* ((spec (car (array-dimensions matrix)))
	 (solve! (mat:solver! (array-copy matrix))))
    (lambda (rhs)
      (or (equal? (list spec) (array-dimensions rhs))
	  (error "matrix-solver: bad shape" (array-dimensions rhs)))
      (solve! (array-copy rhs)))))

(define (pp-ra ra)
  (require 'pretty-print)
  (set! pp-ra
	(lambda (ra)
	  (pretty-print (array->list ra))))
  (pp-ra ra))