packages feed

repa-examples-1.0.0.0: MMult/src/Main.hs

{-# LANGUAGE PatternGuards #-}

import Solver
import Data.Array.Repa	
import Data.Maybe
import Matrix
import System.Environment
import Control.Monad

-- Arg Parsing ------------------------------------------------------------------------------------
data Arg
	= ArgSolver       String
	| ArgMatrixRandom Int Int
	| ArgMatrixFile   FilePath
	| ArgOutFile	  FilePath
	deriving Show

isArgMatrix arg
 = case arg of
	ArgMatrixRandom{}	-> True
	ArgMatrixFile{}		-> True
	_			-> False

parseArgs []		= []
parseArgs (flag:xx)
	| "-file"	<- flag
	, f:rest	<- xx
	= ArgMatrixFile f : parseArgs rest

	| "-out"	<- flag
	, f:rest	<- xx
	= ArgOutFile f	: parseArgs rest
	
	| "-random"	<- flag
	, x:y:rest	<- xx
	= ArgMatrixRandom (read x) (read y) : parseArgs rest
	
	| otherwise	
	= error $ "bad arg " ++ flag ++ "\n"

printHelp
	= putStr 	
	$ unlines
	[ "Usage: mmult [args..]"
	, ""
	, "  -random <height> <width>   Use a random matrix of this size."
	, "  -file   <filename>         Read a matrix from this file."
	, "  -out    <filename>         Write resulting matrix to this file."
	, ""
	, "  Format of matrix file:"
	, "    MATRIX"
	, "    <width> <height>"
	, "    <whitespace separated values..>"
	, "" ]


-- | Get a matrix from a file, or generate a random one.
getMatrix :: Arg -> IO (Array DIM2 Double)
getMatrix arg
 = case arg of
	ArgMatrixFile   fileName	
	 -> readMatrixFromTextFile fileName

	ArgMatrixRandom height width	
	 -> genRandomMatrix (Z :. height :. width)	

			
-- Main -------------------------------------------------------------------------------------------
main :: IO ()
main 
 = do	args	<- liftM parseArgs $ getArgs
	main' args

main' args
	| [argMat1, argMat2]	<- filter isArgMatrix args
	, mArgOut		<- listToMaybe [s | ArgOutFile s <- args]
	= do	
		-- Get matrices from files, 
		-- or generate random ones we were asked to.
		mat1		<- getMatrix argMat1
		mat2		<- getMatrix argMat2

        	mat1
          	 `deepSeqArray` mat2
          	 `deepSeqArray` return ()
		
		-- Run the solver.
		let matResult	= mmMult mat1 mat2

		matResult `deepSeqArray` return ()

		-- Write the output to file if requested.
		case mArgOut of 
		 Nothing	-> return ()
		 Just fileOut	-> writeMatrixAsTextFile matResult fileOut
					
	| otherwise
	= printHelp