--{-# LANGUAGE Strict #-}
{-# LANGUAGE DeriveFoldable, StandaloneDeriving #-}
module Picture.Render
    where
import Control.Lens
import Control.Monad
--import Control.Monad.Trans.State
--
--
import Linear.Matrix
import Linear.V4

import qualified Control.Foldl as F

import Data.Bifunctor

import Picture.Data
import Geometry

import Picture.Preload
--import Control.Lens

import Foreign hiding (rotate)

import Codec.Picture

import Graphics.Rendering.OpenGL hiding (Line,translate,scale,imageHeight,imageWidth,Polygon,Color,T)
import qualified Graphics.Rendering.OpenGL as GL

import Data.Foldable
import Data.List
import qualified Data.Vector.Storable as V
import qualified Data.IntMap as IM

import Control.DeepSeq

import qualified SDL as SDL

white = (1,1,1,1)
black = (0,0,0,1)

polyToTris :: [s] -> [s]
{-# INLINE polyToTris #-}
polyToTris (a:b:c:as) = a : intercalate [a] (zipWith (\x y->[x,y]) (init (b:c:as)) (c:as))
polyToTris _ = []

tripFirst :: (a -> a') -> (a,b,c) -> (a',b,c)
{-# INLINE tripFirst #-}
tripFirst f (x,y,z) = (f x,y,z)

tripSecond :: (b -> b') -> (a,b,c) -> (a,b',c)
{-# INLINE tripSecond #-}
tripSecond f (x,y,z) = (x,f y,z)

scaleT :: Float -> (Point3,Point4,Point2) -> (Point3,Point4,Point2)
{-# INLINE scaleT #-}
scaleT x (a,b,(o,s)) = (a,b,(o,s*x))

overPos :: (Point3 -> Point3) -> RenderType -> RenderType
overPos f (RenderPoly vs) = RenderPoly $ map (first $ f) vs
overPos f (RenderLine vs) = RenderLine $ map (first $ f) vs
overPos f (RenderText vs) = RenderText $ map (\(a,b,c) -> (f a,b,c)) vs
overPos f (RenderCirc (a,b,c)) = RenderCirc (f a,b,c)
overPos f (RenderArc (a,b,c))  = RenderArc (f a,b,c)
overPos _ RenderBlank = RenderBlank

overCol :: (Point4 -> Point4) -> RenderType -> RenderType
overCol f (RenderPoly vs) = RenderPoly $ map (second $ f) vs
overCol f (RenderLine vs) = RenderLine $ map (second $ f) vs
overCol f (RenderText vs) = RenderText $ map (\(a,b,c) -> (a,f b,c)) vs
overCol f (RenderCirc (a,b,c)) = RenderCirc (a,f b,c)
overCol f (RenderArc (a,b,c)) = RenderArc (a,f b,c)
overCol _ RenderBlank = RenderBlank

scaleRen,translateRen :: Float -> Float -> RenderType -> RenderType
{-# INLINE scaleRen #-}
scaleRen x y (RenderText vs) = overPos (scale3 x y) $ RenderText $ map (scaleT x) vs
scaleRen x y rt = overPos (scale3 x y) rt
{-# INLINE translateRen #-}
translateRen x y = overPos $ translate3 x y
rotateRen,setDepthRen :: Float -> RenderType -> RenderType
{-# INLINE rotateRen #-}
rotateRen a (RenderArc (p,c,(as,ae,r,w))) = overPos (rotate3 a) $ RenderArc (p,c,(f as,f ae,r,w))
    --where f b = normalizeAngle $ a + b
    where f b = a + b
rotateRen a pic = overPos (rotate3 a) pic
{-# INLINE setDepthRen #-}
setDepthRen d = overPos $ \(x,y,_) -> (x,y,-d)
{-# INLINE colorRen #-}
colorRen :: RGBA -> RenderType -> RenderType
colorRen c = overCol $ const c

stringToList :: String -> [(Point3,Point4,Point2)]
{-# INLINE stringToList #-}
stringToList s = zipWith (\x (a,b,c) -> (translate3 x 0 a,b,c))
                         [0,0.9*dimText..]
                         $ map charToTuple s
    where dimText = 100

charToTuple :: Char -> (Point3,Point4,Point2)
{-# INLINE charToTuple #-}
charToTuple c = ((0,0,0),white,(offset,100))
    where offset = fromIntegral (fromEnum c) - 32

picToList :: Int -> Picture -> [RenderType]
--{-# INLINE picToList #-}
picToList x (Polygon i ps) 
    | i == x = [RenderPoly $ zip (map zeroZ $ polyToTris ps) $ repeat black]
    | otherwise = []
picToList x (PolygonCol i vs) 
    | i /= x = []
    | otherwise = 
        let (ps,cs) = unzip vs
        in  [RenderPoly $ zip (map zeroZ $ polyToTris ps) $ polyToTris cs]
picToList x (Circle i r) 
    | i == x = [RenderCirc $ ((0,0,0),black,r)]
    | otherwise = []
picToList x (ThickArc i startA endA rad wdth) 
    | i == x = [RenderArc $ ((0,0,0),black,(startA,endA,rad,wdth))]
    | otherwise = []
picToList x (Line i ps) 
    | i == x = [RenderLine $ zip (map zeroZ $ doubleLine ps) $ repeat white]
    | otherwise = []
picToList x (Text i s) 
    | i == x = [RenderText $ stringToList s]
    | otherwise = []
picToList j Blank = []
picToList j (Scale x y pic) = fmap (scaleRen x y) $ picToList j pic
picToList j (Translate x y pic) = fmap (translateRen x y) $ picToList j pic
picToList j (Rotate a pic) = fmap (rotateRen a) $ picToList j pic
picToList j (SetDepth a pic) = fmap (setDepthRen a) $ picToList j pic
picToList j (Color c pic) = fmap (colorRen c) $ picToList j pic
picToList j (Pictures pics) = concatMap (picToList j) pics

picToFTree :: Int -> Picture -> FTree RenderType
--{-# INLINE picToFTree #-}
picToFTree x (Polygon i ps) 
    | i == x = FLeaf $ RenderPoly $ zip (map zeroZ $ polyToTris ps) $ repeat black
    | otherwise = FLeaf RenderBlank
picToFTree x (PolygonCol i vs) 
    | i /= x = FLeaf RenderBlank
    | otherwise = 
        let (ps,cs) = unzip vs
        in  FLeaf $ RenderPoly $ zip (map zeroZ $ polyToTris ps) $ polyToTris cs
picToFTree x (Circle i r) 
    | i == x = FLeaf $ RenderCirc $ ((0,0,0),black,r)
    | otherwise = FLeaf RenderBlank
picToFTree x (ThickArc i startA endA rad wdth) 
    | i == x = FLeaf $ RenderArc $ ((0,0,0),black,(startA,endA,rad,wdth))
    | otherwise = FLeaf RenderBlank
picToFTree x (Line i ps) 
    | i == x = FLeaf $ RenderLine $ zip (map zeroZ $ doubleLine ps) $ repeat white
    | otherwise = FLeaf RenderBlank
picToFTree x (Text i s) 
    | i == x = FLeaf $ RenderText $ stringToList s
    | otherwise = FLeaf RenderBlank
picToFTree j Blank = FLeaf RenderBlank
picToFTree j (Scale x y pic) = collapseBranch (scaleRen x y) $ picToFTree j pic
picToFTree j (Translate x y pic) = collapseBranch (translateRen x y) $ picToFTree j pic
picToFTree j (Rotate a pic) = collapseBranch (rotateRen a) $ picToFTree j pic
picToFTree j (SetDepth a pic) = collapseBranch (setDepthRen a) $ picToFTree j pic
picToFTree j (Color c pic) = collapseBranch (colorRen c) $ picToFTree j pic
picToFTree j (Pictures pics) = FBranches $ map (picToFTree j) pics

collapseBranch :: (RenderType -> RenderType) -> FTree RenderType -> FTree RenderType
collapseBranch f (FBranch g t) = FBranch (f . g) t
collapseBranch f (FBranches ts) = FBranches $ map (collapseBranch f) ts
collapseBranch f t = FBranch f t

doubleLine :: [Point2] -> [Point2]
doubleLine (x:y:xs) = concat $ zipWith (:) (init (x:y:xs)) $ map (\a -> [a]) (y:xs)
doubleLine _ = []

theFold :: TwoPtrs
        -> ThreePtrs
        -> ThreePtrs
        -> TwoPtrs
        -> ThreePtrs
        -> F.FoldM IO RenderType (Int,Int,Int,Int,Int)
theFold pas pbs pcs pds pes 
   -- = (,,,,) <$> pokeFold pas <*> pokeTextFold pbs <*> pokeCircFold pcs
    = (,,,,) <$> pokeTwoPtrsWith pokePoly pas 
             <*> pokeThreePtrsWith pokeText pbs 
             <*> pokeThreePtrsWith pokeCirc pcs
             <*> pokeTwoPtrsWith pokeLine pds
             <*> pokeThreePtrsWith pokeArc pes

type ThreePtrs = (Ptr Float,Ptr Float,Ptr Float)
type TwoPtrs = (Ptr Float,Ptr Float)

pokeThreePtrsWith :: (ThreePtrs -> Int -> RenderType -> IO Int)
    -> ThreePtrs -> F.FoldM IO RenderType Int
pokeThreePtrsWith pokeF ptrs = F.FoldM (pokeF ptrs) (return 0) return
pokeTwoPtrsWith :: (TwoPtrs -> Int -> RenderType -> IO Int)
    -> TwoPtrs -> F.FoldM IO RenderType Int
pokeTwoPtrsWith  pokeF ptrs = F.FoldM (pokeF ptrs) (return 0) return

pokeArc:: ThreePtrs -> Int -> RenderType -> IO Int
pokeArc (pa,pb,pc) n (RenderArc (p,c,s))
    | n > 20000 * 2 = return n
    | otherwise = do
        pokeThreeOff pa n p
        pokeFourOff  pb n c
        pokeFourOff  pc n s
        return $ n + 1
pokeArc _ n _ = return n

pokeTwoOff :: Ptr Float -> Int -> (Float,Float) -> IO ()
pokeTwoOff ptr n (x,y) = do
    pokeElemOff ptr (2*n+0) x
    pokeElemOff ptr (2*n+1) y
pokeThreeOff :: Ptr Float -> Int -> (Float,Float,Float) -> IO ()
pokeThreeOff ptr n (x,y,z) = do
    pokeElemOff ptr (3*n+0) x
    pokeElemOff ptr (3*n+1) y
    pokeElemOff ptr (3*n+2) z
pokeFourOff :: Ptr Float -> Int -> (Float,Float,Float,Float) -> IO ()
pokeFourOff ptr n (x,y,z,w) = do
    pokeElemOff ptr (4*n+0) x
    pokeElemOff ptr (4*n+1) y
    pokeElemOff ptr (4*n+2) z
    pokeElemOff ptr (4*n+3) w


pokeLine :: TwoPtrs -> Int -> RenderType -> IO Int
pokeLine (pa,pb) n (RenderLine vs) = foldM (pokeLineVert pa pb) n vs
pokeLine _       n _               = return n

pokeLineVert :: Ptr Float -> Ptr Float -> Int -> (Point3, Point4) -> IO Int
pokeLineVert pa pb n (p,c)
    | n > 20000 * 2 = return n
    | otherwise = do
        pokeThreeOff pa n p
        pokeFourOff  pb n c
        return (n+1)

pokeCirc :: ThreePtrs -> Int -> RenderType -> IO Int
pokeCirc (pa,pb,pc) n (RenderCirc (p,c,s)) 
    | n > 20000 * 2 = return n
    | otherwise = do
        pokeThreeOff pa n p
        pokeFourOff  pb n c
        pokeElemOff  pc n s
        return (n+1)
pokeCirc _          n _               = return n

pokeText :: (Ptr Float, Ptr Float, Ptr Float) -> Int -> RenderType -> IO Int
pokeText (pa,pb,pc) n  (RenderText vs) = foldM (pokeTextVert pa pb pc) n vs
pokeText _          n _               = return n

pokeTextVert :: Ptr Float -> Ptr Float -> Ptr Float -> Int -> (Point3, Point4, Point2) -> IO Int
pokeTextVert pa pb pc n (p,c,t)
    | n > 20000 * 2 = return n
    | otherwise = do
        pokeThreeOff pa n p
        pokeFourOff  pb n c
        pokeTwoOff   pc n t
        return (n+1)

pokePoly :: TwoPtrs -> Int -> RenderType -> IO Int
pokePoly (pa,pb) n (RenderPoly vs) = foldM (pokeVert pa pb) n vs
pokePoly _       n _               = return n

pokeVert :: Ptr Float -> Ptr Float -> Int -> (Point3, Point4) -> IO Int
pokeVert pa pb n (p,c)
    | n > 20000 * 2 = return n
    | otherwise = do
        pokeThreeOff pa n p
        pokeFourOff  pb n c
        return (n+1)

translate3 :: Float -> Float -> Point3 -> Point3
{-# INLINE translate3 #-}
translate3 a b (x,y,z) = (x+a,y+b,z)
scale3 :: Float -> Float -> Point3 -> Point3
{-# INLINE scale3 #-}
scale3 a b (x,y,z) = (x*a,y*b,z)
rotate3 :: Float -> Point3 -> Point3
{-# INLINE rotate3 #-}
rotate3 a (x,y,z) = (x',y',z)
    where (x',y') = rotateV a (x,y)

fSize = sizeOf (0 :: Float)

bindArrayBuffers :: Int -> [(BufferObject,Ptr Float,Int)] -> IO ()
bindArrayBuffers numVs ps = do
    forM_ ps $ \(bo,ptr,i) -> do
        bindBuffer ArrayBuffer $= Just bo
        bufferData ArrayBuffer $= (fromIntegral $ fSize * numVs * i, ptr, StreamDraw)
        
twoPtrsVAO :: VAO -> (Ptr Float, Ptr Float)
twoPtrsVAO vao = case (\(_,ps,_) -> ps) $ unzip3 $ _vaoBufferTargets vao of
    (a:b:_) -> (a,b)
threePtrsVAO :: VAO -> (Ptr Float, Ptr Float,Ptr Float)
threePtrsVAO vao = case (\(_,ps,_) -> ps) $ unzip3 $ _vaoBufferTargets vao of
    (a:b:c:_) -> (a,b,c)


renderPicture' :: RenderData -> Float -> Float -> (Float,Float) -> (Float,Float) ->
                    [(Point2,Point2)] -> [Point4] -> Picture -> IO (Word32,Word32)
renderPicture' pdata rot zoom (tranx,trany) (winx,winy) wallPoints lightPoints pic = do
    ticksAfterL <- SDL.ticks
    startWallTicks <- SDL.ticks
    depthFunc $= Just Lequal

    aticks <- SDL.ticks
-- calculate world transformation matrix
    let scalMat = Linear.Matrix.transpose $
            V4 (V4 (2*zoom/winx) 0             0 (0::GLfloat))
                (V4 0             (2*zoom/winy) 0 0)
                (V4 0             0             1 0)
                (V4 0             0             0 1)
    let rotMat = Linear.Matrix.transpose $
            V4 (V4 (cos rot) (sin (-rot)) 0 0)
                (V4 (sin rot) (cos rot)    0 0)
                (V4 0             0        1 0)
                (V4 0             0        0 1)
    let tranMat = Linear.Matrix.transpose $
            V4 (V4 1        0        0 0)
                (V4 0        1        0 0)
                (V4 0        0        1 0)
                (V4 (-tranx) (-trany) 0 1)
    let wmat = scalMat !*! rotMat !*! tranMat
        vToL (V4 a b c d) = [a,b,c,d]
    wmata <- (newMatrix RowMajor $ concatMap vToL $ vToL wmat) :: IO (GLmatrix GLfloat)

-- set common uniforms
    forM_ [_basicShader      pdata
          ,_textShader       pdata
          ,_circShader       pdata
          ,_arcShader        pdata
          ,_fadeCircleShader pdata
          ,_backShader       pdata
          ,_wallShadowShader pdata
          ] $ \shad -> do
        currentProgram $= Just (fst shad)
        uniform (snd shad !! 0) $= Vector2 winx winy
        uniform (snd shad !! 1) $= zoom
        uniform (snd shad !! 2) $= rot
        uniform (snd shad !! 3) $= Vector2 tranx trany
        uniform (snd shad !! 4) $= wmata
    bticks <- SDL.ticks

-- draw lightmap
    bindVertexArrayObject $= Just (_vao $ _wallVAO pdata)
    let wallPtr = (\(_,x,_) -> x) $ head $ _vaoBufferTargets $ _wallVAO pdata
--        wallPtr2 = (\(_,x,_) -> x) $ (_vaoBufferTargets $ _wallVAO pdata) !! 1
        foldWalls n ((x,y),(z,w)) = do
                pokeFourOff wallPtr  n (x,y,z,w)
--                pokeFourOff wallPtr2 n (a,b,c,d)
                return $ n+1
    nWalls <- foldM foldWalls 0 wallPoints 

    forM_ lightPoints $ \(x,y,r,lum) -> do
        cullFace $= Just Front
        clear [DepthBuffer]
        currentProgram $= Just (fst $ _wallShadowShader pdata)
        bindArrayBuffers (length wallPoints) $ _vaoBufferTargets $ _wallVAO pdata
        bindVertexArrayObject $= Just (_vao $ _wallVAO pdata)
        uniform (_wssLightPos pdata)  $= Vector2 (x) (y)
        blendFunc $= (Zero,One)
        drawArrays Points (fromIntegral 0) (fromIntegral $ length wallPoints)
        cullFace $= Nothing
        currentProgram $= Just (fst $ _fadeCircleShader pdata)
        bindVertexArrayObject $= Just (_vao $ _fadeCircVAO pdata)
        let fadeCircPtr = (\(_,ptr,_) -> ptr) $ head $ _vaoBufferTargets $ _fadeCircVAO pdata
        pokeFourOff fadeCircPtr 0 (x,y,r,lum)
        bindArrayBuffers (1) $ _vaoBufferTargets $ _fadeCircVAO pdata
        blendFuncSeparate    $= ((Zero,Zero),(Zero, OneMinusSrcAlpha))
        drawArrays Points (fromIntegral 0) (fromIntegral 1)
--    ticksAfterL <- SDL.ticks
    endWallTicks <- SDL.ticks
-- draw picture
-- set drawing for on top
    blendFuncSeparate    $= ((SrcAlphaSaturate, OneMinusSrcAlpha), (Zero,One))
    clear [DepthBuffer]
-- draw layer 0
    ticks2 <- renderTree pdata rot zoom (tranx,trany) (winx,winy) (picToFTree 0 pic)
-- reset blend so that light map doesn't apply
    blendFunc $= (SrcAlpha,OneMinusSrcAlpha)
    ticks3 <- renderTree pdata rot zoom (tranx,trany) (winx,winy) $ picToFTree 1 pic
-- set drawing for on top
    blendFuncSeparate    $= ((SrcAlphaSaturate, OneMinusSrcAlpha), (Zero,One))
    ticks4 <- renderTree pdata rot zoom (tranx,trany) (winx,winy) $ picToFTree 2 pic
-- reset uniforms (hacky for now)
    idmat <- (newMatrix RowMajor [1,0,0,0
                                 ,0,1,0,0
                                 ,0,0,1,0
                                 ,0,0,0,1
                                 ]
             ) :: IO (GLmatrix GLfloat)
    forM_ [_basicShader      pdata
          ,_textShader       pdata
          ,_circShader       pdata
          ,_arcShader        pdata
--          ,_fadeCircleShader pdata
--          ,_backShader       pdata
--          ,_wallShadowShader pdata
          ] $ \shad -> do
        currentProgram $= Just (fst shad)
        uniform (snd shad !! 0)  $= Vector2 (2::Float) 2
        uniform (snd shad !! 1)  $= (1::Float)
        uniform (snd shad !! 2)  $= (0::Float)
        uniform (snd shad !! 3)  $= Vector2 (0::Float) 0
        uniform (snd shad !! 4)  $= idmat
    --return (ticksAfterL, ticks2+ticks3+ticks4, endWallTicks - startWallTicks)
    return (ticksAfterL, endWallTicks - startWallTicks)

bufferOffset :: Integral a => a -> Ptr b
bufferOffset = plusPtr nullPtr . fromIntegral


-- the following code draws a picture tree
-- it does not set nor change the blend function or depth buffer
-- nor does it set uniforms
renderTree :: Foldable f => RenderData -> Float -> Float -> (Float,Float) -> (Float,Float)
                    -> f RenderType -> IO Word32
renderTree pdata rot zoom (tranx,trany) (winx,winy) tree = do

    pokeStartTicks <- SDL.ticks
-- poke necessary data
    (nTriVs,nTextVs,numCircVs,nLineVs,nArcVs) 
--        <- F.foldM (theFold (_ptrPosVBO pdata, _ptrColVBO pdata)  
        <- F.foldM (theFold (twoPtrsVAO $ _triVAO pdata)  
                            (threePtrsVAO $ _textVAO pdata)
                            (threePtrsVAO $ _circVAO pdata)
                            (twoPtrsVAO $ _lineVAO pdata)
                            (threePtrsVAO $ _arcVAO pdata)
                   ) $ tree
    pokeEndTicks <-SDL.ticks

    depthFunc $= Just Less

    currentProgram $= Just (fst $ _backShader pdata)
    bindVertexArrayObject $= Just (_vao $ _backVAO pdata)
    let backPtr = (\(_,x,_) -> x) $ head $ _vaoBufferTargets $ _backVAO pdata
        backPtr2 = (\(_,x,_) -> x) $ (_vaoBufferTargets $ _backVAO pdata) !! 1
    pokeFourOff backPtr 0 (tranx,trany,rot,zoom)
    pokeTwoOff backPtr2 0  (winx,winy)
    bindArrayBuffers 1 $ _vaoBufferTargets $ _backVAO pdata
    textureBinding Texture2D $= Just (_textures pdata !! 1)
    drawArrays Points (fromIntegral 0) (fromIntegral 1)

    depthFunc $= Just Lequal
-- draw triangles
    currentProgram $= Just (fst $ _basicShader pdata)
    bindVertexArrayObject $= Just (_vao $ _triVAO pdata)
    bindArrayBuffers nTriVs $ _vaoBufferTargets $ _triVAO pdata
    drawArrays Triangles 0 (fromIntegral $ nTriVs)
-- draw circles
    currentProgram $= Just (fst $ _circShader pdata)
    bindVertexArrayObject $= Just (_vao $ _circVAO pdata)
    bindArrayBuffers numCircVs $ _vaoBufferTargets $ _circVAO pdata
    drawArrays Points 0 (fromIntegral $ numCircVs)
-- draw arcs
-- assumes that the uniforms are set
    currentProgram $= Just (fst $ _arcShader pdata)
    bindVertexArrayObject $= Just (_vao $ _arcVAO pdata)
    bindArrayBuffers nArcVs $ _vaoBufferTargets $ _arcVAO pdata
    drawArrays Points 0 (fromIntegral $ nArcVs)
-- draw lines
    currentProgram $= Just (fst $ _basicShader pdata)
    bindVertexArrayObject $= Just (_vao $ _lineVAO pdata)
    bindArrayBuffers nLineVs $ _vaoBufferTargets $ _lineVAO pdata
    drawArrays Lines 0 (fromIntegral $ nLineVs)
-- draw text
    currentProgram $= Just (fst $ _textShader pdata)
    bindVertexArrayObject $= Just (_vao $ _textVAO pdata)
    bindArrayBuffers nTextVs $ _vaoBufferTargets $ _textVAO pdata
    textureBinding Texture2D $= Just (_textures pdata !! 0)
    drawArrays Points 0 (fromIntegral $ nTextVs)
    return (pokeEndTicks - pokeStartTicks)