Use cmdargs to parse command-line arguments.

[spline3.git] / src / Main.hs

diff --git a/src/Main.hs b/src/Main.hs
index 1a716836c5db52096e7c15e6caaf45e307ba83cc..951c9c36bf7a819b51f758de42a158484e00d90b 100644 (file)
--- a/src/Main.hs
+++ b/src/Main.hs
@@ -1,15 +1,22 @@
+{-# LANGUAGE RecordWildCards, DoAndIfThenElse #-}
+
 module Main
 where
 
+import Data.Maybe (fromJust)
+import Control.Monad (when)
 import qualified Data.Array.Repa as R
-import System.Environment (getArgs)
+import Data.Maybe (isJust)
+import GHC.Conc (getNumProcessors, setNumCapabilities)
+import System.IO (hPutStrLn, stderr)
+import System.Exit (exitSuccess, exitWith, ExitCode(..))
 
+import CommandLine (Args(..), apply_args)
+import ExitCodes
 import Grid (zoom)
 import MRI (
   flip_x,
   flip_y,
-  mri_shape,
-  mri_slice3d,
   read_word16s,
   round_array,
   swap_bytes,
@@ -18,43 +25,91 @@ import MRI (
   z_slice
   )
 
-in_file :: FilePath
-in_file = "./data/mri.bin"
+
+validate_args :: Args -> IO ()
+validate_args Args{..} = do
+  when (scale <= 0) $ do
+    hPutStrLn stderr "ERROR: scale must be greater than zero."
+    exitWith (ExitFailure exit_arg_not_positive)
+
+  when (width <= 0) $ do
+    hPutStrLn stderr "ERROR: width must be greater than zero."
+    exitWith (ExitFailure exit_arg_not_positive)
+
+  when (height <= 0) $ do
+    hPutStrLn stderr "ERROR: height must be greater than zero."
+    exitWith (ExitFailure exit_arg_not_positive)
+
+  when (depth <= 0) $ do
+    hPutStrLn stderr "ERROR: depth must be greater than zero."
+    exitWith (ExitFailure exit_arg_not_positive)
+
+  case slice of
+    Just s ->
+      when (s < 0 || s > depth) $ do
+        hPutStrLn stderr "ERROR: slice must be between zero and depth."
+        exitWith (ExitFailure exit_arg_out_of_bounds)
+    Nothing -> return ()
 
 
 main :: IO ()
-main = main3d
+main = do
+  args@Args{..} <- apply_args
+  -- validate_args will simply exit if there's a problem.
+  validate_args args
+
+  -- The first thing we do is set the number of processors. We get the
+  -- number of processors (cores) in the machine with
+  -- getNumProcessors, and set it with setNumCapabilities. This is so
+  -- we don't have to pass +RTS -Nfoo on the command line every time.
+  num_procs <- getNumProcessors
+  setNumCapabilities num_procs
 
-main3d :: IO ()
-main3d = do
-  (s:_) <- getArgs
-  let scale = read s :: Int
+  -- Determine whether we're doing 2d or 3d. If we're given a slice,
+  -- assume 2d.
+  let mri_shape = (R.Z R.:. depth R.:. height R.:. width) :: R.DIM3
+
+  if (isJust slice) then
+    main2d args mri_shape
+  else
+    main3d args mri_shape
+
+  exitSuccess
+
+  where
+
+
+
+main3d :: Args -> R.DIM3 -> IO ()
+main3d Args{..} mri_shape = do
   let zoom_factor = (scale, scale, scale)
-  let out_file = "output.bin"
-  arr <- read_word16s in_file
+  arr <- read_word16s input mri_shape
   let arr'          = swap_bytes arr
   let arrMRI        = R.reshape mri_shape arr'
   dbl_data <- R.computeUnboxedP $ R.map fromIntegral arrMRI
-  output <- zoom dbl_data zoom_factor
-  word16_output <- R.computeUnboxedP $ round_array output
-  write_word16s out_file word16_output
-  return ()
+  raw_output <- zoom dbl_data zoom_factor
+  word16_output <- R.computeUnboxedP $ round_array raw_output
+  write_word16s output word16_output
 
 
-main2d :: IO ()
-main2d = do
-  (s:_) <- getArgs
-  let scale = read s :: Int
+main2d :: Args -> R.DIM3 -> IO ()
+main2d Args{..} mri_shape = do
   let zoom_factor = (1, scale, scale)
-  let out_file = "output.bmp"
-  arr <- read_word16s in_file
-  arrSlice <- R.computeUnboxedP $ z_slice 50 $ flip_x $ flip_y $ swap_bytes arr
+  arr <- read_word16s input mri_shape
+  arrSlice <- R.computeUnboxedP
+               $ z_slice (fromJust slice)
+               $ flip_x width
+               $ flip_y height
+               $ swap_bytes arr
   let arrSlice' = R.reshape mri_slice3d arrSlice
 
   -- If zoom isn't being inlined we need to extract the slice before hand,
   -- and convert it to the require formed.
   dbl_data      <- R.computeUnboxedP $ R.map fromIntegral arrSlice'
-  output        <- zoom dbl_data zoom_factor
-  arrSlice0     <- R.computeUnboxedP $ z_slice 0 output
-  
-  write_values_slice_to_bitmap arrSlice0 out_file
+  raw_output    <- zoom dbl_data zoom_factor
+  arrSlice0     <- R.computeUnboxedP $ z_slice 0 raw_output
+
+  write_values_slice_to_bitmap arrSlice0 output
+  where
+    mri_slice3d :: R.DIM3
+    mri_slice3d = (R.Z R.:. 1 R.:. height  R.:. width)