#|
Procedures in this file solve the transport-dominated parabolic problem:
$$
\alignedat2
u_t-\nabla\cdot(a\nabla u) + b\cdot\nabla u + cu&=f,&&\quad x\in\Omega,\ 0<t<T,
\\
\frac{\partial u}{\partial n}&=0,&&\quad x\in\partial\Omega,\ 0<t<T
\\
u(x,0)&=u_0(x),&&\quad x\in\Omega
\endalignedat
$$
where $T=N\Delta t$.  It uses backward differences in time, linear finite elements,
on triangles, and treats the transport term with the modified method of characteristics.
When a characteristic base lies outside a triangle, we find the value of $\Delta t b\cdot u/c$
at the base of the characteristic by reflecting the base point across the boundary edge.
|#

(define (add-displaced-capacity-terms! rhs u characteristic-bases capacity)
  (let ((triangulation (Finite-element-space-triangulation (Finite-element-vector-space u))))
    (let ((displaced-midpoint-values (cu-at-displaced-edge-midpoints capacity u characteristic-bases)))
      (let ((rhs-data (Finite-element-vector-data rhs)))
	(Triangulation-for-each-triangle
	 (lambda (t)
	   (Triangle-interior-integral t displaced-midpoint-values rhs-data))
	 triangulation)))))

(define (cu-at-displaced-edge-midpoints c u characteristic-bases)
  (declare (flonum))
  (with-access
   u (Finite-element-vector data space)
   (let ((data data)
	 (triangulation (Finite-element-space-triangulation space)))
     (let ((midpoint-values (make-f64vector (Triangulation-edges-length triangulation))))
       (Triangulation-for-each-edge
	(lambda (e)
	  (with-access
	   e (Edge vertex-1 vertex-2 index)
	   (cond ((vector-ref characteristic-bases index)   ;; if the base of the characteristic is in a triangle, use it
		  => (lambda (base)
		       (let ((t (car base))
			     (coefficients (cdr base)))
			 (with-access
			  t (Triangle vertex-1 vertex-2 vertex-3)
			  (f64vector-set! midpoint-values
					  index
					  (* (c (Edge-midpoint e))
					     (+ (* (f64vector-ref coefficients 0) (f64vector-ref data (Vertex-index vertex-1)))
						(* (f64vector-ref coefficients 1) (f64vector-ref data (Vertex-index vertex-2)))
						(* (f64vector-ref coefficients 2) (f64vector-ref data (Vertex-index vertex-3))))))))))
		 (else
		  (error "Shouldn't happen"))))) ;; now that we've reflected characteristic bases on the boundary back into the domain.
	triangulation)
       midpoint-values))))

(define (calculate-characteristic-bases triangulation flow)
  (let ((bases (make-vector (Triangulation-edges-length triangulation) #f)))
    (Triangulation-for-each-triangle
     (lambda (t)
       (Triangle-for-each-edge
	(lambda (e)
	  (with-access
	   e (Edge index)
	   (if (not (vector-ref bases index))
	       (let ((midpoint (Edge-midpoint e)))
		 (let ((displaced-midpoint (Point-subtract midpoint (flow midpoint))))
		   (cond ((Point-in-triangle? displaced-midpoint t)
			  => (lambda (coefficients)
			       (vector-set! bases index (cons t coefficients))))))))))
	t))
     triangulation)

    ;; OK, now for some boundary edges, the characteristic base will be outside the domain.
    ;; We deal with this by reflecting the displaced-midpoint in the boundary edge and
    ;; calculating the coordinates of this new point.

    (Triangulation-for-each-boundary-edge
     (lambda (e)
       (with-access
	e (Boundary-edge index triangle vertex-1 vertex-2)
	(if (not (vector-ref bases index))
	    (let* ((midpoint (Edge-midpoint e))
		   (displaced-midpoint (Point-subtract midpoint (flow midpoint)))
		   (point-1 (Vertex-point vertex-1))
		   (point-2 (Vertex-point vertex-2))
		   (v (Point-subtract displaced-midpoint point-1))
		   (e (Point-subtract point-2 point-1))
		   (reflected-displaced-midpoint
		    (Point-add point-1
			       (Point-subtract (Point-scale (FLOAT (* 2.0 (/ (Point-dot-product v e) (Point-dot-product e e)))) e)
					       v))))
		   (cond ((Point-in-triangle? reflected-displaced-midpoint triangle)
			  => (lambda (coefficients)
			       (vector-set! bases index (cons triangle coefficients))))
			(else
			 (error "Shouldn't happen")))))))
     triangulation)
    bases))

(define (Point-in-triangle? p t)
  (declare (flonum))
  (with-access
   t (Triangle vertex-1 vertex-2 vertex-3)
   ;; returns #f or the barycentric coordinates of p in terms
   ;; of point-1, point-2, and point-3
   (let ((point-1 (Vertex-point vertex-1))
	 (point-2 (Vertex-point vertex-2))
	 (point-3 (Vertex-point vertex-3)))
     ;; map point-3 to zero
     (let ((e1 (Point-subtract point-1 point-3))
	   (e2 (Point-subtract point-2 point-3))
	   (p (Point-subtract p point-3)))
       ;; map e1 to (1 0) and e2 to (0 1), calculate coordinates of p
       (let ((a (Point-x e1))
	     (b (Point-x e2))
	     (c (Point-y e1))
	     (d (Point-y e2)))
	 (let ((determinant (- (* a d) (* b c))))
	   (let ((p1 (/ (- (* d (Point-x p)) (* b (Point-y p))) determinant))
		 (p2 (/ (- (* a (Point-y p)) (* c (Point-x p))) determinant)))
	     ;; if (p1 p2) is in the standard triangle, return coordinates of
	     ;; original p.
	     (and (<= 0.0 p1)
		  (<= 0.0 p2)
		  (<= (+ p1 p2) 1.0)
		  (f64vector p1 p2 (- 1.0 p1 p2))))))))))

(define (solve-transport-problem c a b u0 f delta-t space N)
  (let* ((m (->Linear-finite-element-operator space))
	 (f-term (instantiate Linear-finite-element-vector
			      space: space))
	 (u (instantiate Linear-finite-element-vector
			 space: space))
	 (stop-iterations? (lambda (A preconditioner z p r x c m)
			     (or (FIX (= m 200))
				 (FLOAT (< c 1.e-20)))))
	 (t (Finite-element-space-triangulation space))
	 (characteristic-bases 
	  (calculate-characteristic-bases t 
					  (FLOAT 
					   (lambda (p) 
					     (Point-scale (/ delta-t (c p)) (b p)))))))
    (Triangulation-for-each-vertex
     (let ((data (Finite-element-vector-data u)))
       (lambda (v)
	 (with-access
	  v (Vertex point index)
	  (f64vector-set! data index (u0 point)))))
     t)
    (add-interior-integral-terms! f-term (lambda (p) (FLOAT (* delta-t (f p)))))
    (add-stiffness-terms! m (FLOAT (lambda (p) (* delta-t (a p)))))
    (add-mass-terms! m c)
    (let loop ((i 0) (u u) (old-u #f))
      (if (FIX (< i N))
	  (let ((rhs (instantiate Linear-finite-element-vector
				  space: space)))
	    ;; here u_i is available
            (if (zero? (modulo i 10)) (display (list i #\newline)))
	    (add-displaced-capacity-terms! rhs u characteristic-bases c)
	    (loop (FIX (+ i 1))
		  (if old-u
		      (let* ((predicted-u (Vector-subtract (Vector-scale 2.0 u) old-u))
			     (new-rhs (Vector-subtract (Vector-add f-term rhs)
						       (Operator-apply m predicted-u))))
			(Vector-add predicted-u
				    (conjugate-gradient m 
							new-rhs
							stop-iterations?)))
		      (conjugate-gradient m (Vector-add f-term rhs)  stop-iterations?))
		  u))
	  u))))