{-# LANGUAGE BangPatterns #-}
module Solver 
	(solveLaplace)
where	
import Data.Array.Repa		as A
import qualified Data.Array.Repa.Shape	as S
import Data.Array.Parallel.Unlifted				(Elt)
import qualified Data.Array.Parallel.Unlifted			as U


-- | Solver for the Laplace equation.
solveLaplace
	:: Int			-- ^ Number of iterations to use.
	-> Array DIM2 Double	-- ^ Boundary value mask.
	-> Array DIM2 Double	-- ^ Boundary values.
	-> Array DIM2 Double	-- ^ Initial state.
	-> Array DIM2 Double

{-# INLINE solveLaplace #-}
solveLaplace steps arrBoundMask@Manifest{} arrBoundValue@Manifest{} arrInit
 = go steps arrInit
 where	go i arr@Manifest{}
	   | i == 0	= arr
	   | otherwise	= go (i - 1) 
			$ force
			$ applyBoundary arrBoundMask arrBoundValue
			$ relaxLaplace arr


-- | Perform matrix relaxation for the Laplace equation,
--	using a stencil function.
--
--   Computation fn is
--	u'(i,j) = (u(i-1,j) + u(i+1,j) + u(i,j-1) + u(i,j+1)) / 4
--
relaxLaplace
	:: Array DIM2 Double
	-> Array DIM2 Double

{-# INLINE relaxLaplace #-}
relaxLaplace !arr
 = traverse arr id elemFn
 where
	_ :. height :. width	
		= extent arr

	{-# INLINE elemFn #-}
	elemFn get d@(sh :. i :. j)
	 = if isBorder i j
		 then  get d
		 else (get (sh :. (i-1) :. j)
		   +   get (sh :. i     :. (j-1))
		   +   get (sh :. (i+1) :. j)
	 	   +   get (sh :. i     :. (j+1))) / 4

	-- Check if this element is on the border of the matrix.
	-- If so we can't apply the stencil because we don't have all the neighbours.
	{-# INLINE isBorder #-}
	isBorder i j
	 	=  (i == 0) || (i >= width  - 1) 
	 	|| (j == 0) || (j >= height - 1) 


-- | Apply the boundary conditions to this matrix.
--	The mask  matrix has 0 in places where boundary conditions hold
--	and 1 otherwise.
--
--	The value matrix has the boundary condition value in places where it holds,
--	and 0 otherwise.
-- 
applyBoundary
	:: Array DIM2 Double	-- ^ Boundary condition mask.
	-> Array DIM2 Double	-- ^ Boundary condition values.
	-> Array DIM2 Double	-- ^ Initial matrix.
	-> Array DIM2 Double	-- ^ Matrix with boundary conditions applied.

{-# INLINE applyBoundary #-}
applyBoundary arrBoundMask arrBoundValue arr
	= A.zipWith (+) arrBoundValue
	$ A.zipWith (*) arrBoundMask  arr