content: {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE PatternGuards    #-}
  {-# LANGUAGE TemplateHaskell  #-}
  --
  -- Configuration parameters
  --
  
  module Config where
  
  import Type
  
  import Data.Label
  import Control.Monad
  import Prelude                                          as P
  
  import Data.Array.Accelerate                            as A
  import Data.Array.Accelerate.IO.Codec.BMP               as A
  import Data.Array.Accelerate.Data.Colour.RGBA           as A
  import Data.Array.Accelerate.Examples.Internal          as A
  
  
  data Initial a
      = FromFile          FilePath
      | FromFunction      (Backend -> Int -> Int -> a)
  
  
  data Config = Config
    {
      -- simulation
      _viscosity          :: !Float
    , _diffusion          :: !Float
    , _timestep           :: !Float
    , _inputDensity       :: !Float
    , _inputVelocity      :: !Float
    , _simulationSteps    :: !Int
    , _simulationWidth    :: !Int
    , _simulationHeight   :: !Int
  
    -- visualisation
    , _displayScale       :: !Int
    , _displayFramerate   :: !Int
    , _initialDensity     :: DensityField
    , _initialVelocity    :: VelocityField
  
    -- extra options to specify initial conditions for command parsing
    , _setupDensity       :: Initial DensityField
    , _setupVelocity      :: Initial VelocityField
    }
  
  $(mkLabels [''Config])
  
  
  defaults :: Config
  defaults = Config
    { _viscosity          = 0
    , _diffusion          = 0
    , _timestep           = 0.1
    , _inputDensity       = 50
    , _inputVelocity      = 20
    , _simulationSteps    = 40
    , _simulationWidth    = 100
    , _simulationHeight   = 100
  
    , _displayScale       = 5
    , _displayFramerate   = 25
    , _initialDensity     = error "initial density??"
    , _initialVelocity    = error "initial velocity??"
  
    , _setupDensity       = FromFunction makeField_empty
    , _setupVelocity      = FromFunction makeField_empty
    }
  
  
  -- | The set of available command-line options
  --
  options :: [OptDescr (Config -> Config)]
  options =
    -- Simulation options
    [ Option  [] ["viscosity"]
              (ReqArg (parse viscosity) "FLOAT")
              (describe viscosity "viscosity for velocity damping")
  
    , Option  [] ["diffusion"]
              (ReqArg (parse diffusion) "FLOAT")
              (describe diffusion "diffusion rate for mass dispersion")
  
    , Option  [] ["delta"]
              (ReqArg (parse timestep) "FLOAT")
              (describe timestep "simulation time between each frame")
  
    , Option  [] ["density"]
              (ReqArg (parse inputDensity) "FLOAT")
              (describe inputDensity "magnitude of user input density")
  
    , Option  [] ["velocity"]
              (ReqArg (parse inputVelocity) "FLOAT")
              (describe inputVelocity "magnitude of user input velocity")
  
    , Option  [] ["iterations"]
              (ReqArg (parse simulationSteps) "INT")
              (describe simulationSteps "number of iterations of the linear solver")
  
    , Option  [] ["width"]
              (ReqArg (parse simulationWidth) "INT")
              (describe simulationWidth "grid width of simulation")
  
    , Option  [] ["height"]
              (ReqArg (parse simulationHeight) "INT")
              (describe simulationHeight "grid height of simulation")
  
    -- Display options
    , Option  [] ["scale"]
              (ReqArg (parse displayScale) "INT")
              (describe displayScale "feature size of visualisation")
  
    , Option  [] ["framerate"]
              (ReqArg (parse displayFramerate) "INT")
              (describe displayFramerate "frame rate for visualisation")
  
    -- Initial conditions
    , Option  [] ["bmp-density"]
              (ReqArg (set setupDensity . FromFile) "FILE.bmp")
              "file for initial fluid density"
  
    , Option  [] ["bmp-velocity"]
              (ReqArg (set setupVelocity . FromFile) "FILE.bmp")
              "file for initial fluid velocity"
  
    , Option  [] ["init-checks"]
              (NoArg init_checks)
              "initial density field with zero velocity field"
  
    , Option  [] ["init-man"]
              (NoArg init_man)
              "initial density field with swirling velocity"
  
    , Option  [] ["init-elk"]
              (NoArg init_elk)
              "initial density field with swirling velocity"
    ]
    where
      parse f x           = set f (read x)
      describe f msg      = msg P.++ " (" P.++ show (get f defaults) P.++ ")"
  
      init_checks         = set setupDensity  (FromFunction makeDensity_checks)
                          . set setupVelocity (FromFunction makeField_empty)
  
      init_man            = set setupDensity  (FromFunction makeDensity_checks)
                          . set setupVelocity (FromFunction makeVelocity_man)
  
      init_elk            = set setupDensity  (FromFunction makeDensity_checks)
                          . set setupVelocity (FromFunction makeVelocity_elk)
  
  
  header :: [String]
  header =
    [ "accelerate-fluid (c) [2011..2013] The Accelerate Team"
    , ""
    , "Usage: accelerate-fluid [OPTIONS]"
    , ""
    ]
  
  footer :: [String]
  footer =
    [ ""
    , "Runtime usage:"
    , "     ESC          quit"
    , "     click        add density sources to the image"
    , "     shift-click  add velocity sources"
    , "     r            reset the image"
    , "     d            toggle display of density field"
    , "     v            toggle display of velocity field lines"
    , ""
    ]
  
  
  -- Initial conditions
  -- ------------------
  
  initialiseConfig :: (Config, Options, [String]) -> IO (Config, Options, [String])
  initialiseConfig (conf, opts, rest) = do
    let backend   = get optBackend opts
        width     = get simulationWidth conf
        height    = get simulationHeight conf
  
    dens  <- case get setupDensity conf of
                FromFile fn       -> loadDensity_bmp backend fn width height
                FromFunction f    -> return (f backend width height)
  
    velo  <- case get setupVelocity conf of
                FromFile fn       -> loadVelocity_bmp backend fn width height
                FromFunction f    -> return (f backend width height)
  
    let conf'     = set initialDensity  dens
                  . set initialVelocity velo
                  $ conf
  
    return (conf', opts, rest)
  
  
  makeField_empty :: FieldElt e => Backend -> Int -> Int -> Field e
  makeField_empty backend width height
    = run backend
    $ A.fill (constant (Z:.height:.width)) zero
  
  
  makeDensity_checks :: Backend -> Int -> Int -> DensityField
  makeDensity_checks backend width height
    = let width'  = constant $ P.fromIntegral width
          height' = constant $ P.fromIntegral height
          yc      = constant $ P.fromIntegral (height `div` 2)
          xc      = constant $ P.fromIntegral (width  `div` 2)
  
          checks ix
            = let Z :. y :. x     = unlift ix
                  x'              = A.fromIntegral x
                  y'              = A.fromIntegral y
                  tx              = 10 * (x' - xc) / width'
                  ty              = 10 * (y' - yc) / height'
                  xk1             = abs tx A.> 3*pi/2 ? (0 , cos tx)
                  yk1             = abs ty A.> 3*pi/2 ? (0 , cos ty)
                  d1              = xk1 * yk1
              in
              0 `A.max` d1
      in
      run backend $ A.generate (constant (Z:.height:.width)) checks
  
  
  makeVelocity_man :: Backend -> Int -> Int -> VelocityField
  makeVelocity_man backend width height
    = let width'  = constant $ P.fromIntegral width
          height' = constant $ P.fromIntegral height
          yc      = constant $ P.fromIntegral (height `div` 2)
          xc      = constant $ P.fromIntegral (width  `div` 2)
  
          man ix
            = let Z :. y :. x     = unlift ix
                  x'              = A.fromIntegral x
                  y'              = A.fromIntegral y
                  xk2             = cos (19 * (x' - xc) / width')
                  yk2             = cos (17 * (y' - yc) / height')
                  d2              = xk2 * yk2 / 5
              in
              lift (constant 0, d2)
      in
      run backend $ A.generate (constant (Z:.height:.width)) man
  
  
  makeVelocity_elk :: Backend -> Int -> Int -> VelocityField
  makeVelocity_elk backend width height
    = let width'  = constant $ P.fromIntegral width
          height' = constant $ P.fromIntegral height
          yc      = constant $ P.fromIntegral (height `div` 2)
          xc      = constant $ P.fromIntegral (width  `div` 2)
  
          elk ix
            = let Z :. y :. x     = unlift ix
                  x'              = A.fromIntegral x
                  y'              = A.fromIntegral y
                  tx              = 12 * (x' - xc) / width'
                  ty              = 12 * (y' - yc) / height'
                  xk2             =  cos tx
                  yk2             = -cos ty
                  d2              = xk2 * yk2 / 5
              in
              lift (constant 0, d2)
      in
      run backend $ A.generate (constant (Z:.height:.width)) elk
  
  
  loadDensity_bmp :: Backend -> FilePath -> Int -> Int -> IO DensityField
  loadDensity_bmp backend filepath width height
    = do  arr             <- either (error . show) id `fmap` readImageFromBMP filepath
          let Z:.h:.w     =  arrayShape arr
  
          when (w P./= width P.|| h P./= height)
            $ error "accelerate-fluid: density-bmp does not match width x height"
  
          return . run backend $ A.map (luminance . unpackRGBA) (use arr)
  
  
  loadVelocity_bmp :: Backend -> FilePath -> Int -> Int -> IO VelocityField
  loadVelocity_bmp backend filepath width height
    = do  arr             <- either (error . show) id `fmap` readImageFromBMP filepath
          let Z:.h:.w     =  arrayShape arr
  
          when (w P./= width P.|| h P./= height)
            $ error "accelerate-fluid: velocity-bmp does not match width x height"
  
          let conv rgb =
                let RGBA r g _ _ = unlift (unpackRGBA8 rgb) :: RGBA (Exp Word8)
                    r'           = A.fromIntegral (-128 + A.fromIntegral r :: Exp Int)
                    g'           = A.fromIntegral (-128 + A.fromIntegral g :: Exp Int)
                in lift (r' * 0.0001, g' * 0.0001)
  
          return . run backend $ A.map conv (use arr)
  
  