;;
;; 
;; Generation of special utility matrices.
;;
;;
;; Copyright 2007-2010 Ivan Raikov and the Okinawa Institute of
;; Science and Technology.
;;
;;
;; 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 3 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.
;;
;; A full copy of the GPL license can be found at
;; <http://www.gnu.org/licenses/>.
;;
;;

(module matrix-utils

                (matrix-utils:error
		 make-matrix-map
		 make-matrix-fold
		 make-matrix-fold-partial
		 make-fill-matrix
		 make-matrix-ones
		 make-matrix-zeros
		 make-matrix-eye
		 make-matrix-diag
		 make-linspace
		 make-logspace
		 
		 define-utility-matrices
		 with-utility-matrices)

   (import scheme chicken data-structures srfi-4)
   (require-extension srfi-4 srfi-42 srfi-4-comprehensions blas)


(define (matrix-utils:error x . rest)
  (let ((port (open-output-string)))
    (let loop ((objs (cons x rest)))
      (if (null? objs)
	  (begin
	    (newline port)
	    (error 'matrix-utils (get-output-string port)))
	  (begin (display (car objs) port)
		 (display " " port)
		 (loop (cdr objs)))))))

;;
;; Given a procedure VECTOR-SET! returns a procedure FILL-MATRIX! of
;; the form FILL-MATRIX!:: M * N * A * F * F0 * [IX * IY * EX * EY] -> A
;;
;; Where procedure FILL-MATRIX! fills matrix A of size M x N with the
;; values returned by applying procedure F to each pair of indices in
;; the matrix.
;;
;; Procedure F is of the form LAMBDA I J AX -> VAL * AX1, where I and
;; J are matrix indices, and AX is accumulator value (like in
;; fold). The initial value of AX is given by F0. Procedure F is
;; expected to return two values: the value to be placed in matrix A
;; at position [I,J], and the new accumulator value (or #f). 
;;
;; Optional arguments IX IY EX EY may specify a sub-matrix in matrix A
;; to be filled. These arguments are checked to make sure they specify
;; a valid sub-matrix.
;;
;; VECTOR-SET! is one of the homogeneous vector setting procedure from
;; SRFI-4. Procedure F must ensure that it returns values that are
;; within the range of the SRFI-4 type used.
;;
(define (make-fill-matrix vector-set!)
  (lambda (order M N A f f0 . rest) 
    ;; optional arguments to specify a sub-matrix to be filled
    (let-optionals rest ((ix 0) (iy 0) (ex M) (ey N))
     (if (not (and (fx>= ix 0) (fx>= iy 0) (fx<= ex M) (fx<= ey N)
		   (fx<= ix ex) (fx<= iy ey)))
	 (matrix-utils:error 'fill-matrix! "invalid sub-matrix dimensions: " (list ix iy ex ey)))

     (cond ((= order blas:RowMajor)
	    (fold-ec f0 (:parallel (:range b (fx* N ix) (fx* N M) N) (:range x ix ex)) (cons x b)
		     (lambda (x+b ax)
		       (fold-ec ax (:range y iy ey) y
				(lambda (y ax)
				  (let ((i (fx- (car x+b) ix)) (j (fx- y iy)))
				    (let-values (((val ax1) (f i j ax)))
				      (vector-set! A (fx+ (cdr x+b) y) val)
				      ax1)))))))
	   
	   ((= order blas:ColMajor)
	    (fold-ec f0 (:parallel (:range b (fx* N ix) (fx* N M) M) (:range y iy ey)) (cons y b)
		     (lambda (y+b ax)
		       (fold-ec ax (:range x ix ex) x
				(lambda (x ax)
				  (let ((i (fx- x ix)) (j (fx- (car y+b) iy)))
				    (let-values (((val ax1) (f j i ax)))
				      (vector-set! A (fx+ (cdr y+b) x) val)
				      ax1)))))))
	   
	   (else (matrix-utils:error 'fill-matrix! "unknown order " order)))
     A)))

;;
;; Given procedures MAKE-VECTOR and FILL-MATRIX!, returns a procedure
;; ONES of the form ONES:: M * N [* ORDER] -> A
;;
;; Where procedure ONES returns a matrix A of size M x N, in which all
;; elements are 1.  Optional argument ORDER specifies the matrix
;; layout: blas:ColMajor or blas:RowMajor, default is blas:RowMajor.
;;
;; MAKE-VECTOR is one of the homogeneous vector creation procedures
;; from SRFI-4, and FILL-MATRIX! is a procedure created by
;; MAKE-FILL-MATRIX, above. FILL-MATRIX! must operate on the same type
;; of vector as MAKE-VECTOR.
;;
(define (make-matrix-ones make-vector fill-matrix!)
  (lambda (m n . rest)
    (let-optionals rest ((order blas:RowMajor))
      (let ((A (make-vector (fx* m n))))
	(fill-matrix! order m n A (lambda (i j ax) (values 1.0 #f)) #f)))))

;;
;; Given procedures MAKE-VECTOR and FILL-MATRIX!, returns a procedure
;; ZEROS of the form ZEROS:: M * N [* ORDER] -> A
;;
;; Where procedure ZEROS returns a matrix A of size M x N, in which all
;; elements are 0.  Optional argument ORDER specifies the matrix
;; layout: blas:ColMajor or blas:RowMajor, default is blas:RowMajor.
;;
;; MAKE-VECTOR is one of the homogeneous vector creation procedures
;; from SRFI-4, and FILL-MATRIX! is a procedure created by
;; MAKE-FILL-MATRIX, above. FILL-MATRIX! must operate on the same type
;; of vector as MAKE-VECTOR.
;;
(define (make-matrix-zeros make-vector fill-matrix!)
  (lambda (m n . rest)
    (let-optionals rest ((order blas:RowMajor))
      (let ((A (make-vector (fx* m n))))
	(fill-matrix! order m n A (lambda (i j ax) (values 0.0 #f)) #f)))))


;;
;; Given procedures MAKE-VECTOR and FILL-MATRIX!, returns a procedure
;; EYE of the form EYE:: M * N [* ORDER] -> I
;;
;; Where procedure EYE returns an identity matrix of size M x N.
;; Optional argument ORDER specifies the matrix layout: blas:ColMajor
;; or blas:RowMajor, default is blas:RowMajor.
;;
;; MAKE-VECTOR is one of the homogeneous vector creation procedures
;; from SRFI-4, and FILL-MATRIX! is a procedure created by
;; MAKE-FILL-MATRIX, above. FILL-MATRIX! must operate on the same type
;; of vector as MAKE-VECTOR.
;;
(define (make-matrix-eye make-vector fill-matrix!)
  (lambda (m n . rest)
    (let-optionals rest ((order blas:RowMajor))
      (let ((A  (make-vector (fx* m n))))
	(fill-matrix! order m n A (lambda (i j ax) (values (if (fx= i j) 1.0 0.0) #f)) #f)))))

;;
;; Given procedures VECTOR-REF, MAKE-VECTOR and FILL-MATRIX!, returns
;; a procedure DIAG of the form DIAG:: M * N * V [* K * ORDER] -> D
;;
;; Where procedure DIAG returns a diagonal matrix D of size M x N,
;; with vector V on diagonal K.  Argument K is optional.  If it is
;; positive, the vector is placed on the K-th super-diagonal of matrix
;; D.  If it is negative, it is placed on the -K-th sub-diagonal of
;; matrix D.  The default value of K is 0, and the vector is placed on
;; the main diagonal of matrix D. Optional argument ORDER specifies
;; the matrix layout: blas:ColMajor or blas:RowMajor, default is
;; blas:RowMajor. Vector V is always assumed to be a row vector.
;;
;; VECTOR-REF and MAKE-VECTOR are two of the homogeneous vector
;; procedures from SRFI-4, and FILL-MATRIX! is a procedure created by
;; MAKE-FILL-MATRIX, above. FILL-MATRIX! must operate on the same type
;; of vector as VECTOR-REF and MAKE-VECTOR.
;;
(define (make-matrix-diag vector-ref make-vector fill-matrix!)
  (lambda (m n v . rest)
    (let-optionals rest ((k 0) (order blas:RowMajor))
      (let ((A  (make-vector (fx* m n)))
	    (k (if (eq? order blas:RowMajor) k (- k))))
	(fill-matrix! order m n A
		      (lambda (i j vi) 
			(if (fx= k (fx- j i))
			    (values (vector-ref v vi) (fx+ 1 vi))
			    (values 0.0 vi)))
		      0)))))

;;
;; Given procedures MAKE-VECTOR and FILL-MATRIX!, returns a procedure
;; LINSPACE of the form LINSPACE:: N * BASE * LIMIT -> V
;;
;; Where LINSPACE returns a row vector with N linearly spaced elements
;; between BASE and LIMIT.  The number of elements, N, must be greater
;; than 1.  The BASE and LIMIT are always included in the range.  If
;; BASE is greater than LIMIT, the elements are stored in decreasing
;; order.
;;
;; MAKE-VECTOR is one of the homogeneous vector creation procedures
;; from SRFI-4, and FILL-MATRIX! is a procedure created by
;; MAKE-FILL-MATRIX, above. FILL-MATRIX! must operate on the same type
;; of vector as MAKE-VECTOR.
;;
(define (make-linspace make-vector fill-matrix!)
  (lambda (n base limit)
    (if (not (> n 1)) 
	(matrix-utils:error 'linspace "vector size N must be greater than 1"))
    (let ((step  (/ (- limit base) (fx- n 1)))
	  (a     (make-vector n)))
      (fill-matrix! blas:RowMajor 1 n a 
		    (lambda (i j ax) (values (+ base (* 1.0 j step)) ax))  #f))))
      

;; Given procedures VECTOR-REF, MAKE-VECTOR and FILL-MATRIX!, returns
;; a procedure LOGSPACE of the form LOGSPACE:: N * BASE * LIMIT -> V
;;
;; Where LOGSPACE returns a row vector with elements that are
;; logarithmically spaced from 10^BASE to 10^LIMIT. The number of
;; elements, N, must be greater than 1.  The BASE and LIMIT are always
;; included in the range.  If BASE is greater than LIMIT, the elements
;; are stored in decreasing order.
;;
;; MAKE-VECTOR is one of the homogeneous vector creation procedures
;; from SRFI-4, and FILL-MATRIX! is a procedure created by
;; MAKE-FILL-MATRIX, above. FILL-MATRIX! must operate on the same type
;; of vector as MAKE-VECTOR.
;;
(define (make-logspace linspace vector-ref make-vector fill-matrix!)
  (lambda (n base limit)
    (if (not (> n 1)) 
	(matrix-utils:error 'logspace "vector size N must be greater than 1"))
    (let ((step  (/ (- limit base) (fx- n 1)))
	  (a     (make-vector n))
	  (b     (linspace n base limit)))
      (fill-matrix! blas:RowMajor 1 n a 
		    (lambda (i j ax) 
		      (let ((v  (expt 10 (vector-ref b j))))
			(values v  ax))) #f))))



;;
;; Given a procedure VECTOR-REF returns a procedure FILL-MATRIX! of
;; the form MATRIX-FOLD-PARTIAL:: M * N * A * F * P * X0 * [IX * IY * EX * EY] -> XN
;;
;; Where procedure MATRIX-FOLD-PARTIAL applies the fold operation on a
;; matrix A of size M x N with the values returned by applying
;; procedure F to each pair of indices and the corresponding value at
;; that position in the matrix. MATRIX-FOLD-PARTIAL first applies the
;; predicate P to the indices, and if P returns true, then F is
;; applied.
;;
;; Procedure F is of the form LAMBDA V AX -> AX1, where V is a
;; matrix element at position (I,J) and AX is accumulator value. The
;; initial value of AX is given by X0. Procedure F is expected to
;; return the new accumulator value.
;;
;; Procedure P is of the form LAMBDA I J -> boolean, where I and J are
;; matrix indices.
;;
;; Optional arguments IX IY EX EY may specify a sub-matrix in matrix A
;; to be filled. These arguments are checked to make sure they specify
;; a valid sub-matrix.
;;
;; VECTOR-REF is one of the homogeneous vector accessor procedure from
;; SRFI-4. 
;;

(define (make-matrix-fold-partial vector-ref)
  (lambda (order M N A f p x0 . rest) 
    ;; optional arguments to specify a sub-matrix 
    (let-optionals rest ((ix 0) (iy 0) (ex M) (ey N))
     (if (not (and (fx>= ix 0) (fx>= iy 0) (fx<= ex M) (fx<= ey N)
		   (fx<= ix ex) (fx<= iy ey)))
	 (matrix-utils:error 'matrix-fold-partial "invalid sub-matrix dimensions: " (list ix iy ex ey)))

     (cond ((= order blas:RowMajor)
	    (fold-ec x0 (:parallel (:range b (fx* N ix) (fx* N M) N) (:range x ix ex)) (cons x b)
		     (lambda (x+b ax)
		       (fold-ec ax (:range y iy ey) y
				(lambda (y ax)
				  (let ((i (fx- (car x+b) ix)) (j (fx- y iy)))
				    (if (p i j)
					(f (vector-ref A (fx+ (cdr x+b) y)) ax) ax)))))))
	   
	   ((= order blas:ColMajor)
	    (fold-ec x0 (:parallel (:range b (fx* N ix) (fx* N M) M) (:range y iy ey)) (cons  b)
		     (lambda (y+b ax)
		       (fold-ec ax (:range x ix ex) x
				(lambda (x ax)
				  (let ((i (fx- x ix)) (j (fx- (car y+b) iy)))
				    (if (p j i)
					(f (vector-ref A (fx+ (cdr y+b) x)) ax) ax)))))))
	   
	   (else (matrix-utils:error 'matrix-fold-partial "unknown order " order))))))
     
  

;;
;; Given a procedure VECTOR-REF returns a procedure MATRIX-FOLD of
;; the form MATRIX-FOLD:: M * N * A * F * X0 * [IX * IY * EX * EY] -> XN
;;
;; Where procedure MATRIX-FOLD applies the fold operation on a
;; matrix A of size M x N with the values returned by applying
;; procedure F to each pair of indices and the corresponding value at
;; that position in the matrix. 
;;
;; Procedure F is of the form LAMBDA I J V AX -> AX1, where V is a
;; matrix element at position (I,J) and AX is accumulator value. The
;; initial value of AX is given by X0. Procedure F is expected to
;; return the new accumulator value.
;;
;; Optional arguments IX IY EX EY may specify a sub-matrix in matrix A
;; to be filled. These arguments are checked to make sure they specify
;; a valid sub-matrix.
;;
;; VECTOR-REF is one of the homogeneous vector accessor procedure from
;; SRFI-4. 
;;

(define (make-matrix-fold vector-ref)
  (lambda (order M N A f x0 . rest) 
    ;; optional arguments to specify a sub-matrix 
    (let-optionals rest ((ix 0) (iy 0) (ex M) (ey N))
     (if (not (and (fx>= ix 0) (fx>= iy 0) (fx<= ex M) (fx<= ey N)
		   (fx<= ix ex) (fx<= iy ey)))
	 (matrix-utils:error 'matrix-fold "invalid sub-matrix dimensions: " (list ix iy ex ey)))

     (cond ((= order blas:RowMajor)
	    (fold-ec x0 (:parallel (:range b (fx* N ix) (fx* N M) N) (:range x ix ex)) (cons x b)
		     (lambda (x+b ax)
		       (fold-ec ax (:range y iy ey) y
				(lambda (y ax)
				  (let ((i (fx- (car x+b) ix)) (j (fx- y iy)))
				    (f i j (vector-ref A (fx+ (cdr x+b) y)) ax) ax))))))
	   
	   ((= order blas:ColMajor)
	    (fold-ec x0 (:parallel (:range b (fx* N ix) (fx* N M) M) (:range y iy ey)) (cons  b)
		     (lambda (y+b ax)
		       (fold-ec ax (:range x ix ex) x
				(lambda (x ax)
				  (let ((i (fx- x ix)) (j (fx- (car y+b) iy)))
				    (f j i (vector-ref A (fx+ (cdr y+b) x)) ax) ax))))))
	   
	   (else (matrix-utils:error 'matrix-fold "unknown order " order))))))
     

  
(define (make-matrix-map vector-ref vector-set!)
  (define fill-matrix! (make-fill-matrix vector-set!))
  (lambda (order M N A f . rest) 
    ;; optional arguments to specify a sub-matrix to be mapped
    (let-optionals rest ((ix 0) (iy 0) (ex M) (ey N))
     (if (not (and (fx>= ix 0) (fx>= iy 0) (fx<= ex M) (fx<= ey N)
		   (fx<= ix ex) (fx<= iy ey)))
	 (matrix-utils:error 'matrix-map "invalid sub-matrix dimensions: " (list ix iy ex ey)))
     (cond ((= order blas:RowMajor)
	    (fill-matrix! order M N A (lambda (i j ax) 
					(let ((v  (vector-ref A (fx+ j (fx* i M))))) (values (f i j v) #f))) #f))
	   ((= order blas:ColMajor)
	    (fill-matrix! order M N A (lambda (i j ax) 
					(let ((v  (vector-ref A (fx+ i (fx* j N))))) (values (f j i v) #f))) #f))
	   (else (matrix-utils:error 'matrix-map "unknown order " order))))))
			    


(define-syntax define-utility-matrices 
  (lambda (x r c)
    (let* ((type (cadr x))
	   (prims  (case type
		     ((f64)    `(list f64vector-ref make-f64vector (make-fill-matrix f64vector-set!)))
		     ((f32)    `(list f32vector-ref make-f32vector (make-fill-matrix f32vector-set!)))
		     ((s32)    `(list s64vector-ref make-s64vector (make-fill-matrix s64vector-set!)))
		     ((u32)    `(list u32vector-ref make-u32vector (make-fill-matrix u32vector-set!)))
		     (else     `(if (list? ,type)
				    (list (alist-ref 'vector-ref ,type)
					  (alist-ref 'make-vector ,type)
					  (make-fill-matrix (alist-ref 'vector-set! ,type)))
				    (matrix-utils:error 'define-matrix-type "invalid type " ,type)))))
	   (%begin  (r 'begin))
	   (%define (r 'define)))

     `(,%begin
	(,%define matrix-primitives ,prims)
	(,%define fill-matrix! (caddr matrix-primitives))
	(,%define matrix-ones  (make-matrix-ones (cadr matrix-primitives) fill-matrix!))
	(,%define matrix-zeros (make-matrix-zeros (cadr matrix-primitives) fill-matrix!))
	(,%define matrix-eye   (make-matrix-eye (cadr matrix-primitives) fill-matrix!))
	(,%define matrix-diag  (make-matrix-diag (car matrix-primitives) (cadr matrix-primitives) fill-matrix!))
	(,%define linspace (make-linspace (cadr matrix-primitives) fill-matrix!))
	(,%define logspace (make-logspace linspace (car matrix-primitives) (cadr matrix-primitives) fill-matrix!)))))
  )
	

(define-syntax with-utility-matrices
  (lambda (x r c)
    (let* ((type (cadr x))
	   (expr (caddr x))
	   (prims  (case type
		     ((f64)    `(list f64vector-ref make-f64vector (make-fill-matrix f64vector-set!)))
		     ((f32)    `(list f32vector-ref make-f32vector (make-fill-matrix f32vector-set!)))
		     ((s32)    `(list s64vector-ref make-s64vector (make-fill-matrix s64vector-set!)))
		     ((u32)    `(list u32vector-ref make-u32vector (make-fill-matrix u32vector-set!)))
		     (else     `(if (list? ,type)
				    (list (alist-ref 'vector-ref ,type)
					  (alist-ref 'make-vector ,type)
					  (make-fill-matrix (alist-ref 'vector-set! ,type)))
				    (matrix-utils:error 'define-matrix-type "invalid type " ,type)))))
	   (%let* (r 'let*)))
     `(,%let* ((matrix-primitives ,prims)
	       (fill-matrix! (caddr matrix-primitives))
	       (matrix-ones  (make-matrix-ones (cadr matrix-primitives) fill-matrix!))
	       (matrix-zeros (make-matrix-zeros (cadr matrix-primitives) fill-matrix!))
	       (matrix-eye   (make-matrix-eye (cadr matrix-primitives) fill-matrix!))
	       (matrix-diag  (make-matrix-diag (car matrix-primitives) (cadr matrix-primitives) fill-matrix!))
	       (linspace (make-linspace (cadr matrix-primitives) fill-matrix!))
	       (logspace (make-logspace linspace (car matrix-primitives) (cadr matrix-primitives) fill-matrix!)))
	      ,expr)))
  )

)