;; The multigrid algorithm for solving linear systems requires a lot of structure---
;; a family of nested (in our case) triangulations, the associated finite-element spaces,
;; projectors on the duals of the spaces, injectors on the spaces themselves, the finite-element
;; operator and right-hand-side at each level to be solved, and a bunch of solvers.

;; Once you have that ;-), the algorithm is trivial.

(define-class Multigrid-data Object
  ((= triangulations maybe-uninitialized:)
   (= spaces maybe-uninitialized:)

   ;; to project on coarser grid use (vector-ref projectors k) v

   (= projectors maybe-uninitialized:)

   ;; to inject from a coarser grid, use (vector-ref injectors k) v

   (= injectors maybe-uninitialized:)

   ;; we're trying to solve (vector-ref operators k) z = (vector-ref right-hand-sides k)

   (= operators maybe-uninitialized:)
   (= right-hand-sides maybe-uninitialized:)

   ;; a vector of iterative solvers to be used at each level

   (= solvers maybe-uninitialized:)
   (= solver-description maybe-uninitialized:)))


(define-generic (MG-solver (data)))

(define-handy-method (MG-solver (data Multigrid-data))
  (declare (fixnum))
  (let* ((n (vector-length operators))
	 (results (make-vector n)))
    (let main-loop ((k 0))
      (if (< k n)
	  (begin
	    
	    ;; Solve at each level, using the injection of the solution at the previous level
	    ;; as the initial guess unless k is zero, where we use zero as the initial guess.
	    
	    (vector-set! results k ((vector-ref solvers k)
				    (vector-ref operators k)
				    (vector-ref right-hand-sides k)
				    (if (eq? k 0)
					(Vector-zero (vector-ref right-hand-sides 0))
					(Operator-apply (vector-ref injectors k)
							(vector-ref results (- k 1))))))
	    (main-loop (+ k 1)))
	  (vector-ref results (- n 1))))))

;; Starts building a Multigrid-data object, adding the triangulation and finite-element-space
;; information.  This is set up for Linear-finite-element-spaces, but could easily be made
;; generic by passing a different function for Triangulation->Linear-finite-element-space,
;; and makint coarse<-fine and transpose generic.

(define-generic (MG-build-data (data) coarse-triangulation number-of-refinements))

(define-handy-method (MG-build-data (data Multigrid-data) coarse-triangulation number-of-refinements)
  (declare (fixnum))
  (let ((n (+ number-of-refinements 1)))
    (let ((new-triangulations (make-vector n))
	  (new-spaces (make-vector n))
	  (new-injectors (make-vector n))
	  (new-projectors (make-vector n)))
      (vector-set! new-triangulations 0 coarse-triangulation)
      (do ((i 1 (+ i 1)))
	  ((= i n))
	(vector-set! new-triangulations i (refine (vector-ref new-triangulations (- i 1)))))
      (set! triangulations new-triangulations)
      (do ((i 0 (+ i 1)))
	  ((= i n))
	(vector-set! new-spaces i (Triangulation->Linear-finite-element-space 
				   (vector-ref new-triangulations i))))
      (set! spaces new-spaces)
      (do ((i 1 (+ i 1)))
	  ((= i n))
	(vector-set! new-projectors i (coarse<-fine (vector-ref new-spaces (- i 1)) 
						    (vector-ref new-spaces i))))
      (set! projectors new-projectors)
      (do ((i 1 (+ i 1)))
	  ((= i n))
	(vector-set! new-injectors i (projector->injector (vector-ref new-projectors i))))
      (set! injectors new-injectors)
      data)))

;; Here we use the problem-data to build the matrix and right-hand-side for the FEM
;; at the finest level.

(define-generic (MG-add-operator-and-rhs-information! (multigrid-data) problem-data))

(define-handy-method (MG-add-operator-and-rhs-information! (multigrid-data Multigrid-data) problem-data)
  (with-access
   problem-data (Problem-data a c b gamma f g)
   
   (declare (fixnum))
   
   (if b  ; if b is defined, we assume m is nonsymmetric, and we solve the normal equaltions
       (error "MG-add-operator-and-rhs-information!: We cannot handle a nonsymmetric operator (b != 0)"))

   (let* ((n (vector-length triangulations))
	  (space (vector-ref spaces (- n 1))))
     
     (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! operators (make-vector n))
       (set! right-hand-sides (make-vector n))
       
       (vector-set! operators (- n 1) m)
       (vector-set! right-hand-sides (- n 1) rhs)))))

;; We separate the code that propagates the operator and right-hand-side data to
;; all the coarser levels.

(define-generic (MG-propagate-level-information! (multigrid-data)))

(define-handy-method (MG-propagate-level-information! (multigrid-data Multigrid-data))
  (declare (fixnum))
  
  (do ((i (- (vector-length operators) 1) (- i 1)))
      ((= i 0))
    (vector-set! operators (- i 1) (Operator-compose (vector-ref projectors i)
						     (Operator-compose (vector-ref operators i)
								       (vector-ref injectors i))))
    (vector-set! right-hand-sides (- i 1) (Operator-apply (vector-ref projectors i)
							  (vector-ref right-hand-sides i)))))

;;; Single multigrid step.  See Brenner and Scott for notation.

(define (MG-step
	 data                       ; Multigrid-data
	 smoothers!                 ; call with ((vector-ref smoothers! k) rhs z)
	 m1                         ; number of pre-smoothings
	 m2                         ; number of post-smoothings
	 p                          ; number of coarse corrections
	 k                          ; current refinement level, 0 is coarsest grid, k > 0
	 z0                         ; initial guess
	 rhs                        ; right-hand-side
	 )
  (declare (fixnum))

  (with-access
   data (Multigrid-data operators projectors injectors solvers)
   (let ((smoother! (vector-ref smoothers! k)))
     
     (let pre-smoothing-loop ((zl z0) (l 0))
       (if (< l m1)
	   (pre-smoothing-loop (smoother! rhs zl)
			       (+ l 1))
	   
	   (let ((post-coarse-correction
		  (if (= k 0)
		      zl
		      (let ((coarse-rhs   (Operator-apply (vector-ref projectors k) 
							  (Vector-subtract rhs (Operator-apply (vector-ref operators k) zl)))))
			(let correction-loop ((q (Vector-zero coarse-rhs)) (l 0))
			  (if (< l p)
			      (correction-loop 
			       (MG-step data smoothers! m1 m2 p (- k 1) q coarse-rhs)
			       (+ l 1))
			      (Vector-add zl (Operator-apply (vector-ref injectors k) q))))))))
	     
	     (let post-smoothing-loop ((zl post-coarse-correction) (l 0))
	       (if (< l m2)
		   (post-smoothing-loop (smoother! rhs zl)
					(+ l 1))
		   zl))))))))

;; for the finest level, uses conjugate-gradient with a ridiculously small limit on the
;; residual; for the other levels, uses MG-step (with parameters m1 m2 and p) r times.

(define-generic (MG-add-solver-information! (multigrid-data) smoother-maker smoother-name r m1 m2 p))

(define-handy-method (MG-add-solver-information! (multigrid-data Multigrid-data) smoother-maker smoother-name r m1 m2 p)
  (declare (fixnum))
  (set! solver-description (string-append "MG: "
					  smoother-name
					  ", r = "
					  (number->string r)
					  ", m1 = "
					  (number->string m1)
					  ", m2 = "
					  (number->string m2)
					  ", p = "
					  (number->string p)
					  ".  "))
  (let ((n (vector-length operators)))
    (let ((smoothers! (make-vector n)))
      (do ((k 0 (+ k 1)))
	  ((= k n))
	(vector-set! smoothers! k (smoother-maker (vector-ref operators k))))
      (set! solvers (make-vector n))
      
      (do ((k 0 (+ k 1)))
	  ((= k n))
	(vector-set! solvers k
		     (lambda (A b initial-guess)
		       (do ((i 0 (+ i 1))
			    (z initial-guess (MG-step multigrid-data
						      smoothers!
						      m1
						      m2
						      p
						      k 
						      z
						      b)))
			   ((= i r) z))))))))