Args.hs

module Args where
import Config
import Data.Array.Repa                  as R
import Data.Array.Repa.Algorithms.Pixel as R
import Data.Array.Repa.IO.BMP           as R
import System.Console.GetOpt
import Data.IORef
import Prelude                          as P
import Control.Monad

-- | Command line options.
loadConfig :: [String] -> IO Config
loadConfig args
 = do   
        batchModeArg    <- newIORef False
        maxStepsArg     <- newIORef 0
        widthArg        <- newIORef 100
        scaleArg        <- newIORef 5
        rateArg         <- newIORef 25
        deltaArg        <- newIORef 0.1
        diffArg         <- newIORef 0
        viscArg         <- newIORef 0
        densArg         <- newIORef 100
        velArg          <- newIORef (20, 20)

        densityBmpArg   <- newIORef Nothing
        velocityBmpArg  <- newIORef Nothing
        
        let setWidthArg arg     = writeIORef widthArg         (read arg)
        let setScaleArg arg     = writeIORef scaleArg         (read arg)
        let setDeltaArg arg     = writeIORef deltaArg         (read arg)
        let setDiffArg  arg     = writeIORef diffArg          (read arg)
        let setViscArg  arg     = writeIORef viscArg          (read arg)
        let setDensArg  arg     = writeIORef densArg          (read arg)
        let setVelArg   arg     = let a = read arg in writeIORef velArg (a, a)
        let setRate     arg     = writeIORef rateArg          (read arg)
        let setMaxSteps arg     = writeIORef maxStepsArg      (read arg)
        let setBatchArg         = writeIORef batchModeArg     True
        let setDensityBMP  arg  = writeIORef densityBmpArg    (Just arg)
        let setVelocityBMP arg  = writeIORef velocityBmpArg   (Just arg)

        let options :: [OptDescr (IO ())]
            options
             = [Option [] ["batch"]             (NoArg  setBatchArg            )
                        "Run a fixed number of steps instead of displaying in a window.",

                Option [] ["max"]               (ReqArg setMaxSteps     "INT")
                        "Quit after this number of steps.",

                Option [] ["width"]             (ReqArg setWidthArg     "INT")
                        "Size of simulation (100)",

                Option [] ["scale"]             (ReqArg setScaleArg     "INT")
                        "Width of cell in window (4)",

                Option [] ["rate"]              (ReqArg setRate         "INT")
                        "Frame rate of simulator (25)",

                Option [] ["delta"]             (ReqArg setDeltaArg     "FLOAT")
                        "Length of time steps (0.1)",

                Option [] ["diffusion"]         (ReqArg setDiffArg      "FLOAT")
                        "Diffusion rate for the density (0)",

                Option [] ["viscosity"]         (ReqArg setViscArg      "FLOAT")
                        "Viscosity rate for the velocity (0)",

                Option [] ["user-density"]      (ReqArg setDensArg      "FLOAT")
                        "Magnitude of a user inserted density (100)",

                Option [] ["user-velocity"]     (ReqArg setVelArg       "FLOAT")
                        "Magnitude of a user inserted velocity (20)",

                Option [] ["bmp-density"]       (ReqArg setDensityBMP   "FILE.bmp")
                        "File for initial fluid density",

                Option [] ["bmp-velocity"]      (ReqArg setVelocityBMP  "FILE.bmp")
                        "File for initial fluid velocity"
                ]

        case getOpt RequireOrder options args of
         (actions,      [],   []) 
          -> foldl (>>) (return ()) actions

         (      _, nonOpts,   []) 
          -> error $ "unrecognized arguments: " P.++ unwords nonOpts

         (      _,       _, msgs) 
          -> error $ concat msgs P.++ usageInfo "Usage: fluid [OPTION...]" options


        batchMode       <- readIORef batchModeArg
        maxSteps        <- readIORef maxStepsArg
        width           <- readIORef widthArg
        let height      = width
        scale           <- readIORef scaleArg
        rate            <- readIORef rateArg
        delta           <- readIORef deltaArg
        diff            <- readIORef diffArg
        visc            <- readIORef viscArg
        dens            <- readIORef densArg
        vel             <- readIORef velArg
        densityBmp      <- readIORef densityBmpArg
        velocityBmp     <- readIORef velocityBmpArg


        -- Load the initial desity bmp if we were given one.
        initialDensity
         <- case densityBmp of
                -- No density file given, so just set the field to zero.
                Nothing 
                 -> return
                        $ R.fromListUnboxed (Z :. height :. width)
                        $ replicate (height * width) 0

                -- Load density from a .bmp, using the luminance as
                -- the scalar density value.
                Just filePath
                 -> do  result   <- readImageFromBMP filePath
                        let arr  =  case result of
                                        Right arr'      -> arr'
                                        Left err        -> error $ show err

                        let Z :. height' :. width' 
                                 = extent arr

                        when (height /= height' || width /= width')
                         $ error "Fluid: bmp size does not match --width argument"

                        density  <- computeUnboxedP 
                                  $ R.map floatLuminanceOfRGB8 arr
                        return density


        -- Load the initial velocity bmp if we were given one
        initialVelocity
         <- case velocityBmp of
                -- No velocity file given, so just set the field to zero.
                Nothing
                 -> return 
                        $ R.fromListUnboxed (Z :. height :. width)
                        $ replicate (height * width) (0, 0)

                -- Load 
                Just filePath
                 -> do  result  <- readImageFromBMP filePath
                        let arr  =  case result of
                                        Right arr'      -> arr'
                                        Left err        -> error $ show err

                        let Z :. height' :. width' 
                                 = extent arr

                        when (height /= height' || width /= width')
                         $ error "Fluid: bmp size does not match --width argument"

                        let {-# INLINE conv #-}
                            conv (r, g, _b) 
                             = let r'   = fromIntegral (-128 + fromIntegral r :: Int)
                                   g'   = fromIntegral (-128 + fromIntegral g :: Int)
                               in  (r' * 0.0001, g' * 0.0001)

                        velocity  <- computeUnboxedP $ R.map conv arr
                        return velocity

        return  Config
                { configRate            = rate
                , configWindowSize      = (scale * width, scale * width)
                , configMaxSteps        = maxSteps
                , configBatchMode       = batchMode 
                , configModelSize       = (width, width)
                , configDelta           = delta
                , configDiffusion       = diff
                , configViscosity       = visc
                , configDensity         = dens
                , configVelocity        = vel 
                , configInitialDensity  = initialDensity
                , configInitialVelocity = initialVelocity }