(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 solve 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 (fl< c residual-limit)
					    (fx= 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.
      (if (stop-iterations? A preconditioner z p r x c m)
	  x
	  (let* ((Ap (Operator-apply A p))
		 (alpha (fl/ c (Vector-dot-product Ap p))))
	    (let ((x (Vector-*axpy alpha p x))
		  (r (Vector-*axpy (fl- alpha) Ap r)))
	      (let* ((z (Operator-apply preconditioner r))
		     (d (Vector-dot-product r z))
		     (p (Vector-*axpy (fl/ d c) p z)))
		(loop z p r x d (fx+ 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))
  (let ((indices indices)
	(coefficients coefficients)
	(offsets offsets))
    (let ((diagonal-inverse 
	   (let* ((n (fx- (vector-length offsets) 1))
		  (result (make-f64vector n)))
	     (do ((i 0 (fx+ i 1)))
		 ((fx= i n) result)
	       (f64vector-set! result i (fl/ (f64vector-ref coefficients (fx- (vector-ref offsets (fx+ 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))
				(do ((i (fx- n 1) (fx- i 1)))
				    ((fx< i 0) result)
				  (f64vector-set! result i (fl* (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 (Sparse-matrix-operator-range-dimension m) 0.0)))
    (let ((offsets offsets)
	  (coefficients coefficients)
	  (indices indices))
      (let ((n range-dimension))
	(let lower-loop ((i 0) (offset 0))
	  (if (fx< i n)
	      (let ((next-offset (vector-ref offsets (fx+ i 1))))
		(cond ((fx= (fx- next-offset offset) 7)
		       (f64vector-set! result i (fl/ (fl- (f64vector-ref rhs i)
							  (fl* (f64vector-ref coefficients offset) 
							       (f64vector-ref result (indices-ref indices offset)))
							  (fl* (f64vector-ref coefficients                (fx+ offset 1))
							       (f64vector-ref result (indices-ref indices (fx+ offset 1))))
							  (fl* (f64vector-ref coefficients                (fx+ offset 2)) 
							       (f64vector-ref result (indices-ref indices (fx+ offset 2))))
							  (fl* (f64vector-ref coefficients                (fx+ offset 3)) 
							       (f64vector-ref result (indices-ref indices (fx+ offset 3))))
							  (fl* (f64vector-ref coefficients                (fx+ offset 4)) 
							       (f64vector-ref result (indices-ref indices (fx+ offset 4))))
							  (fl* (f64vector-ref coefficients                (fx+ offset 5)) 
							       (f64vector-ref result (indices-ref indices (fx+ offset 5)))))
						     (f64vector-ref coefficients (fx+ offset 6))))
		       (lower-loop (fx+ i 1) next-offset))
		      (else
		       (do ((j (fx- next-offset 2) (fx- j 1))
			    (diff (f64vector-ref rhs i) (fl- diff (fl* (f64vector-ref coefficients j) 
								       (f64vector-ref result (indices-ref indices j))))))
			   ((fx< j offset)
			    (f64vector-set! result i (fl/ diff (f64vector-ref coefficients (fx- next-offset 1))))
			    (lower-loop (fx+ i 1) next-offset))))))
	      (let upper-loop ((i (fx- n 1)) (previous-offset (f64vector-length coefficients)))
		(if (fx>= i 0)
		    (let ((offset (vector-ref offsets i)))
		      (cond ((fx= (fx- previous-offset offset) 7)
			     (f64vector-set! result i (fl/ (fl- (f64vector-ref rhs i)
								(fl* (f64vector-ref coefficients offset) 
								     (f64vector-ref result (indices-ref indices offset)))
								(fl* (f64vector-ref coefficients                (fx+ offset 1))
								     (f64vector-ref result (indices-ref indices (fx+ offset 1))))
								(fl* (f64vector-ref coefficients                (fx+ offset 2)) 
								     (f64vector-ref result (indices-ref indices (fx+ offset 2))))
								(fl* (f64vector-ref coefficients                (fx+ offset 3)) 
								     (f64vector-ref result (indices-ref indices (fx+ offset 3))))
								(fl* (f64vector-ref coefficients                (fx+ offset 4)) 
								     (f64vector-ref result (indices-ref indices (fx+ offset 4))))
								(fl* (f64vector-ref coefficients                (fx+ offset 5)) 
								     (f64vector-ref result (indices-ref indices (fx+ offset 5)))))
							   (f64vector-ref coefficients (fx+ offset 6))))
			     (upper-loop (fx- i 1) offset))
			    (else
			     (do ((j (fx- previous-offset 2) (fx- j 1))
				  (diff (f64vector-ref rhs i) (fl- diff (fl* (f64vector-ref coefficients j) 
									     (f64vector-ref result (indices-ref indices j))))))
				 ((fx< j offset)
				  (f64vector-set! result i (fl/ diff (f64vector-ref coefficients (fx- previous-offset 1))))
				  (upper-loop (fx- 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 (fl/ (Vector-Linf-norm m))))
    (lambda (rhs input)
      (let ((n range-dimension))
	(let ((offsets offsets)
	      (coefficients coefficients)
	      (indices indices)
	      (result (Vector-copy input)))
	  (let loop ((i 0) (offset 0))
	    (if (fx< i n)
		(let ((next-offset (vector-ref offsets (fx+ i 1))))
		  (cond ((fx= (fx- next-offset offset) 7)
			 (f64vector-set! result i (fl+ (f64vector-ref result i)
						       (fl* (fl- (f64vector-ref rhs i)
								 (fl* (f64vector-ref coefficients offset) 
								      (f64vector-ref input (indices-ref indices offset)))
								 (fl* (f64vector-ref coefficients               (fx+ offset 1))
								      (f64vector-ref input (indices-ref indices (fx+ offset 1))))
								 (fl* (f64vector-ref coefficients               (fx+ offset 2)) 
								      (f64vector-ref input (indices-ref indices (fx+ offset 2))))
								 (fl* (f64vector-ref coefficients               (fx+ offset 3)) 
								      (f64vector-ref input (indices-ref indices (fx+ offset 3))))
								 (fl* (f64vector-ref coefficients               (fx+ offset 4)) 
								      (f64vector-ref input (indices-ref indices (fx+ offset 4))))
								 (fl* (f64vector-ref coefficients               (fx+ offset 5)) 
								      (f64vector-ref input (indices-ref indices (fx+ offset 5))))
								 (fl* (f64vector-ref coefficients               (fx+ offset 6)) 
								      (f64vector-ref input (indices-ref indices (fx+ offset 6)))))
							    lambda-inverse)))
			 (loop (fx+ i 1) next-offset))
			(else
			 (do ((j (fx- next-offset 1) (fx- j 1))
			      (diff (f64vector-ref rhs i) (fl- diff (fl* (f64vector-ref coefficients j) 
									 (f64vector-ref input (indices-ref indices j))))))
			     ((fx< j offset)
			      (f64vector-set! result i (fl+ (f64vector-ref result i) (fl* diff lambda-inverse)))
			      (loop (fx+ 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)
    (let ((n range-dimension))
      (let ((offsets offsets)
	    (coefficients coefficients)
	    (indices indices)
	    (result (make-f64vector n 0.0)))
	(let loop ((i 0) (offset 0))
	  (if (fx< i n)
	      (let ((next-offset (vector-ref offsets (fx+ i 1))))
		(cond ((fx= (fx- next-offset offset) 7)
		       (f64vector-set! result i (fl/ (fl- (f64vector-ref rhs i)
							  (fl* (f64vector-ref coefficients offset) 
							       (f64vector-ref input (indices-ref indices offset)))
							  (fl* (f64vector-ref coefficients               (fx+ offset 1))
							       (f64vector-ref input (indices-ref indices (fx+ offset 1))))
							  (fl* (f64vector-ref coefficients               (fx+ offset 2)) 
							       (f64vector-ref input (indices-ref indices (fx+ offset 2))))
							  (fl* (f64vector-ref coefficients               (fx+ offset 3)) 
							       (f64vector-ref input (indices-ref indices (fx+ offset 3))))
							  (fl* (f64vector-ref coefficients               (fx+ offset 4)) 
							       (f64vector-ref input (indices-ref indices (fx+ offset 4))))
							  (fl* (f64vector-ref coefficients               (fx+ offset 5)) 
							       (f64vector-ref input (indices-ref indices (fx+ offset 5)))))
						     (f64vector-ref coefficients (fx+ offset 6))))
		       (loop (fx+ i 1) next-offset))
		      (else
		       (do ((j (fx- next-offset 2) (fx- j 1))
			    (diff (f64vector-ref rhs i) (fl- diff (fl* (f64vector-ref coefficients j) 
								       (f64vector-ref input (indices-ref indices j))))))
			   ((fx< j offset)
			    (f64vector-set! result i (fl/ diff (f64vector-ref coefficients (fx- next-offset 1))))
			    (loop (fx+ 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.