{-# LANGUAGE BangPatterns, RankNTypes, FlexibleContexts #-}
-- | Solver for the Laplace equation
--	You supply a BMP files specifying the boundary conditions.
--	The output is written back to another BMP file.
--
import Data.Array.Repa.Algorithms.Pixel
import Data.Array.Repa.Algorithms.ColorRamp
import Data.Array.Repa.IO.BMP	
import Data.Array.Repa.IO.Timing
import System.Environment
import Data.Word
import Control.Monad
import SolverGet                as SG
import SolverStencil            as SS
import Data.Array.Repa          as R
import Prelude                  as P

type Solver m 
	=  Monad m
        => Int			-- ^ Number of iterations to use.
	-> Array U DIM2 Double	-- ^ Boundary value mask.
	-> Array U DIM2 Double	-- ^ Boundary values.
	-> Array U DIM2 Double	-- ^ Initial state.
	-> m (Array U DIM2 Double)

solvers 
 = 	[ ("get", 	SG.solveLaplace)
	, ("stencil", 	SS.solveLaplace) ]


main :: IO ()
main 
 = do	args	<- getArgs
	case args of
	  [strSolver, steps, fileInput, fileOutput]	
	    -> 	let Just solver	= lookup strSolver solvers
	  	in  laplace solver (read steps) fileInput fileOutput

	  _ -> do
		putStr usage
		return ()


-- | Command line usage information.
usage	:: String
usage	= unlines
	[ "Usage: laplace <solver> <iterations> <input.bmp> <output.bmp>"
	, ""
	, "  iterations  :: Int       Number of iterations to use in the solver."
	, "  input.bmp   :: FileName  Uncompressed RGB24 or RGBA32 BMP file for initial and boundary values."
	, "  output.bmp  :: FileName  BMP file to write output to."
	, "" 
	, "  solver = one of " P.++ show (P.map fst solvers)
	, ""
	, "  Format of input file:"
	, "      Boundary values are indicated in greyscale,"
	, "        ie from the list [(x, x, x) | x <- [0 .. 255]]"
	, "      Non-boundary values are indicated in blue,"
	, "        ie (0, 0, 255)"
	, "      Any other pixel value is an error." 
	, ""
	, "  NOTE: For GHC 7.0.3, this runs better when you turn off the parallel"
	, "        garbage collector. Run with +RTS -qg"
	, "" ]
			

-- | Solve it.
laplace :: Solver IO
	-> Int			-- ^ Number of iterations to use.
	-> FilePath 		-- ^ Input file.
	-> FilePath		-- ^ Output file
	-> IO ()

laplace solve steps fileInput fileOutput
 = do
	-- Load up the file containing boundary conditions.
	arrImage	<- liftM (either (error . show) id)
			$  readImageFromBMP fileInput

	arrBoundValue   <- computeP $ R.map slurpBoundValue arrImage
	arrBoundMask	<- computeP $ R.map slurpBoundMask  arrImage
	let arrInitial	= arrBoundValue		
	
	-- Run the Laplace solver and print how long it took.
	(arrFinal, t)   <- time $ solve steps arrBoundMask arrBoundValue arrInitial

	putStr (prettyTime t)

	-- Write out the result to a file.
	arrImageOut     <- computeP
	                $  R.map rgb8OfDouble
	                $  R.map (rampColorHotToCold 0.0 1.0) arrFinal

	writeImageToBMP	fileOutput arrImageOut


-- | Extract the boundary value from a RGB triple.
slurpBoundValue :: (Word8, Word8, Word8) -> Double
{-# INLINE slurpBoundValue #-}
slurpBoundValue (!r, !g, !b)
	-- A non-boundary value.
 	| r == 0 && g == 0 && b == 255	
	= 0

	-- A boundary value.
	| (r == g) && (r == b) 
	= fromIntegral (fromIntegral r :: Int) / 255
	
	| otherwise
	= error $ "Unhandled pixel value in input " P.++ show (r, g, b)


-- | Extract boundary mask from a RGB triple.
slurpBoundMask :: (Word8, Word8, Word8) -> Double
{-# INLINE slurpBoundMask #-}
slurpBoundMask (!r, !g, !b)
	-- A non-boundary value.
 	| r == 0 && g == 0 && b == 255	
	= 1

	-- A boundary value.
	| (r == g) && (r == b) 
	= 0
	
	| otherwise
	= error $ "Unhandled pixel value in input " P.++ show (r, g, b)