(declare (standard-bindings) (extended-bindings) (not safe) (block))

(declare (not core))
(include "points.scm")
(declare (core))

(define number-of-refinements 8)

(define mg-data (MG-build-data (instantiate Multigrid-data) coarse-square-triangulation  number-of-refinements))

(define t (vector-ref (Multigrid-data-triangulations mg-data)  number-of-refinements))

(Triangulation-consistent? t)

(define fes (vector-ref (Multigrid-data-spaces mg-data) number-of-refinements))

(define cg-data (CG-build-data t fes))

(define fine-t (refine t))

(define fine-fes (Triangulation->Linear-finite-element-space fine-t))

(define L2-inner-product-operator (->Linear-finite-element-operator fine-fes))

(add-mass-terms! L2-inner-product-operator (lambda (p) 1.0))

(define H1-inner-product-operator (->Linear-finite-element-operator fine-fes))

(add-stiffness-terms! H1-inner-product-operator (lambda (p) 1.0))

(add-mass-terms! H1-inner-product-operator (lambda (p) 1.0))

(define coarse->fine-operator (Operator-adjoint (coarse<-fine fes fine-fes)))

(define symmetric-tests
  (FLOAT
   (list 
    ;; u(x,y) = x^2 - y^2 (is harmonic), Neumann boundary conditions,
    ;; a=c=1, b = 0, f = u, g = 0, 2, or -2 depending on the boundary
    (instantiate Problem-data
		 description: "a=1, c=1, u=x^2-y^2, Neumann"
		 a: (lambda (p) 1.0)
		 c: (lambda (p) 1.0)
		 f: (lambda (p) 
		      (let ((x (Point-x p))
			    (y (Point-y p)))
			(- (* x x) (* y y))))
		 g: (lambda (p)
		      (let ((x (Point-x p))
			    (y (Point-y p)))
			(cond ((= x 0.0) 0.0)
			      ((= x 1.0) 2.0)
			      ((= y 0.0) 0.0)
			      ((= y 1.0) -2.0)
			      (else (error "shouldn't happen")))))
		 u: (lambda (p) 
		      (let ((x (Point-x p))
			    (y (Point-y p)))
			(- (* x x) (* y y)))))
    (let* ((pi (* 4.0 (atan 1.0)))
	   (coeff (+ 1.0 (* pi pi))))
      (instantiate Problem-data
		   description: "a=1, c=1, u=\\cos(\\pi x), Neumann"
		   a: (lambda (p) 1.0)
		   c: (lambda (p) 1.0)
		   f: (lambda (p) 
			(let ((x (Point-x p)))
			  (* coeff (cos (* pi x)))))
		   u: (lambda (p) 
			(let ((x (Point-x p)))
			  (cos (* pi x))))))
    (let* ((pi (* 4.0 (atan 1.0)))
	   (coeff (+ 1.0 (* pi pi)))
	   (minus-pi (- pi)))
      (instantiate Problem-data
		   description: "a=1, c=1, u=\\sin(\\pi x)+\\sin(\\pi y), Neumann"
		   a: (lambda (p) 1.0)
		   c:  (lambda (p) 1.0)
		   f: (lambda (p) 
			(let ((x (Point-x p))
			      (y (Point-y p)))
			  (* coeff (+ (sin (* pi x)) (sin (* pi y))))))
		   g: (lambda (p) minus-pi)
		   u: (lambda (p)  
			(let ((x (Point-x p))
			      (y (Point-y p)))
			  (+ (sin (* pi x)) (sin (* pi y)))))))
    (let ((pi (* 4.0 (atan 1.0))))
      (instantiate Problem-data
		   description: "a=11.0+10.0*\\cos\\pi x, c=1, u=xy, Neumann"
		   a: (lambda (p) 
			(let ((x (Point-x p)))
			  (+ 11.0 (* 10. (cos (* pi x))))))
		   c: (lambda (p) 1.0)
		   f: (lambda (p) 
			(let ((x (Point-x p))
			      (y (Point-y p)))
			  (+ (* x y)
			     (* 10.0 pi (sin (* pi x)) y))))
		   g: (lambda (p)
			(let ((x (Point-x p))
			      (y (Point-y p)))
			  (cond ((= x 0.0) (- (* 21.0 y)))
				((= x 1.0) y)
				((= y 0.0) (* (+ 11.0 (* 10.0 (cos (* pi x))))
					      (- x)))
				((= y 1.0) (* (+ 11.0 (* 10.0 (cos (* pi x))))
					      x))
				(else
				 (error "shouldn't happen")))))
		   u: (lambda (p) 
			(let ((x (Point-x p))
			      (y (Point-y p)))
			  (* x y)))))
    (instantiate Problem-data
		 description: "a=1, c=1+xy, u=x^2-y^2, Neumann"
		 a: (lambda (p) 1.0)
		 c: (lambda (p)  
		      (let ((x (Point-x p))
			    (y (Point-y p)))
			(+ 1.0 (* x y))))
		 f: (lambda (p)  
		      (let ((x (Point-x p))
			    (y (Point-y p)))
			(* (+ 1.0 (* x y))
			   (- (* x x) (* y y)))))
		 g: (lambda (p)
		      (let ((x (Point-x p))
			    (y (Point-y p)))
			(cond ((= x 0.0) 0.0)
			      ((= x 1.0) 2.0)
			      ((= y 0.0) 0.0)
			      ((= y 1.0) -2.0)
			      (else (error "shouldn't happen")))))
		 u: (lambda (p)
		      (let ((x (Point-x p))
			    (y (Point-y p)))
			(- (* x x) (* y y)))))
    (let* ((pi (* 4.0 (atan 1.0))))
      (instantiate Problem-data
		   description: "a=1, c=2+\\cos(\\pi y), u=\\cos(\\pi x), Neumann"
		   a: (lambda (p) 1.0)
		   c: (lambda (p) 
			(let ((x (Point-x p))
			      (y (Point-y p)))
			  (+ 2.0 (cos (* pi y)))))
		   f: (lambda (p) 
			(let ((x (Point-x p))
			      (y (Point-y p)))
			  (* (cos (* pi x))
			     (+ (* pi pi)
				(* (+ 2.0 (cos (* pi y))))))))
		   u: (lambda (p) 
			(let ((x (Point-x p))
			      (y (Point-y p)))
			  (cos (* pi x))))))
    (let* ((pi (* 4.0 (atan 1.0)))
	   (coeff (+ 1.0 (* pi pi)))
	   (minus-pi (- pi)))
      (instantiate Problem-data
		   description: "a=1, c=1, gamma=2+\\cos(\\pi x), u=\\sin(\\pi x)+\\sin(\\pi y), Robin"
		   a: (lambda (p) 1.0)
		   c:  (lambda (p) 1.0)
		   gamma: (lambda (p)  
			    (let ((x (Point-x p))
				  (y (Point-y p)))
			      (+ 2.0 (cos (* pi x)))))
		   f: (lambda (p)  
			(let ((x (Point-x p))
			      (y (Point-y p)))
			  (* coeff (+ (sin (* pi x)) (sin (* pi y))))))
		   g: (lambda (p)
			(let ((x (Point-x p))
			      (y (Point-y p)))
			  (+ minus-pi
			     (* (+ 2.0 (cos (* pi x)))
				(+ (sin (* pi x)) (sin (* pi y)))))))
		   u: (lambda (p)  
			(let ((x (Point-x p))
			      (y (Point-y p)))
			  (+ (sin (* pi x)) (sin (* pi y)))))))
    (let ((pi (* 4.0 (atan 1.0))))
      (instantiate Problem-data
		   description: "a=11.0+10.0*\\cos\\pi x, c=1+xy, gamma=2+\\cos(\\pi x), u=xy, Robin"
		   a: (lambda (p) 
			(let ((x (Point-x p)))
			  (+ 11.0 (* 10. (cos (* pi x))))))
		   c: (lambda (p) 1.0)
		   f: (lambda (p) 
			(let ((x (Point-x p))
			      (y (Point-y p)))
			  (+ (* x y)
			     (* 10.0 pi (sin (* pi x)) y))))
		   gamma: (lambda (p) 
			    (let ((x (Point-x p))
				  (y (Point-y p)))
			      (+ 2.0 (cos (* pi x)))))
		   g: (lambda (p)
			(let ((x (Point-x p))
			      (y (Point-y p)))
			  (+ (* (* x y)
				(+ 2.0 (cos (* pi x))))
			     (cond ((= x 0.0) (- (* 21.0 y)))
				   ((= x 1.0) y)
				   ((= y 0.0) (* (+ 11.0 (* 10.0 (cos (* pi x))))
						 (- x)))
				   ((= y 1.0) (* (+ 11.0 (* 10.0 (cos (* pi x))))
						 x))
				   (else
				    (error "shouldn't happen"))))))
		   u: (lambda (p) 
			(let ((x (Point-x p))
			      (y (Point-y p)))
			  (* x y))))))))

;; we've defined the following nonsymmetric tests, and tried to solve them by setting up
;; the normal equations in MG-add-operator-and-rhs-information! (i.e., calculating m-transpose
;; and multiplying both m and rhs by m-transpose), but it didn't work at all, except for the
;; first problem for some reason (and even then it didn't work well).  So support
;; for nonsymmetric problems in multigrid has been disabled.

(define nonsymmetric-tests
  (FLOAT
   (define flow-vector (make-Point 10.0 5.0))
   (list 
    ;; u(x,y) = x^2 - y^2 (is harmonic), Neumann boundary conditions,
    ;; a=c=1, b = (10 5), f = u, g = 0, 2, or -2 depending on the boundary
    (instantiate Problem-data
		 description: "a=1, b = (10 5), c=1, u=x^2-y^2, Neumann"
		 a: (lambda (p) 1.0)
		 b: (lambda (p) flow-vector)
		 c: (lambda (p) 1.0)
		 f: (lambda (p) 
		      (let ((x (Point-x p))
			    (y (Point-y p)))
			(- (+ (* x x) (* 20.0 x)) (+ (* y y) (* 10.0 y)))))
		 g: (lambda (p)
		      (let ((x (Point-x p))
			    (y (Point-y p)))
			(cond ((= x 0.0) 0.0)
			      ((= x 1.0) 2.0)
			      ((= y 0.0) 0.0)
			      ((= y 1.0) -2.0)
			      (else (error "shouldn't happen")))))
		 u: (lambda (p) 
		      (let ((x (Point-x p))
			    (y (Point-y p)))
			(- (* x x) (* y y)))))
    (let* ((pi (* 4.0 (atan 1.0)))
	   (coeff (+ 1.0 (* pi pi))))
      (instantiate Problem-data
		   description: "a=1, b = (10 5), c=1, u=\\cos(\\pi x), Neumann"
		   a: (lambda (p) 1.0)
		   b: (lambda (p) flow-vector)
		   c: (lambda (p) 1.0)
		   f: (lambda (p) 
			(let ((x (Point-x p)))
			  (- (* coeff (cos (* pi x))) (* 10.0 pi (sin (* pi x))))))
		   u: (lambda (p) 
			(let ((x (Point-x p)))
			  (cos (* pi x))))))
    (let* ((pi (* 4.0 (atan 1.0)))
	   (coeff (+ 1.0 (* pi pi)))
	   (minus-pi (- pi)))
      (instantiate Problem-data
		   description: "a=1, b = (10 5), c=1, u=\\sin(\\pi x)+\\sin(\\pi y), Neumann"
		   a: (lambda (p) 1.0)
		   b: (lambda (p) flow-vector)
		   c:  (lambda (p) 1.0)
		   f: (lambda (p) 
			(let ((x (Point-x p))
			      (y (Point-y p)))
			  (+ (* coeff (+ (sin (* pi x)) (sin (* pi y))))
			     (* 10.0 pi (cos (* pi x)))
			     (* 5.0 pi (cos (* pi y))))))
		   g: (lambda (p) minus-pi)
		   u: (lambda (p)  
			(let ((x (Point-x p))
			      (y (Point-y p)))
			  (+ (sin (* pi x)) (sin (* pi y)))))))
    (let ((pi (* 4.0 (atan 1.0))))
      (instantiate Problem-data
		   description: "a=11.0+10.0*\\cos\\pi x, b = (10 5), c=1, u=xy, Neumann"
		   a: (lambda (p)  
			(let ((x (Point-x p)))
			  (+ 11.0 (* 10. (cos (* pi x))))))
		   b: (lambda (p) flow-vector)
		   c: (lambda (p) 1.0)
		   f: (lambda (p) 
			(let ((x (Point-x p))
			      (y (Point-y p)))
			  (+ (* x y)
			     (* 10.0 pi (sin (* pi x)) y)
			     (* 10.0 y)
			     (* 5.0 x))))
		   g: (lambda (p)
			(let ((x (Point-x p))
			      (y (Point-y p)))
			  (cond ((= x 0.0) (- (* 21.0 y)))
				((= x 1.0) y)
				((= y 0.0) (* (+ 11.0 (* 10.0 (cos (* pi x))))
					      (- x)))
				((= y 1.0) (* (+ 11.0 (* 10.0 (cos (* pi x))))
					      x))
				(else
				 (error "shouldn't happen")))))
		   u: (lambda (p) 
			(let ((x (Point-x p))
			      (y (Point-y p)))
			  (* x y)))))
    (instantiate Problem-data
		 description: "a=1, b = (10 5), c=1+xy, u=x^2-y^2, Neumann"
		 a: (lambda (p) 1.0)
		 b: (lambda (p) flow-vector)
		 c: (lambda (p) 
		      (let ((x (Point-x p))
			    (y (Point-y p)))
			(+ 1.0 (* x y))))
		 f: (lambda (p) 
		      (let ((x (Point-x p))
			    (y (Point-y p)))
			(+ (* (+ 1.0 (* x y))
			      (- (* x x) (* y y)))
			   (- (* 20.0 x)
			      (* 10.0 y)))))
		 g: (lambda (p)
		      (let ((x (Point-x p))
			    (y (Point-y p)))
			(cond ((= x 0.0) 0.0)
			      ((= x 1.0) 2.0)
			      ((= y 0.0) 0.0)
			      ((= y 1.0) -2.0)
			      (else (error "shouldn't happen")))))
		 u: (lambda (p)
		      (let ((x (Point-x p))
			    (y (Point-y p)))
			(- (* x x) (* y y)))))
    (let* ((pi (* 4.0 (atan 1.0))))
      (instantiate Problem-data
		   description: "a=1, b = (10 5), c=2+\\cos(\\pi y), u=\\cos(\\pi x), Neumann"
		   a: (lambda (p) 1.0)
		   b: (lambda (p) flow-vector)
		   c: (lambda (p) 
			(let ((y (Point-y p)))
			  (+ 2.0 (cos (* pi y)))))
		   f: (lambda (p) 
			(let ((x (Point-x p))
			      (y (Point-y p)))
			  (- (* (cos (* pi x))
				(+ (* pi pi)
				   (* (+ 2.0 (cos (* pi y))))))
			     (* 10.0 pi (sin (* pi x))))))
		   u: (lambda (p) 
			(let ((x (Point-x p)))
			  (cos (* pi x))))))
    (let* ((pi (* 4.0 (atan 1.0)))
	   (coeff (+ 1.0 (* pi pi)))
	   (minus-pi (- pi)))
      (instantiate Problem-data
		   description: "a=1, b = (10 5), c=1, gamma=2+\\cos(\\pi x), u=\\sin(\\pi x)+\\sin(\\pi y), Robin"
		   a: (lambda (p) 1.0)
		   b: (lambda (p) flow-vector)
		   c:  (lambda (p) 1.0)
		   gamma: (lambda (p)  
			    (let ((x (Point-x p)))
			      (+ 2.0 (cos (* pi x)))))
		   f: (lambda (p) 
			(let ((x (Point-x p))
			      (y (Point-y p)))
			  (+ (* coeff (+ (sin (* pi x)) (sin (* pi y))))
			     (* 10.0 pi (cos (* pi x)))
			     (* 5.0 pi (cos (* pi y))))))
		   g: (lambda (p)
			(let ((x (Point-x p))
			      (y (Point-y p)))
			  (+ minus-pi
			     (* (+ 2.0 (cos (* pi x)))
				(+ (sin (* pi x)) (sin (* pi y)))))))
		   u: (lambda (p) 
			(let ((x (Point-x p))
			      (y (Point-y p)))
			  (+ (sin (* pi x)) (sin (* pi y)))))))
    (let ((pi (* 4.0 (atan 1.0))))
      (instantiate Problem-data
		   description: "a=11.0+10.0*\\cos\\pi x, b = (10 5), c=1+xy, gamma=2+\\cos(\\pi x), u=xy, Robin"
		   a: (lambda (p) 
			(let ((x (Point-x p)))
			  (+ 11.0 (* 10. (cos (* pi x))))))
		   b: (lambda (p) flow-vector)
		   c: (lambda (p) 1.0)
		   f: (lambda (p)
			(let ((x (Point-x p))
			      (y (Point-y p)))
			  (+ (* x y)
			     (* 10.0 pi (sin (* pi x)) y)
			     (* 10.0 y)
			     (* 5.0 x))))
		   gamma: (lambda (p) 
			    (let ((x (Point-x p)))
			      (+ 2.0 (cos (* pi x)))))
		   g: (lambda (p)
			(let ((x (Point-x p))
			      (y (Point-y p)))
			  (+ (* (* x y)
				(+ 2.0 (cos (* pi x))))
			     (cond ((= x 0.0) (- (* 21.0 y)))
				   ((= x 1.0) y)
				   ((= y 0.0) (* (+ 11.0 (* 10.0 (cos (* pi x))))
						 (- x)))
				   ((= y 1.0) (* (+ 11.0 (* 10.0 (cos (* pi x))))
						 x))
				   (else
				    (error "shouldn't happen"))))))
		   u: (lambda (p) 
			(let ((x (Point-x p))
			      (y (Point-y p)))
			  (* x y))))))))


(define (run-test check solver-thunk)
  
  (define (format-for-display x)
    (if (< x 1.0)
	(let ((x (+ 1.0 x)))
	  (let ((string-x (number->string x)))
	    (substring string-x 1 12)))
	(number->string x)))
  
  (let* ((at-start (##process-statistics))
	 (result (solver-thunk))
	 (at-end (##process-statistics))
	 (user+sys-ns (+ (- (f64vector-ref at-end 0) (f64vector-ref at-start 0))
			 (- (f64vector-ref at-end 1) (f64vector-ref at-start 1))))
	 (gc-ns (+ (- (f64vector-ref at-end 3) (f64vector-ref at-start 3))
		   (- (f64vector-ref at-end 4) (f64vector-ref at-start 4))))
	 (difference (Vector-subtract (Operator-apply coarse->fine-operator result) check)))
    
    (for-each display (list #\tab
			    (format-for-display (sqrt (Vector-inner-product difference L2-inner-product-operator difference)))))
    (for-each display (list #\tab
			    (format-for-display (sqrt (Vector-inner-product difference H1-inner-product-operator difference)))))
    (for-each display (list #\tab
			    (round (* 1000 (- user+sys-ns gc-ns)))))
    (for-each display (list #\tab
			    (round (* 1000 gc-ns))
			    #\newline))))

(define (run-multigrid-tests tests)
  
  (for-each (lambda (problem-data)
	      
	      (define check (instantiate Linear-finite-element-vector
					 data: (let ((u (Problem-data-u problem-data))
						     (result (make-f64vector (Triangulation-vertices-length fine-t))))
						 (vector-for-each-with-index-1 
						  (lambda (i v)
						    (with-access
						     v (Vertex point)
						     (f64vector-set! result i (u point))))
						  (Triangulation-vertices fine-t))
						 result)
					 space: fine-fes))

	      (for-each display (list #\newline (Problem-data-description problem-data) #\newline))
  
	      (MG-add-operator-and-rhs-information! mg-data problem-data)
	      (MG-propagate-level-information! mg-data)

	      (for-each (lambda (smoother-maker smoother-name)
			  (for-each (lambda (p)
				      (for-each (lambda (m1)
						  
						  (define (iota m)
						    (let loop ((result '()) (m m))
						      (if (< m 0)
							  result
							  (loop (cons m result) (- m 1)))))
						  
						  (for-each (lambda (m2)
							      (for-each (lambda (r)
									  (MG-add-solver-information! mg-data smoother-maker smoother-name r m1 m2 p)
									  (display (Multigrid-data-solver-description mg-data))
									  (run-test check (lambda () (MG-solver mg-data))))
									'(1 2 3 4)))
							    (list m1)));(iota m1)))
						'(1 2 3 4)))
				    '(1 2)))
			(list richardson-smoother!
			      jacobi-smoother!
			      gauss-seidel-smoother!)
			(list "  Richardson"
			      "      Jacobi"
			      "Gauss-Seidel")))
	    tests))

(define (run-cg-tests tests)
  
  (for-each (lambda (problem-data)

	      (for-each display (list #\newline (Problem-data-description problem-data) #\newline))
	      
	      (CG-add-operator-and-rhs-information! cg-data problem-data)
	      
	      (for-each (lambda (preconditioner preconditioner-name)
			  
			  (define check (instantiate Linear-finite-element-vector
					 data: (let ((u (Problem-data-u problem-data))
						     (result (make-f64vector (Triangulation-vertices-length fine-t))))
						 (vector-for-each-with-index-1 
						  (lambda (i v)
						    (with-access
						     v (Vertex point)
						     (f64vector-set! result i (u point))))
						  (Triangulation-vertices fine-t))
						 result)
					 space: fine-fes))
			  
			  (CG-add-solver-information! cg-data
						      (preconditioner (Conjugate-gradient-data-operator cg-data))
						      preconditioner-name 
						      2000 
						      (expt (Vector-dimension (Conjugate-gradient-data-right-hand-side cg-data)) -2.0))
			  (display (Conjugate-gradient-data-solver-description cg-data))
			  (run-test check (Conjugate-gradient-data-solver cg-data)))
									    
			(list (lambda (m) identity-operator)
			      diagonal-preconditioner
			      gauss-seidel-preconditioner)
			'("(Identity preconditioner)    "
			  "(Diagonal preconditioner)    "
			  "(Gauss-Seidel preconditioner)")))
	    tests))