;;; This file contains a lot of hackery to make the Sparse-vector implementation
;;; run quickly.  The macros are defined as internal-meroon-macros so they can't be used
;;; outside this file (or files that include this file).

;;; An Entry is an index-coefficient pair

(define-internal-meroon-macro (make-Entry index coefficient)
  `(let ((index ,index) (coefficient ,coefficient))
     (cons index coefficient)))

(define-internal-meroon-macro (Entry-index e)
  `(let ((e ,e))
     (car e)))

(define-internal-meroon-macro (Entry-coefficient e)
  `(let ((e ,e))
     (cdr e)))

;;; A Sparse-vector is a *sorted* list of entries with nonzero Entry-coefficients.  That's it.
;;; The overhead of defining and using a Meroon class for this structure is not worth
;;; the extra ease in programming since these Sparse-vectors are usually really short.

(define-internal-meroon-macro (add-to-entry-list e l)
  `(let ((e ,e) (l ,l))
     (if (FLOAT (zero? (Entry-coefficient e)))
	 l
	 (cons e l))))

(define Sparse-vector-zero '())

(define (Sparse-vector-add x-entries y-entries)
  (declare (fixnum))
  (let loop ((x-entries x-entries) 
	     (y-entries y-entries)
	     (result-entries Sparse-vector-zero))
    (cond ((null? x-entries)
	   (reverse-append! result-entries 
			    y-entries))
	  ((null? y-entries)
	   (reverse-append! result-entries 
			    x-entries))
	  (else
	   (let ((x-entry (car x-entries))
		 (y-entry (car y-entries)))
	     (let ((x-index (Entry-index x-entry))
		   (y-index (Entry-index y-entry)))
	       (cond ((< x-index y-index)
		      (loop (cdr x-entries)
			    y-entries
			    (cons x-entry result-entries)))
		     ((< y-index x-index)
		      (loop x-entries
			    (cdr y-entries)
			    (cons y-entry result-entries)))
		     (else
		      (loop (cdr x-entries)
			    (cdr y-entries)
			    (add-to-entry-list (make-Entry x-index
							   (FLOAT (+ (Entry-coefficient x-entry)
								     (Entry-coefficient y-entry))))
					       result-entries))))))))))

(define (Sparse-vector-subtract x-entries y-entries)
  (declare (fixnum))
  (let loop ((x-entries x-entries) 
	     (y-entries y-entries)
	     (result-entries Sparse-vector-zero))
    (cond ((null? x-entries)
	   (reverse-append! result-entries (map1 (lambda (e)  ; this map is OK, because the Entry-coefficients in y are nonzero.
						   (make-Entry (Entry-index e)
							       (FLOAT (- (Entry-coefficient e)))))
						 y-entries)))
	  ((null? y-entries)
	   (reverse-append! result-entries x-entries))
	  (else
	   (let ((x-entry (car x-entries))
		 (y-entry (car y-entries)))
	     (let ((x-index (Entry-index x-entry))
		   (y-index (Entry-index y-entry)))
	       (cond ((< x-index y-index)
		      (loop (cdr x-entries)
			    y-entries
			    (cons x-entry result-entries)))
		     ((< y-index x-index)
		      (loop x-entries
			    (cdr y-entries)
			    (cons (make-Entry y-index
					      (FLOAT (- (Entry-coefficient y-entry)))) 
				  result-entries)))
		     (else
		      (loop (cdr x-entries)
			    (cdr y-entries)
			    (add-to-entry-list (make-Entry x-index
							   (FLOAT (- (Entry-coefficient x-entry)
								     (Entry-coefficient y-entry))))
					       result-entries))))))))))

(define (Sparse-vector-scale x y-entries)
  (if (FLOAT (zero? x))
      Sparse-vector-zero
      (let loop ((y-entries y-entries)
		 (result-entries Sparse-vector-zero))
	(if (null? y-entries)
	    (reverse-append! result-entries '())
	    (let ((y-entry (car y-entries)))
	      (loop (cdr y-entries)
		    ;; we need to use add-to-entry-list here because (* x (Entry-coefficient y-entry))
		    ;; can be zero even when x and (Entry-coefficient y-entry) are not.
		    (add-to-entry-list (make-Entry (Entry-index y-entry)
						   (FLOAT (* x (Entry-coefficient y-entry))))
				       result-entries)))))))

(define (Sparse-vector-*axpy a x-entries y-entries)
  (declare (fixnum))
  (if (FLOAT (zero? a))
      y-entries
      (let loop ((x-entries x-entries) 
		 (y-entries y-entries)
		 (result-entries Sparse-vector-zero))
	(cond ((null? x-entries)
	       (reverse-append! result-entries y-entries))
	      ((null? y-entries)
	       (reverse-append! result-entries (Sparse-vector-scale a x-entries)))
	      (else
	       (let ((x-entry (car x-entries))
		     (y-entry (car y-entries)))
		 (let ((x-index (Entry-index x-entry))
		       (y-index (Entry-index y-entry)))
		   (cond ((< x-index y-index)
			  (loop (cdr x-entries)
				y-entries
				(add-to-entry-list (make-Entry x-index
							       (FLOAT (* a (Entry-coefficient x-entry))))
						   result-entries)))
			 ((< y-index x-index)
			  (loop x-entries
				(cdr y-entries)
				(cons y-entry result-entries)))
			 (else
			  (loop (cdr x-entries)
				(cdr y-entries)
				(add-to-entry-list (make-Entry x-index
							       (FLOAT (+ (* a (Entry-coefficient x-entry))
									 (Entry-coefficient y-entry))))
						   result-entries)))))))))))

;; I've decided to define a Sparse-matrix class that
;; cannot be applied as an operator, and whose rows are Sparse-vectors.  It should be easier to
;; define the Vector-* methods on this class than on the Sparse-matrix-operator class, and we can
;; convert back and forth as needed.

(define-class Sparse-matrix Matrix-operator
  ((= rows)))                 ; a list of Sparse-vectors

(define-handy-method (check-same-class (m Sparse-matrix) n)
  (and (call-next-method)
       (eq? domain-dimension (Matrix-operator-domain-dimension n))
       (eq? range-dimension (Matrix-operator-range-dimension n))))
       

(define-handy-method (Vector-add (m Sparse-matrix) n)
  (if (check-same-class m n)
      (instantiate Sparse-matrix
		   rows: (map2 Sparse-vector-add rows (Sparse-matrix-rows n))
		   domain-dimension: domain-dimension
		   range-dimension: range-dimension)
      (error "Vector-add: second argument is not the same class as first: " m n)))

(define-handy-method (Vector-subtract (m Sparse-matrix) n)
  (if (check-same-class m n)
      (instantiate Sparse-matrix
		   rows: (map2 Sparse-vector-subtract rows (Sparse-matrix-rows n))
		   domain-dimension: domain-dimension
		   range-dimension: range-dimension)
      (error "Vector-subtract: second argument is not the same class as first: " m n)))

(define-handy-method (Vector-scale a (m Sparse-matrix))
  (instantiate Sparse-matrix
	       rows: (map1 (lambda (r) (Sparse-vector-scale a r)) rows)
	       domain-dimension: domain-dimension
	       range-dimension: range-dimension))

(define-handy-method (Vector-*axpy a (m Sparse-matrix) n)
  (if (check-same-class m n)
      (instantiate Sparse-matrix
		   rows: (map2 (lambda (r1 r2) (Sparse-vector-*axpy a r1 r2))
			      rows (Sparse-matrix-rows n))
		   domain-dimension: domain-dimension
		   range-dimension: range-dimension)
      (error "Vector-*axpy: third argument is not the same class as second: " m n)))

(define-handy-method (Operator-adjoint (m Sparse-matrix))
  (FIX
   (let ((old-rows (list->vector rows))
	 (new-rows (make-vector domain-dimension '())))
     (do ((i (- range-dimension 1) (- i 1)))
	 ((< i 0) (instantiate Sparse-matrix
			       domain-dimension: range-dimension
			       range-dimension: domain-dimension
			       rows: (vector->list new-rows)))
       (for-each (lambda (e)
		   (let ((index (Entry-index e))
			 (coefficient (Entry-coefficient e)))
		     (vector-set! new-rows index (cons (make-Entry i coefficient) (vector-ref new-rows index)))))
		 (vector-ref old-rows i))))))

(define-method (Operator-compose (m Sparse-matrix) (n Sparse-matrix))
  (if (eq? (Sparse-matrix-range-dimension n)
	   (Sparse-matrix-domain-dimension m))
      (let ((rows-of-n (list->vector (Sparse-matrix-rows n))))
	(instantiate Sparse-matrix
		     domain-dimension: (Sparse-matrix-domain-dimension n)
		     range-dimension: (Sparse-matrix-range-dimension m)
		     rows: (map1 (lambda (r)
				  (reduce-from-left
				   Sparse-vector-add
				   Sparse-vector-zero
				   (map1 (lambda (e)
					  (let ((index (Entry-index e))
						(coeff (Entry-coefficient e)))
					    (Sparse-vector-scale coeff (vector-ref rows-of-n index))))
					r)))
				(Sparse-matrix-rows m))))
      (error "Operator-compose: These operators do not have compatible domains and ranges" m n)))

;; An applyable (in the sense of Operator-apply) sparse matrix class.

(define-macro (setup-indices)
  (if (##fixnum? (expt 2 32))
      `(begin
	 (define indices-ref u32vector-ref)
	 (define indices-set! u32vector-set!)
	 (define list->indices list->u32vector)
	 (define indices->list u32vector->list)
	 (define make-indices make-u32vector)
	 (define indices-length u32vector-length))
      `(begin
	 (define indices-ref vector-ref)
	 (define indices-set! vector-set!)
	 (define list->indices list->vector)
	 (define indices->list vector->list)
	 (define make-indices make-vector)
	 (define indices-length vector-length))))

(setup-indices)

;;; A Sparse-matrix-operator requires the diagonal index to be at the *end* of each row.

(define-class Sparse-matrix-operator Matrix-operator
  ((= offsets)                              ; a vector of offsets into indices and coefficients where the row elements start
   (= indices)                              ; a vector of nonzero indices
   (= coefficients)))                       ; a f64vector of nonzero coefficients

;; Many of our operators will have matrix implementations, and I want a way to refer to it
;; generically.

(define-class Operator-with-matrix-implementation Operator
  ((= matrix-implementation maybe-uninitialized:)))

(define-handy-method (initialize! (m Sparse-matrix-operator))
  
  ;; This must be efficient.
  
  (if (not (field-defined? m 'mapping))
      (initialize-field-value!
       m
       (lambda (v)
	 (declare (fixnum))
	 (let ((indices indices)
	       (coefficients coefficients)
	       (offsets offsets)
	       (data v))
	   (let* ((n range-dimension)
		  (result (make-f64vector n)))
             (declare (not interrupts-enabled))
	     (let outer ((i (- n 1)) (next-offset (vector-ref offsets n)))
	       (if (< i 0)
		   result
		   (let ((offset (vector-ref offsets i)))
		     (cond ((= (- next-offset offset) 7) ; interior vertices in a refined triangulation
			    (f64vector-set! 
			     result 
			     i 
			     (FLOAT
			      (+ (* (f64vector-ref coefficients offset) 
				    (f64vector-ref data (indices-ref indices offset)))
				 (let ((offset (FIX (+ offset 1))))
				   (* (f64vector-ref coefficients offset) 
				      (f64vector-ref data (indices-ref indices offset))))
				 (let ((offset (FIX (+ offset 2))))
				   (* (f64vector-ref coefficients offset) 
				      (f64vector-ref data (indices-ref indices offset))))
				 (let ((offset (FIX (+ offset 3))))
				   (* (f64vector-ref coefficients offset) 
				      (f64vector-ref data (indices-ref indices offset))))
				 (let ((offset (FIX (+ offset 4))))
				   (* (f64vector-ref coefficients offset) 
				      (f64vector-ref data (indices-ref indices offset))))
				 (let ((offset (FIX (+ offset 5))))
				   (* (f64vector-ref coefficients offset) 
				      (f64vector-ref data (indices-ref indices offset))))
				 (let ((offset (FIX (+ offset 6))))
				   (* (f64vector-ref coefficients offset) 
				      (f64vector-ref data (indices-ref indices offset)))))))
			    (outer (- i 1) offset))
			   ((= (- next-offset offset) 5) ; boundary vertices in a refined triangulation
			    (f64vector-set! 
			     result 
			     i 
			     (FLOAT
			      (+ (* (f64vector-ref coefficients offset) 
				    (f64vector-ref data (indices-ref indices offset)))
				 (let ((offset (FIX (+ offset 1))))
				   (* (f64vector-ref coefficients offset) 
				      (f64vector-ref data (indices-ref indices offset))))
				 (let ((offset (FIX (+ offset 2))))
				   (* (f64vector-ref coefficients offset) 
				      (f64vector-ref data (indices-ref indices offset))))
				 (let ((offset (FIX (+ offset 3))))
				   (* (f64vector-ref coefficients offset) 
				      (f64vector-ref data (indices-ref indices offset))))
				 (let ((offset (FIX (+ offset 4))))
				   (* (f64vector-ref coefficients offset) 
				      (f64vector-ref data (indices-ref indices offset)))))))
			    (outer (- i 1) offset))
			   (else
			    (do ((j (- next-offset 1) (- j 1))
				 (inner-result 0.0 (FLOAT
						    (+ inner-result 
						       (* (f64vector-ref coefficients j)
							  (f64vector-ref data (indices-ref indices j)))))))
				((< j offset) 
				 (f64vector-set! result i inner-result)
				 (outer (- i 1) offset)))))))))))
       'mapping))
  
  (call-next-method))

(define (Sparse-matrix-operator-increment-entry m i j value)
  (declare (fixnum))
  (with-access
   m (Sparse-matrix-operator indices coefficients offsets)
   (let ((indices indices)
	 (coefficients coefficients)
	 (offset (vector-ref offsets i))
	 (next-offset (vector-ref offsets (+ i 1))))
     (let loop ((k offset))
       (if (= k next-offset)
	   (error "Sparse-matrix-operator-increment-entry: can't find entry" m i j value)
	   (if (= j (indices-ref indices k))
	       (f64vector-set! coefficients k (FLOAT (+ (f64vector-ref coefficients k)
							value)))
	       (loop (+ k 1))))))))


;; these routines convert between Sparse-matrices (a list of rows, each of which is a list)
;; and Sparse-matrix-operators (with separate vectors of indices and coefficients).

(define-handy-method (->Sparse-matrix (m Sparse-matrix-operator))
  (instantiate Sparse-matrix
	       rows: (let ((indices indices)
			   (coefficients coefficients)
			   (offsets offsets))
		       (declare (fixnum))
		       (let loop ((result '()) (i (- range-dimension 1)))
			 (if (< i 0)
			     result
			     (let ((offset (vector-ref offsets i))
				   (next-offset (vector-ref offsets (+ i 1))))
			       (let inner ((entries '()) 
					   (j (- next-offset 1)))
				 (if (< j offset)
				     (loop (cons (sort! entries (lambda (x y) 
								  (< (Entry-index x) 
								     (Entry-index y))))
						 result)
					   (- i 1))
				     (inner (cons (make-Entry (indices-ref indices j)
							      (f64vector-ref coefficients j))
						  entries)
					    (- j 1))))))))
	       domain-dimension: domain-dimension
	       range-dimension: range-dimension))
	       

(define-handy-method (->Sparse-matrix-operator (m Sparse-matrix))
  (declare (fixnum))
  ;; first, put all diagonal entries, if they exist, at the end of each row
  (let ((rows (let outer ((index 0)
			  (rows rows)
			  (result '()))
		(if (null? rows)
		    (reverse-append! result '())
		    (let inner ((row (car rows))
				(non-diagonal-entries '())
				(diagonal-entry '()))
		      (if (null? row)
			  (outer (+ index 1)
				 (cdr rows)
				 (cons (reverse-append! non-diagonal-entries diagonal-entry)
				       result))
			  (if (eq? (Entry-index (car row)) index)
			      (inner (cdr row)
				     non-diagonal-entries
				     (cons (car row) diagonal-entry))
			      (inner (cdr row)
				     (cons (car row) non-diagonal-entries)
				     diagonal-entry))))))))
    (let* ((offsets (let* ((n (+ range-dimension 1))
			   (result (make-vector n)))
		      (vector-set! result 0 0)
		      (let loop ((i 1)
				 (rows rows))
			(cond ((= i n) result)
			      (else
			       (vector-set! result i (+ (vector-ref result (- i 1))
							(length (car rows))))
			       (loop (+ i 1) (cdr rows))))))))
      
      (instantiate Sparse-matrix-operator
		   indices: (list->indices (mappend1 (lambda (entries)
						       (map1 (lambda (e) (Entry-index e))
							     entries))
						     rows))
		   coefficients: (list->f64vector (mappend1 (lambda (entries)
							      (map1 (lambda (e) (Entry-coefficient e))
								    entries))
							    rows))
		   offsets: offsets
		   domain-dimension: domain-dimension
		   range-dimension: range-dimension))))

;; These methods rely on the associated methods for Sparse-matrices.

(define-method (Vector-add (m Sparse-matrix-operator) n)
  (->Sparse-matrix-operator (Vector-add (->Sparse-matrix m) (->Sparse-matrix n))))

(define-method (Vector-subtract (m Sparse-matrix-operator) n)
  (->Sparse-matrix-operator (Vector-subtract (->Sparse-matrix m) (->Sparse-matrix n))))

(define-method (Vector-scale a (m Sparse-matrix-operator))
  (->Sparse-matrix-operator (Vector-scale a (->Sparse-matrix m))))

(define-method (Vector-*axpy a (m Sparse-matrix-operator) n)
  (->Sparse-matrix-operator (Vector-*axpy a (->Sparse-matrix m) (->Sparse-matrix n))))

(define-handy-method (Vector-Linf-norm (m Sparse-matrix-operator))
  (declare (fixnum))
  (let* ((coefficients coefficients)
	 (offsets offsets)
	 (n (- (vector-length offsets) 1)))
    (let outer ((i 0) (row-max 0.0))
      (if (= i n)
	  row-max
	  (let ((next-offset (vector-ref offsets (+ i 1))))
	    (let inner ((offset (vector-ref offsets i))
			(column-sum 0.0))
	      (if (= offset next-offset)
		  (outer (+ i 1)
			 (FLOAT (max row-max column-sum)))
		  (inner (+ offset 1)
			 (FLOAT (+ column-sum (abs (f64vector-ref coefficients offset))))))))))))

;; it isn't worth converting m to a Sparse-matrix and defining this only on Sparse-matrices

(define-handy-method (Operator->identity (m Sparse-matrix-operator))
  (FIX
   (if (= domain-dimension range-dimension)
       (let ((indices (let* ((n range-dimension)
			     (result (make-indices n)))
			(do ((i 0 (+ i 1)))
			    ((= i n) result)
			  (indices-set! result i i))))
	     (offsets (let* ((n (+ range-dimension 1))
			     (result (make-vector n)))
			(do ((i 0 (+ i 1)))
			    ((= i n) result)
			  (vector-set! result i i)))))
	     (instantiate Sparse-matrix-operator
			  domain-dimension: domain-dimension
			  range-dimension: range-dimension
			  indices: indices
			  coefficients: (make-f64vector range-dimension 1.0)
			  offsets: offsets))
       (call-next-method))))

(define-method (Operator-compose (op1 Sparse-matrix-operator) (op2 Sparse-matrix-operator))
  (->Sparse-matrix-operator (Operator-compose (->Sparse-matrix op1) (->Sparse-matrix op2))))		      

(define-method (Operator-adjoint (m Sparse-matrix-operator))
  (->Sparse-matrix-operator (Operator-adjoint (->Sparse-matrix m))))