(define-class Obstacle-multigrid-data Multigrid-data
  ((= obstacles maybe-uninitialized:)
   (= interpolating-projectors maybe-uninitialized:)))

(define-class Obstacle-problem-data Problem-data
  ((= obstacle)))

(define-handy-method (MG-solver (data Obstacle-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)
				    (Lattice-max (vector-ref obstacles k)
						 (if (eq? k 0)
						     (Vector-zero (vector-ref right-hand-sides 0))
						     (Operator-apply (vector-ref injectors k)
								     (vector-ref results (- k 1)))))
				    (vector-ref obstacles k)))
	    (main-loop (+ k 1)))
	  (vector-ref results (- n 1))))))

(define-handy-method (MG-build-data (data Obstacle-multigrid-data) coarse-triangulation number-of-refinements)
  (declare (fixnum))
  (call-next-method)
  (let ((n (+ number-of-refinements 1)))
    (let ((interpolating-projectors (make-vector n)))
      (do ((i 1 (+ i 1)))
	  ((= i n))
	(vector-set! interpolating-projectors
		     i
		     (wrap-matrix-implementation-in-linear-finite-element-operator
		      (->Sparse-matrix-operator
		       (instantiate Sparse-matrix
				    rows: (map (lambda (vertex)
						 (list (make-Entry (Vertex-index (Vertex-fine vertex))
								   1.0)))
					       (Triangulation-vertices->list 
						(Triangulation-vertices 
						 (vector-ref triangulations (- i 1)))))
				    domain-dimension: (Triangulation-vertices-length (vector-ref triangulations i))
				    range-dimension:  (Triangulation-vertices-length (vector-ref triangulations (- i 1)))))
		      (vector-ref spaces i)
		      (vector-ref spaces (- i 1)))))
      (initialize-field-value! data interpolating-projectors 'interpolating-projectors)))
  data)

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

(define-handy-method (MG-add-operator-and-rhs-information! (multigrid-data Obstacle-multigrid-data) problem-data)
  (call-next-method)
  (let ((n (vector-length interpolating-projectors)))
    (set! obstacles (make-vector n))
    (vector-set! obstacles 
		 (- n 1) 
		 (instantiate Linear-finite-element-vector
			      data: (let* ((obstacle (Obstacle-problem-data-obstacle problem-data))
					   (fine-t (vector-ref triangulations (- n 1)))
					   (result (make-f64vector (Triangulation-vertices-length fine-t))))
				      (Triangulation-for-each-vertex
				       (lambda (v)
					 (with-access
					  v (Vertex point index)
					  (f64vector-set! result index (obstacle point))))
				       fine-t)
				      result)
			      space: (vector-ref spaces (- n 1))))))

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

(define-handy-method (MG-propagate-level-information! (multigrid-data Obstacle-multigrid-data))
  (declare (fixnum))
  
  (call-next-method)
  (do ((i (- (vector-length operators) 1) (- i 1)))
      ((= i 0))
    (vector-set! obstacles 
		 (- i 1) 
		 (Operator-apply (vector-ref interpolating-projectors i) (vector-ref obstacles i)))))

(define (Obstacle-MG-step
	 data                       ; Multigrid-data
	 smoothers!                 ; call with ((vector-ref smoothers! k) rhs z obstacle)
	 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
	 obstacle)                    ; other arguments (obstacles, etc.)
  (declare (fixnum))

  (with-access
   data (Obstacle-multigrid-data operators projectors injectors solvers interpolating-projectors)
   (let ((smoother! (vector-ref smoothers! k)))
     
     (let pre-smoothing-loop ((zl z0) (l 0))
       (if (< l m1)
	   (pre-smoothing-loop (smoother! rhs zl obstacle)
			       (+ 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))))
			    (coarse-obstacle (Operator-apply (vector-ref interpolating-projectors k) (Vector-subtract obstacle zl))))
			(let correction-loop ((q (Lattice-max coarse-obstacle (Vector-zero coarse-rhs))) (l 0))
			  (if (< l p)
			      (correction-loop 
			       (Obstacle-MG-step data smoothers! m1 m2 p (- k 1) q coarse-rhs coarse-obstacle)
			       (+ l 1))
			      (Lattice-max obstacle (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 obstacle)
					(+ 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-handy-method (MG-add-solver-information! (multigrid-data Obstacle-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 (let ((smoother! (smoother-maker (vector-ref operators k))))
				    (lambda (rhs zl obstacle)
				      (Lattice-max obstacle (smoother! rhs zl))))))
      (set! solvers (make-vector n))
      
      (do ((k 0 (+ k 1)))
	  ((= k n))
	(vector-set! solvers k
		     (lambda (A b initial-guess obstacle)
		       (do ((i 0 (+ i 1))
			    (z initial-guess (Obstacle-MG-step multigrid-data
							       smoothers!
							       m1
							       m2
							       p
							       k 
							       z
							       b
							       obstacle)))
			   ((= i r) z))))))))