{-# LANGUAGE RankNTypes #-}

-- | Rendering of Accelerate arrays as raster images
--
module Graphics.Gloss.Accelerate.Raster.Array (

  module Graphics.Gloss.Accelerate.Data.Point,
  module Graphics.Gloss.Accelerate.Data.Color,

  -- * Display functions
  Render, Display(..),
  animateArray, animateArrayWith,
  animateArrayIO, animateArrayIOWith,
  playArray, playArrayWith,
  playArrayIO, playArrayIOWith,

  -- * Picture creation
  makePicture,

) where

-- Friends
import Graphics.Gloss.Accelerate.Render
import Graphics.Gloss.Accelerate.Data.Color
import Graphics.Gloss.Accelerate.Data.Point
import Graphics.Gloss.Accelerate.Data.Picture

-- Standard library
import Prelude                                          as P

-- Gloss
import Graphics.Gloss.Data.Display                      ( Display(..) )
import Graphics.Gloss.Data.Picture                      ( Picture(..) )
import Graphics.Gloss.Interface.IO.Animate              as G ( animateFixedIO, black )
import Graphics.Gloss.Interface.Pure.Game               as G ( Event, play )
import Graphics.Gloss.Interface.IO.Game                 as G ( playIO )

-- Accelerate
import Data.Array.Accelerate                            as A


-- Animate --------------------------------------------------------------------
-- -------

-- | Animate a bitmap generated by an Accelerate computation, using the default
--   backend.
--
animateArray
    :: Display                          -- ^ Display mode
    -> (Int, Int)                       -- ^ Number of pixels to draw per point
    -> (Exp Float -> Acc (Array DIM2 Color))
            -- ^ A function to construct an array of colours. The function
            --   should return an array of the same extent every time it is
            --   applied.
            --
            --   It is passed the time in seconds since the program started.
    -> IO ()
animateArray = animateArrayWith defaultRender


-- | Animate a bitmap generated by an Accelerate computation, specifying the
--   backend used to render the image.
--
animateArrayWith
    :: Render                           -- ^ Method to render the array
    -> Display                          -- ^ Display mode
    -> (Int, Int)                       -- ^ Number of pixels to draw per point
    -> (Exp Float -> Acc (Array DIM2 Color))
            -- ^ A function to construct an array of colours. The function
            --   should return an array of the same extent every time it is
            --   applied.
            --
            --   It is passed the time in seconds since the program started.
    -> IO ()
animateArrayWith render display (zoomX, zoomY) makeArray
  | zoomX < 1 || zoomY < 1
  = error "Graphics.Gloss.Raster: invalid pixel scalar factor"

  | otherwise
  = let picture         = makePicture render zoomX zoomY (makeArray . the)
                        . fromList Z
                        . return
    in
    animateFixedIO display G.black (return . picture)


-- | Animate a bitmap generated by an Accelerate computation and IO actions, using the default
--   backend.
--
animateArrayIO
    :: Arrays world
    => Display                          -- ^ Display mode
    -> (Int, Int)                       -- ^ Number of pixels to draw per point
    -> (Float -> IO world)              -- ^ Extract world from time in seconds
                                        --   since the program started
    -> (Acc world -> Acc (Array DIM2 Color))
            -- ^ A function to construct an array of colours. The function
            --   should return an array of the same extent every time it is
            --   applied.
            --
            --   It is passed the world
    -> IO ()
animateArrayIO = animateArrayIOWith defaultRender


-- | Animate a bitmap generated by an Accelerate computation and IO actions, specifying the
--   backend used to render the image.
--
animateArrayIOWith
    :: Arrays world
    => Render                           -- ^ Method to render the array
    -> Display                          -- ^ Display mode
    -> (Int, Int)                       -- ^ Number of pixels to draw per point
    -> (Float -> IO world)              -- ^ Extract world from time in seconds
                                        --   since the program started
    -> (Acc world -> Acc (Array DIM2 Color))
            -- ^ A function to construct an array of colours. The function
            --   should return an array of the same extent every time it is
            --   applied.
            --
            --   It is passed the world
    -> IO ()
animateArrayIOWith render display (zoomX, zoomY) makeWorld makeArray
  | zoomX < 1 || zoomY < 1
  = error "Graphics.Gloss.Raster: invalid pixel scalar factor"

  | otherwise
  = let picture = fmap (makePicture render zoomX zoomY makeArray)
                . makeWorld
    in
    animateFixedIO display G.black picture


-- | Play with a bitmap generated by an Accelerate computation, using the
--   default backend.
--
playArray
    :: Arrays world
    => Display          -- ^ Display mode
    -> (Int, Int)       -- ^ Number of pixels to draw per point
    -> Int              -- ^ Number of simulation steps to take for each second of real time
    -> state            -- ^ The initial state
    -> (state -> world) -- ^ Extract the world state
    -> (Acc world -> Acc (Array DIM2 Color))
            -- ^ Compute the colour of the world
    -> (Event -> state -> state)
            -- ^ Handle input events
    -> (Float -> state -> state)
            -- ^ Step the world one iteration.
            --   It is passed the time in seconds since the program started.
    -> IO ()
playArray = playArrayWith defaultRender


-- | Play with a bitmap generated by an Accelerate computation, specifying the
--   method used to render the world.
--
playArrayWith
    :: Arrays world
    => Render           -- ^ Method to render the world
    -> Display          -- ^ Display mode
    -> (Int, Int)       -- ^ Number of pixels to draw per point
    -> Int              -- ^ Number of simulation steps to take for each second of real time
    -> state            -- ^ The initial state
    -> (state -> world) -- ^ Extract the world state
    -> (Acc world -> Acc (Array DIM2 Color))
            -- ^ Compute the colour of the world
    -> (Event -> state -> state)
            -- ^ Handle input events
    -> (Float -> state -> state)
            -- ^ Step the world one iteration.
            --   It is passed the time in seconds since the program started.
    -> IO ()
playArrayWith render display (zoomX, zoomY) stepRate
              initState makeWorld makeArray handleEvent stepState
  | zoomX < 1 || zoomY < 1
  = error "Graphics.Gloss.Raster: invalid pixel scalar factor"

  | otherwise
  = let picture         = makePicture render zoomX zoomY makeArray
                        . makeWorld
    in
    play display G.black stepRate initState picture handleEvent stepState


-- | Play with a bitmap generated by an Accelerate computation and IO actions, using the
--   default backend.
--
playArrayIO
    :: Arrays world
    => Display          -- ^ Display mode
    -> (Int, Int)       -- ^ Number of pixels to draw per point
    -> Int              -- ^ Number of simulation steps to take for each second of real time
    -> state            -- ^ The initial state
    -> (state -> IO world) -- ^ Extract the world state
    -> (Acc world -> Acc (Array DIM2 Color))
            -- ^ Compute the colour of the world
    -> (Event -> state -> IO state)
            -- ^ Handle input events
    -> (Float -> state -> IO state)
            -- ^ Step the world one iteration.
            --   It is passed the time in seconds since the program started.
    -> IO ()
playArrayIO = playArrayIOWith defaultRender


-- | Play with a bitmap generated by an Accelerate computation and IO actions, specifying the
--   method used to render the world.
--
playArrayIOWith
    :: Arrays world
    => Render           -- ^ Method to render the world
    -> Display          -- ^ Display mode
    -> (Int, Int)       -- ^ Number of pixels to draw per point
    -> Int              -- ^ Number of simulation steps to take for each second of real time
    -> state            -- ^ The initial state
    -> (state -> IO world) -- ^ Extract the world state
    -> (Acc world -> Acc (Array DIM2 Color))
            -- ^ Compute the colour of the world
    -> (Event -> state -> IO state)
            -- ^ Handle input events
    -> (Float -> state -> IO state)
            -- ^ Step the world one iteration.
            --   It is passed the time in seconds since the program started.
    -> IO ()
playArrayIOWith render display (zoomX, zoomY) stepRate
              initState makeWorld makeArray handleEvent stepState
  | zoomX < 1 || zoomY < 1
  = error "Graphics.Gloss.Raster: invalid pixel scalar factor"

  | otherwise
  = let picture         = fmap (makePicture render zoomX zoomY makeArray)
                        . makeWorld
    in
    G.playIO display G.black stepRate initState picture handleEvent stepState


-- Internals
-- ---------

-- | Lift an Accelerate computation from a 'world' to an image into a real
--   Haskell-land function that executes the computation of the image and wraps
--   it as a Gloss picture ready for display.
--
makePicture
    :: Arrays world
    => Render                                   -- ^ method to compute the image
    -> Int                                      -- ^ pixel width
    -> Int                                      -- ^ pixel height
    -> (Acc world -> Acc (Array DIM2 Color))    -- ^ function to create the image
    -> (world -> Picture)                       -- ^ new function that generates the picture
makePicture render zoomX zoomY makeArray
  = let -- compute the image
        pixels          = render (A.map (packRGBA . opaque) . makeArray)

        -- Turn the array into a Gloss picture
        picture world   = bitmapOfArray (pixels world) False
    in
    Scale (P.fromIntegral zoomX) (P.fromIntegral zoomY) . picture