-- author: Samuel Gélineau (gelisam)
-- in response to https://www.reddit.com/r/haskell/comments/3u5s4e/is_there_a_way_to_write_the_frames_of_a_gloss/
-- slightly improved

module Graphics.Gloss.Export.Image
    ( Size
    , Animation
    , initialize
    , pictureToImageRGB8
    , pictureToImageRGBA8
    , exportPictureToFormat
    , exportPicturesToFormat
    ) where

import Codec.Picture.Types (Image(..), PixelRGBA8, PixelRGB8, PixelBaseComponent)
import Control.Monad (forM_)
import Data.Vector.Storable (Vector, unsafeFromForeignPtr0, Storable, slice, concat)
import qualified Graphics.Gloss.Rendering as Gloss
import Graphics.GL -- as GL*
import qualified Graphics.UI.GLFW as GLFW
import Foreign (newForeignPtr_, Ptr)
import Foreign.Marshal.Alloc (free)
import Foreign.Marshal.Array (mallocArray)
import Text.Printf (printf)
import GHC.Int
import qualified Graphics.UI.GLUT as GLUT
import Prelude hiding (concat)

type Size = (Int, Int)
type Animation = Float -> Gloss.Picture

-- | Save a gloss Picture to a file.
exportPictureToFormat :: (FilePath -> Image PixelRGBA8 -> IO ()) -- ^ function that saves an intermediate representation to a format. Written with writeXY from Codec.Picture in mind
                      -> Size                -- ^ (width, heigth) in pixels - as in Gloss.Display
                      -> Gloss.Color         -- ^ Background color
                      -> FilePath -> Gloss.Picture -> IO ()
exportPictureToFormat savefunc size bgc f p = do
    s <- initialize size
    --initialize is not part of pictureToImage so in
    --exportPicturesToFormat we only have to initialize once
    (img,ptr) <- pictureToImageRGBA8 size bgc s p
    savefunc f img
    free ptr

-- | If you want to write your own functions, call this function before pictureToImage* 
initialize :: Size -> IO Gloss.State
initialize size = do
    _ <- GLUT.exit                     -- otherwise 'illegal reinitialization'
    (_,_) <- GLUT.getArgsAndInitialize -- needed for text  https://github.com/elisehuard/game-in-haskell/pull/3
    s <- Gloss.initState
    initOpenGL size
    return s


-- let GLFW bother with the OpenGL initialization
initOpenGL :: (Int, Int) -- ^ windowWidth, windowHeight
           -> IO ()        
initOpenGL (windowWidth, windowHeight) = do
    True <- GLFW.init
    GLFW.windowHint (GLFW.WindowHint'Visible False)
    Just w <- GLFW.createWindow
                windowWidth windowHeight
                "gloss-to-file demo"
                Nothing Nothing
    GLFW.makeContextCurrent (Just w)




-- | Save a series of gloss Picture to files of spcified format.
exportPicturesToFormat :: (FilePath -> Image PixelRGBA8 -> IO ()) -- ^ function that saves an intermediate representation to a format. Written with writeXY from Codec.Picture in mind
                       -> Size                -- ^ (width, height) in pixels - as in Gloss.Display
                       -> Gloss.Color         -- ^ background color
                       -> FilePath            -- ^ must contain "%d", will be replaced by frame number
                       -> Animation           -- ^ function that maps from point in time to Picture. analog to Gloss.Animation
                       -> [Float]             -- ^ list of points in time at which to evaluate the animation
                       -> IO ()
exportPicturesToFormat savefunc size bgc f anim ts = do
    s <- initialize size
    forM_ (zip [1..] ts) $ \(n, t) -> do
      let filename = printf f (n :: Int)
      let picture = anim t
      (img,ptr) <- pictureToImageRGBA8 size bgc s picture
      savefunc filename img
      free ptr

-- I couldn't find a way to generalize over the type of pixel so I made RGBA8 and RGB8.
-- Keep in mind: different pixel types have different sizes!
--               GL_RGBA, 4 would also have to be abstracted

-- | convert a gloss 'Picture' into an 'Image'.
-- The pointer should be freed after the image is no longer needed in memory
-- Use 'free ptr' to free the memory.
pictureToImageRGBA8 :: Size                -- ^ (width, height) in pixels - as in Gloss.Display
                    -> Gloss.Color         -- ^ Background color
                    -> Gloss.State         -- ^ Result of 'initialize'
                    -> Gloss.Picture
                    -> IO (Image PixelRGBA8, Ptr (PixelBaseComponent PixelRGBA8)) -- ^ image and a pointer to memory
pictureToImageRGBA8 (windowWidth, windowHeight) bgc s p = do
    drawReadBuffer (windowWidth, windowHeight) bgc s p
    imageData <- mallocArray (windowWidth * windowHeight * 4)
    let wW = fromIntegral windowWidth  :: GHC.Int.Int32
    let wH = fromIntegral windowHeight :: GHC.Int.Int32
    glReadPixels 0 0 wW wH GL_RGBA GL_UNSIGNED_BYTE imageData 
    foreignPtr <- newForeignPtr_ imageData
    let vector = unsafeFromForeignPtr0 foreignPtr (windowWidth * windowHeight * 4)
    let vectorFlipped = reverseImage 4 windowWidth windowHeight vector
    let image :: Image PixelRGBA8
        image = Image windowWidth windowHeight vectorFlipped
    
    return (image, imageData)

drawReadBuffer :: Size
          -> Gloss.Color -- ^ Background color
          -> Gloss.State -> Gloss.Picture -> IO ()
drawReadBuffer size bg s p = do
    glDrawBuffer GL_BACK
    Gloss.withClearBuffer bg $ Gloss.withModelview size $ do
                                                           glColor3f 0 0 0
                                                           Gloss.renderPicture s 1 p
    glReadBuffer GL_BACK 


-- the drawn image is flipped ([rowN,...,row1]) so we need to reverse the order of rows
-- I guess this is because the origin is specified as topleft and bottomleft by different functions
reverseImage :: Storable a => Int -- ^ #Bits per pixel. 4 for rgba8, 3 for rgb8
                           -> Int
                           -> Int -> Vector a -> Vector a
reverseImage components width height vec = concat [slice i widthC vec | i <- map (*widthC) [(height-1),(height-2)..0] ]
    where
      widthC = width*components



-- Below is the version without transparency for making gifs. 
-- ideally there would just be one function and the result would be casted to either ::PixelRGBA8 or ::PixelRGB8

-- | convert a gloss 'Picture' into an 'Image'.
-- The pointer should be freed after the image is no longer needed in memory
-- Use 'free ptr' to free the memory.
pictureToImageRGB8 :: Size                -- ^ (width, height) in pixels - as in Gloss.Display
                   -> Gloss.Color         -- ^ Background color
                   -> Gloss.State         -- ^ Result of 'initializate'
                   -> Gloss.Picture
                   -> IO (Image PixelRGB8, Ptr (PixelBaseComponent PixelRGB8)) -- ^ image and a pointer to memory
pictureToImageRGB8 (windowWidth, windowHeight) bgc s p = do
    drawReadBuffer (windowWidth, windowHeight) bgc s p
    imageData <- mallocArray (windowWidth * windowHeight * 3)
    let wW = fromIntegral windowWidth  :: GHC.Int.Int32
    let wH = fromIntegral windowHeight :: GHC.Int.Int32
    glReadPixels 0 0 wW wH GL_RGB GL_UNSIGNED_BYTE imageData 
    foreignPtr <- newForeignPtr_ imageData
    let vector = unsafeFromForeignPtr0 foreignPtr (windowWidth * windowHeight * 3)
    let vectorFlipped = reverseImage 3 windowWidth windowHeight vector
    let image :: Image PixelRGB8
        image = Image windowWidth windowHeight vectorFlipped
    
    return (image, imageData)