{-# LANGUAGE ScopedTypeVariables, BangPatterns #-}
module Model
        ( Delta
        , Rate
        , Field
        , DensityField
        , VelocityField
        , SourceDensity (..)
        , CurrentButton (..)

        , Model (..)
        , initModel

        , pictureOfModel
        , pixel32OfDensity
        , pixel8OfDensity

        , outputBMP
        , outputPPM)
where
import Data.Array.Repa                  as R
import Data.Array.Repa.IO.BMP           as R
import Data.Array.Repa.Repr.ForeignPtr  as R
import Data.Bits
import Data.Word
import Unsafe.Coerce
import Prelude                          as P
import qualified Graphics.Gloss         as G


-- | Time delta (seconds)
type Delta      = Float

-- | Time rate (or period) (1/seconds)
type Rate       = Float


-- | A 2d field.
type Field a        
        = Array U DIM2 a

-- | Scalar density field.
type DensityField   
        = Field Float

-- | Vector velocity field.
type VelocityField  
        = Field (Float, Float)


-- | Button being pressed by the user.
data CurrentButton  
        = LeftButton
        | RightButton
        | None

-- | A density source added by the user.
data SourceDensity a
        = SourceDensity DIM2 a


-- Model ----------------------------------------------------------------------
-- | The world model.
data Model
        = Model
        { densityField   :: DensityField
        , densitySource  :: Maybe (SourceDensity Float)
        , velocityField  :: VelocityField
        , velocitySource :: Maybe (SourceDensity (Float, Float))
        , clickLoc       :: Maybe (Int, Int)
        , stepsPassed    :: Int
        , currButton     :: CurrentButton
        , drawVelocity   :: Bool
        }


-- | Creates an initial blank model
initModel 
        :: Array U DIM2 Float
        -> Array U DIM2 (Float, Float)
        -> Model

initModel density velocity
 | extent density /= extent velocity
 = error "Fluid: initModel density and velocity extents do not  match"

 | otherwise
 =      Model
        { densityField   = density
        , densitySource  = Nothing
        , velocityField  = velocity
        , velocitySource = Nothing
        , clickLoc       = Nothing
        , stepsPassed    = 0
        , currButton     = None
        , drawVelocity   = False }
{-# INLINE initModel #-}


-- Picture --------------------------------------------------------------------
-- | Function to convert the Model into a Bitmap for displaying in Gloss
pictureOfModel :: Monad m => (Float, Float) -> Model -> m G.Picture
pictureOfModel (scaleX, scaleY) m 
 = let  (Z :. height' :. width') = R.extent $ densityField m
        width           = fromIntegral width'
        height          = fromIntegral height'

   in do
        (arrDensity :: Array F DIM2 Word32)
         <- computeP $ R.map pixel32OfDensity $ densityField m

        let picVel :: G.Picture
            picVel  
             | drawVelocity m
             = G.Translate (- width / 2) (- height / 2)
                    $ G.Color (G.light $ G.light G.red)
                    $ G.Pictures
                        [ G.Line [(xf, yf), (xf + vx', yf + vy')]
                        | x <- [0, 5 .. width'  - 1]
                        , y <- [0, 5 .. height' - 1]
                        , let xf         = fromIntegral x
                        , let yf         = fromIntegral y 
                        , let (vx0, vy0) = velocityField m R.! (Z :. y :. x) 
                        , let vx'        = vx0 * 100
                        , let vy'        = vy0 * 100 ]

             | otherwise
             = G.blank

        let picDens :: G.Picture
            picDens = G.bitmapOfForeignPtr width' height'
                        (R.toForeignPtr $ unsafeCoerce arrDensity)
                        False

        return  $ G.Scale scaleX scaleY
                $ G.pictures    [picDens, picVel ]

{-# NOINLINE pictureOfModel #-}


-- | Converts Float value to Word32 for pixel data
pixel32OfDensity :: Float -> Word32
pixel32OfDensity f
 = let  !fsat
          | f <  0       = 0
          | f >= 1       = 1
          | otherwise    = f

        !x      = truncate $ fsat * 255
        !a      = 255

    in   unsafeShiftL x 24 .|. unsafeShiftL x 16 
     .|. unsafeShiftL x  8 .|. a
{-# INLINE pixel32OfDensity #-}


-- | Converts Float value to a tuple of pixel components.
pixel8OfDensity :: Float -> (Word8, Word8, Word8)
pixel8OfDensity f
 = let  !fsat
          | f <  0       = 0
          | f >= 1       = 1
          | otherwise    = f

        !x      = truncate $ fsat * 255
    in  (x, x, x)
{-# INLINE pixel8OfDensity #-}


-- Dump -----------------------------------------------------------------------
-- Writes bitmap data to test batch-mode ran correctly
outputBMP :: FilePath -> Int -> DensityField -> IO ()
outputBMP path step df 
 = do   arr             <- computeUnboxedP $ R.map pixel8OfDensity df
        let step'       = replicate (6 - length (show step)) '0' P.++ show step
        R.writeImageToBMP (path P.++ step' P.++ ".bmp") arr


outputPPM :: Int -> String -> Float -> Array U DIM2 Float -> IO ()
outputPPM step name scale df
 = do   let (Z :. h :. w) = extent df
        let step'         
             = replicate (4 - length (show step)) '0' P.++ show step

        let getVal x y
             =  let v     = df R.! (Z :. y :. x)
                in  truncate (v * 255 * scale)  :: Int

        let showVal x y
             =  let  v   = getVal x y
                in replicate (4 - length (show v)) ' ' P.++ show v

        let mx
             =  maximum [getVal x y | x <- [0..w-1], y <- [0..h-1]]

        let out = unlines $ 
                [ "P2"
                , show h P.++ " " P.++ show w
                , show mx]
                P.++ [ concat [ showVal x y P.++ " "
                              | x <- [0..w - 1]]
                     | y <- [h - 1, h - 2 .. 0]]

        writeFile ("out/" P.++ step' P.++ "-" P.++ name P.++ ".ppm") out