(define-class Conjugate-gradient-data Object
  ((= solver-description maybe-uninitialized:)   ; can be used to identify the solver
   (= triangulation)
   (= space)
   (= operator maybe-uninitialized:)             ; we want to solve operator x = right-hand-side
   (= right-hand-side maybe-uninitialized:)
   (= solver maybe-uninitialized:)))             ; a thunk to solver the problem

(define (CG-build-data triangulation space)      ; just adds triangulation and space to Conjugate-gradient-data
  (instantiate Conjugate-gradient-data
	       triangulation: triangulation
	       space: space))

(define (CG-add-operator-and-rhs-information! cg-data problem-data) ; builds the operator and the right-hand-side from problem-data
  (with-access
   cg-data (Conjugate-gradient-data triangulation space operator right-hand-side)
   
   (with-access
    problem-data (Problem-data a c b gamma f g)
    
    (let ((m (->Linear-finite-element-operator space))
	  (rhs (instantiate Linear-finite-element-vector
			    space: space)))
      
      (cond (a     => (lambda (a) (add-stiffness-terms! m a))))
      (cond (c     => (lambda (c) (add-mass-terms! m c))))
      (cond (b     => (lambda (b) (add-transport-terms! m b))))
      (cond (gamma => (lambda (gamma) (add-boundary-terms! m gamma))))
      (cond (f     => (lambda (f) (add-interior-integral-terms! rhs f))))
      (cond (g     => (lambda (g) (add-boundary-integral-terms! rhs g))))
      
      (set! operator m)
      (set! right-hand-side rhs)))))

(define (CG-add-solver-information! cg-data preconditioner-operator preconditioner-description iteration-limit residual-limit)
  ;; adds solver-description and builds solver using just the iteration-limit and resitual-limit
  (with-access
   cg-data (Conjugate-gradient-data triangulation space operator right-hand-side solver solver-description)
   (set! solver-description (string-append "CG: " preconditioner-description))
   (set! solver (lambda ()
		  (conjugate-gradient operator
				      right-hand-side
				      (lambda (A preconditioner z p r x c m)
					(or (FLOAT (< c residual-limit))
					    (FIX (= m iteration-limit))))
				      preconditioner-operator)))))
	 
   

;;; This routine uses the notation in Carl de Boor's notes on
;;; the preconditioned conjugate gradient method.

(define (conjugate-gradient A                ; an Operator, we solve Ax=b
			    b                ; a Vector
			    stop-iterations? ; a call-back function, see below
			    #!optional 
			    (preconditioner identity-operator))
  (let ((z (Operator-apply preconditioner b)))
    (let loop ((z z)                        ; (Operator-apply preconditioner r)
	       (p z)                        ; search direction
	       (r b)                        ; residual
	       (x (Vector-zero b))          ; approximate solution
	       (c (Vector-dot-product b z)) ; z \cdot r
	       (m 0))                       ; number of iterations so far.
      (declare (flonum))
      (if (stop-iterations? A preconditioner z p r x c m)
	  x
	  (let* ((Ap (Operator-apply A p))
		 (alpha (/ c (Vector-dot-product Ap p))))
	    (let ((x (Vector-*axpy alpha p x))
		  (r (Vector-*axpy (- alpha) Ap r)))
	      (let* ((z (Operator-apply preconditioner r))
		     (d (Vector-dot-product r z))
		     (p (Vector-*axpy (/ d c) p z)))
		(loop z p r x d (FIX (+ m 1))))))))))

;; If we write m = L + D + L^T for a symmetric, positive-definite matrix m, then
;; this operator calculates D^{-1}.

(define-generic (diagonal-preconditioner (m)))

(define-handy-method (diagonal-preconditioner (m Sparse-matrix-operator))
  (declare (fixnum))
  (let ((indices indices)
	(coefficients coefficients)
	(offsets offsets))
    (let ((diagonal-inverse 
	   (let* ((n (- (vector-length offsets) 1))
		  (result (make-f64vector n)))
	     (do ((i 0 (+ i 1)))
		 ((= i n) result)
	       (f64vector-set! result i (FLOAT (/ (f64vector-ref coefficients (FIX (- (vector-ref offsets (+ i 1)) 1))))))))))
      (instantiate Operator
		   domain?: range?           ; this is an approximate inverse to m, so the domain and range
		   range?: domain?           ; are swapped.
		   mapping: (lambda (data)
			      (let* ((n (f64vector-length data))
				     (result (make-f64vector n))
				     (diagonal-inverse diagonal-inverse))
                                (declare (not interrupts-enabled))
				(do ((i (- n 1) (- i 1)))
				    ((< i 0) result)
				  (f64vector-set! result i (FLOAT (* (f64vector-ref data i) 
								     (f64vector-ref diagonal-inverse i)))))))))))

(define-handy-method (diagonal-preconditioner (m Linear-finite-element-operator))
  (instantiate Operator
	       domain?: range?
	       range?: domain?
	       mapping: (let ((matrix-operator (diagonal-preconditioner matrix-implementation)))
			  (lambda (v)
			    (instantiate-from (v Linear-finite-element-vector)
			       data: (Operator-apply matrix-operator (Linear-finite-element-vector-data v)))))))

(define-generic (gauss-seidel-smoother! (m)))

(define-handy-method (gauss-seidel-smoother! (m Sparse-matrix-operator))
  (lambda (rhs #!optional (result (make-f64vector (FIX (- (vector-length (Sparse-matrix-operator-offsets m)) 1)) 0.0)))
    (declare (fixnum))
    (let ((offsets offsets)
	  (coefficients coefficients)
	  (indices indices))
      (let ((n (- (vector-length offsets) 1)))
        (declare (not interrupts-enabled))
	(let lower-loop ((i 0) (offset 0))
	  (if (< i n)
	      (let ((next-offset (vector-ref offsets (+ i 1))))
		(cond ((= (- next-offset offset) 7)
		       (f64vector-set! result i (FLOAT (/ (- (f64vector-ref rhs i)
							     (* (f64vector-ref coefficients offset) 
								(f64vector-ref result (indices-ref indices offset)))
							     (* (f64vector-ref coefficients (FIX (+ offset 1))) 
								(f64vector-ref result (indices-ref indices (FIX (+ offset 1)))))
  							     (* (f64vector-ref coefficients (FIX (+ offset 2))) 
								(f64vector-ref result (indices-ref indices (FIX (+ offset 2)))))
							     (* (f64vector-ref coefficients (FIX (+ offset 3))) 
								(f64vector-ref result (indices-ref indices (FIX (+ offset 3)))))
							     (* (f64vector-ref coefficients (FIX (+ offset 4))) 
								(f64vector-ref result (indices-ref indices (FIX (+ offset 4)))))
							     (* (f64vector-ref coefficients (FIX (+ offset 5))) 
								(f64vector-ref result (indices-ref indices (FIX (+ offset 5))))))
							  (f64vector-ref coefficients (FIX (+ offset 6))))))
		       (lower-loop (+ i 1) next-offset))
		      (else
		       (do ((j (- next-offset 2) (- j 1))
			    (diff (f64vector-ref rhs i) (FLOAT (- diff (* (f64vector-ref coefficients j) 
									  (f64vector-ref result (indices-ref indices j)))))))
			   ((< j offset)
			    (f64vector-set! result i (FLOAT (/ diff (f64vector-ref coefficients (FIX (- next-offset 1))))))
			    (lower-loop (+ i 1) next-offset))))))
	      (let upper-loop ((i (- n 1)) (previous-offset (f64vector-length coefficients)))
		(if (>= i 0)
		    (let ((offset (vector-ref offsets i)))
		      (cond ((= (- previous-offset offset) 7)
			     (f64vector-set! result i (FLOAT (/ (- (f64vector-ref rhs i)
								   (* (f64vector-ref coefficients offset) 
								      (f64vector-ref result (indices-ref indices offset)))
								   (* (f64vector-ref coefficients (FIX (+ offset 1))) 
								      (f64vector-ref result (indices-ref indices (FIX (+ offset 1)))))
								   (* (f64vector-ref coefficients (FIX (+ offset 2))) 
								      (f64vector-ref result (indices-ref indices (FIX (+ offset 2)))))
								   (* (f64vector-ref coefficients (FIX (+ offset 3))) 
								      (f64vector-ref result (indices-ref indices (FIX (+ offset 3)))))
								   (* (f64vector-ref coefficients (FIX (+ offset 4))) 
								      (f64vector-ref result (indices-ref indices (FIX (+ offset 4)))))
								   (* (f64vector-ref coefficients (FIX (+ offset 5))) 
								      (f64vector-ref result (indices-ref indices (FIX (+ offset 5))))))
								(f64vector-ref coefficients (FIX (+ offset 6))))))
			     (upper-loop (- i 1) offset))
			    (else
			     (do ((j (- previous-offset 2) (- j 1))
				  (diff (f64vector-ref rhs i) (FLOAT (- diff (* (f64vector-ref coefficients j) 
										(f64vector-ref result (indices-ref indices j)))))))
				 ((< j offset)
				  (f64vector-set! result i (FLOAT (/ diff (f64vector-ref coefficients (FIX (- previous-offset 1))))))
				  (upper-loop (- i 1) offset))))))
		    result))))))))

(define-handy-method (gauss-seidel-smoother! (m Linear-finite-element-operator))
  (let ((matrix-smoother! (gauss-seidel-smoother! matrix-implementation)))
    (lambda (rhs #!optional (result (instantiate Linear-finite-element-vector space: (Linear-finite-element-operator-range-space m))))
      (matrix-smoother! (Linear-finite-element-vector-data rhs)
			(Linear-finite-element-vector-data result))
      result)))

(define-generic (richardson-smoother! (m)))

(define-handy-method (richardson-smoother! (m Sparse-matrix-operator))
  (let ((lambda-inverse (FLOAT (/ (* 0.625 (Vector-Linf-norm m))))))
  (lambda (rhs input)
    (declare (fixnum))
    (let ((n (- (vector-length offsets) 1)))
      (let ((offsets offsets)
	    (coefficients coefficients)
	    (indices indices)
	    (result (Vector-copy input)))
        (declare (not interrupts-enabled))
	(let loop ((i 0) (offset 0))
	  (if (< i n)
	      (let ((next-offset (vector-ref offsets (+ i 1))))
		(cond ((= (- next-offset offset) 7)
		       (f64vector-set! result i (FLOAT (+ (f64vector-ref result i)
							  (* (- (f64vector-ref rhs i)
								(* (f64vector-ref coefficients offset) 
								   (f64vector-ref input (indices-ref indices offset)))
								(* (f64vector-ref coefficients (FIX (+ offset 1))) 
								   (f64vector-ref input (indices-ref indices (FIX (+ offset 1)))))
								(* (f64vector-ref coefficients (FIX (+ offset 2))) 
								   (f64vector-ref input (indices-ref indices (FIX (+ offset 2)))))
								(* (f64vector-ref coefficients (FIX (+ offset 3))) 
								   (f64vector-ref input (indices-ref indices (FIX (+ offset 3)))))
								(* (f64vector-ref coefficients (FIX (+ offset 4))) 
								   (f64vector-ref input (indices-ref indices (FIX (+ offset 4)))))
								(* (f64vector-ref coefficients (FIX (+ offset 5))) 
								   (f64vector-ref input (indices-ref indices (FIX (+ offset 5)))))
								(* (f64vector-ref coefficients (FIX (+ offset 6))) 
								   (f64vector-ref input (indices-ref indices (FIX (+ offset 6))))))
							     lambda-inverse))))
		       (loop (+ i 1) next-offset))
		      (else
		       (do ((j (- next-offset 1) (- j 1))
			    (diff (f64vector-ref rhs i) (FLOAT (- diff (* (f64vector-ref coefficients j) 
									  (f64vector-ref input (indices-ref indices j)))))))
			   ((< j offset)
			    (f64vector-set! result i (FLOAT (+ (f64vector-ref result i) (* diff lambda-inverse))))
			    (loop (+ i 1) next-offset))))))
	      result)))))))

(define-handy-method (richardson-smoother! (m Linear-finite-element-operator))
  (let ((matrix-smoother! (richardson-smoother! matrix-implementation)))
    (lambda (rhs input)
      (instantiate-from (input Linear-finite-element-vector)
        data: (matrix-smoother! (Linear-finite-element-vector-data rhs)
				(Linear-finite-element-vector-data input))))))



(define-generic (jacobi-smoother! (m)))

(define-handy-method (jacobi-smoother! (m Sparse-matrix-operator))
  (lambda (rhs input)
    (declare (fixnum))
    (let ((n (- (vector-length offsets) 1)))
      (let ((offsets offsets)
	    (coefficients coefficients)
	    (indices indices)
	    (result (make-f64vector n 0.0)))
        (declare (not interrupts-enabled))
	(let loop ((i 0) (offset 0))
	  (if (< i n)
	      (let ((next-offset (vector-ref offsets (+ i 1))))
		(cond ((= (- next-offset offset) 7)
		       (f64vector-set! result i (FLOAT (/ (- (f64vector-ref rhs i)
							     (* (f64vector-ref coefficients offset) 
								(f64vector-ref input (indices-ref indices offset)))
							     (* (f64vector-ref coefficients (FIX (+ offset 1))) 
								(f64vector-ref input (indices-ref indices (FIX (+ offset 1)))))
  							     (* (f64vector-ref coefficients (FIX (+ offset 2))) 
								(f64vector-ref input (indices-ref indices (FIX (+ offset 2)))))
							     (* (f64vector-ref coefficients (FIX (+ offset 3))) 
								(f64vector-ref input (indices-ref indices (FIX (+ offset 3)))))
							     (* (f64vector-ref coefficients (FIX (+ offset 4))) 
								(f64vector-ref input (indices-ref indices (FIX (+ offset 4)))))
							     (* (f64vector-ref coefficients (FIX (+ offset 5))) 
								(f64vector-ref input (indices-ref indices (FIX (+ offset 5))))))
							  (f64vector-ref coefficients (FIX (+ offset 6))))))
		       (loop (+ i 1) next-offset))
		      (else
		       (do ((j (- next-offset 2) (- j 1))
			    (diff (f64vector-ref rhs i) (FLOAT (- diff (* (f64vector-ref coefficients j) 
									  (f64vector-ref input (indices-ref indices j)))))))
			   ((< j offset)
			    (f64vector-set! result i (FLOAT (/ diff (f64vector-ref coefficients (FIX (- next-offset 1))))))
			    (loop (+ i 1) next-offset))))))
	      result))))))

(define-handy-method (jacobi-smoother! (m Linear-finite-element-operator))
  (let ((matrix-smoother! (jacobi-smoother! matrix-implementation)))
    (lambda (rhs input)
      (instantiate-from (input Linear-finite-element-vector)
        data: (matrix-smoother! (Linear-finite-element-vector-data rhs)
				(Linear-finite-element-vector-data input))))))


(define-generic (gauss-seidel-preconditioner (m)))

(define-handy-method (gauss-seidel-preconditioner (m Linear-finite-element-operator))
  (let ((smoother! (gauss-seidel-smoother! m)))
    (instantiate Operator
		 domain?: range?
		 range?: domain?
		 mapping: (lambda (v) (smoother! v))))) ;; we make sure that smoother! gets only one argument, so it's functional.