content: {-# OPTIONS -cpp #-}
  {-# LANGUAGE OverloadedStrings, PackageImports, TypeOperators, DataKinds, FlexibleContexts #-}
  
  import Control.Applicative hiding (Const)
  import Control.Monad
  import qualified Data.ByteString.Char8 as SB
  import qualified Data.Trie as T
  import Data.Vect hiding (reflect')
  import Data.Vect.Float.Instances ()
  import FRP.Elerea.Param
  import qualified Graphics.UI.GLFW as GLFW
  import Text.Printf
  
  import LambdaCube.GL
  import LambdaCube.GL.Mesh
  
  import Common.Utils
  import Common.GraphicsUtils
  
  #ifdef CAPTURE
  import Graphics.Rendering.OpenGL.Raw.Core32
  import Codec.Image.DevIL
  import Text.Printf
  import Foreign
  
  withFrameBuffer :: Int -> Int -> Int -> Int -> (Ptr Word8 -> IO ()) -> IO ()
  withFrameBuffer x y w h fn = allocaBytes (w*h*4) $ \p -> do
      glReadPixels (fromIntegral x) (fromIntegral y) (fromIntegral w) (fromIntegral h) gl_RGBA gl_UNSIGNED_BYTE $ castPtr p
      fn p
  #endif
  
  main :: IO ()
  main = do
  #ifdef CAPTURE
      ilInit
  #endif
      
      let pipeline :: Exp Obj (Image 1 V4F)
          pipeline = PrjFrameBuffer "outFB" tix0 sceneRender
  
      (win,windowSize) <- initWindow "LambdaCube 3D Cube Map Demo" 1280 720
      let keyIsPressed k = fmap (==GLFW.KeyState'Pressed) $ GLFW.getKey win k
  
      (duration, renderer) <- measureDuration $ compileRenderer (ScreenOut pipeline)
      putStrLn $ "Renderer compiled - " ++ show duration
      
      putStrLn "Renderer uniform slots:"
      forM_ (T.toList (slotUniform renderer)) $ \(name, slot) -> do
          putStrLn $ "  " ++ SB.unpack name
          forM_ (T.toList slot) $ \(inputName, inputType) -> do
              putStrLn $ "    " ++ SB.unpack inputName ++ " :: " ++ show inputType
      
      putStrLn "Renderer stream slots:"
      forM_ (T.toList (slotStream renderer)) $ \(name, (primitive, attributes)) -> do
          putStrLn $ "  " ++ SB.unpack name ++ " - " ++ show primitive
          forM_ (T.toList attributes) $ \(attributeName, attributeType) -> do
              putStrLn $ "    " ++ SB.unpack attributeName ++ " :: " ++ show attributeType
  
      quadMesh <- compileMesh quad
      addMesh renderer "postSlot" quadMesh []
  
      cubeMesh <- compileMesh (cube 1)
      
      (duration, cubeObjects) <- measureDuration $ replicateM 6 $ addMesh renderer "geometrySlot" cubeMesh ["modelMatrix"]
      putStrLn $ "Cube meshes added - " ++ show duration
      
      reflectorMesh <- compileMesh (capsule 2.5 3.5 25) -- (sphere 5 25)
      (duration, reflectorObject) <- measureDuration $ addMesh renderer "reflectSlot" reflectorMesh ["modelMatrix"]
  
      putStrLn $ "Reflector mesh added - " ++ show duration
      
      let objectSlots = reflectorSlot : map objectUniformSetter cubeObjects
          reflectorSlot = objectUniformSetter reflectorObject
          sceneSlots = uniformSetter renderer
  
          draw command = do
              render renderer
              command
              GLFW.swapBuffers win >> GLFW.pollEvents
  
      sceneSignal <- start $ do
          thread <- scene win keyIsPressed (setScreenSize renderer) sceneSlots objectSlots windowSize
          return $ draw <$> thread
      driveNetwork sceneSignal $ readInput keyIsPressed
  
      dispose renderer
      putStrLn "Renderer destroyed."
  
      GLFW.destroyWindow win
      GLFW.terminate
  
  scene win keyIsPressed setSize sceneSlots (reflectorSlot:planeSlot:cubeSlots) windowSize = do
      pause <- toggle =<< risingEdge =<< effectful (keyIsPressed (GLFW.Key'P))
      time <- transfer 0 (\dt paused time -> time + if paused then 0 else dt) pause 
      
      capture <- toggle =<< risingEdge =<< effectful (keyIsPressed (GLFW.Key'C))
      frameCount <- stateful (0 :: Int) (const (+1))
      
      fpsTracking <- stateful (0, 0, Nothing) $ \dt (time, count, _) -> 
          let time' = time + dt
              done = time > 5
          in if done
             then (0, 0, Just (count / time'))
             else (time', count + 1, Nothing)
  
      mousePosition <- effectful $ do
          (x, y) <- GLFW.getCursorPos win
          return $ Vec2 (realToFrac x) (realToFrac y)
      directionControl <- effectful $ (,,,,)
                   <$> keyIsPressed GLFW.Key'Left
                   <*> keyIsPressed GLFW.Key'Up
                   <*> keyIsPressed GLFW.Key'Down
                   <*> keyIsPressed GLFW.Key'Right
                   <*> keyIsPressed GLFW.Key'RightShift
      
      mousePosition' <- delay zero mousePosition
      camera <- userCamera (Vec3 (-4) 0 10) (mousePosition - mousePosition') directionControl
      
      let setViewCameraMatrix = uniformM44F "viewCameraMatrix" sceneSlots . fromMat4
          setViewCameraPosition = uniformV3F "viewCameraPosition" sceneSlots . fromVec3
          setCubeCameraMatrix i = uniformM44F (cubeMatrixName i) sceneSlots . fromMat4
          setCubeCameraPosition = uniformV3F "cubeCameraPosition" sceneSlots . fromVec3
          setLightPosition = uniformV3F "lightPosition" sceneSlots . fromVec3
          setPlaneModelMatrix = uniformM44F "modelMatrix" planeSlot . fromMat4
          setCubeModelMatrices = [uniformM44F "modelMatrix" cubeSlot . fromMat4 | cubeSlot <- cubeSlots]
          setReflectorModelMatrix = uniformM44F "modelMatrix" reflectorSlot . fromMat4
          
          setupRendering ((_, _, fps), frameCount, capture) (windowWidth, windowHeight) (cameraPosition, cameraDirection, cameraUp, _) time = do
              let aspect = fromIntegral windowWidth / fromIntegral windowHeight
                  
                  cameraView = fromProjective (lookat cameraPosition (cameraPosition &+ cameraDirection) cameraUp)
                  cameraProjection = perspective 0.1 50 (pi/2) aspect
  
                  lightPosition = Vec3 (15 * sin time) 2 10
                  reflectorPosition = Vec3 (-8) (5 * sin (time * 0.25)) 0
                  
                  cubeCameraProjection = perspective 0.1 50 (pi/2) 1
                  cubeLookAt dir up = fromProjective (lookat reflectorPosition (reflectorPosition &+ dir) up)
                  cubeCameraMatrix 1 = cubeLookAt (Vec3 1 0 0) (Vec3 0 (-1) 0)
                  cubeCameraMatrix 2 = cubeLookAt (Vec3 (-1) 0 0) (Vec3 0 (-1) 0)
                  cubeCameraMatrix 3 = cubeLookAt (Vec3 0 1 0) (Vec3 0 0 1)
                  cubeCameraMatrix 4 = cubeLookAt (Vec3 0 (-1) 0) (Vec3 0 0 (-1))
                  cubeCameraMatrix 5 = cubeLookAt (Vec3 0 0 1) (Vec3 0 (-1) 0)
                  cubeCameraMatrix 6 = cubeLookAt (Vec3 0 0 (-1)) (Vec3 0 (-1) 0)
              
              case fps of
                  Just value -> putStrLn $ "FPS: " ++ show value
                  Nothing -> return ()
              
              setViewCameraMatrix (cameraView .*. cameraProjection)
              setViewCameraPosition cameraPosition
              setLightPosition lightPosition
              
              setCubeCameraPosition reflectorPosition
              setReflectorModelMatrix (fromProjective (translation reflectorPosition))
              forM_ [1..6] $ \index -> setCubeCameraMatrix index (cubeCameraMatrix index .*. cubeCameraProjection)
      
              setPlaneModelMatrix (fromProjective $ scaling (Vec3 12 12 1) .*. translation (Vec3 0 (-2) (-12)))
              forM_ (zip setCubeModelMatrices [0..]) $ \(setCubeModelMatrix, i) -> do
                  let t = i * 2 * pi / 5
                      s = (t + 2) * 0.3
                      trans = scaling (Vec3 s s s) .*. rotationEuler (Vec3 0 0 s) .*. translation (Vec3 (t * 0.3) (sin t * 4) (cos t * 4))
                  setCubeModelMatrix (fromProjective trans)
              setSize (fromIntegral windowWidth) (fromIntegral windowHeight)
              
              return $ do
  #ifdef CAPTURE
                  when capture $ do
                      glFinish
                      withFrameBuffer 0 0 windowWidth windowHeight $ writeImageFromPtr (printf "frame%08d.jpg" frameCount) (windowHeight, windowWidth)
  #endif
                  return ()
      
      effectful4 setupRendering ((,,) <$> fpsTracking <*> frameCount <*> capture) windowSize camera time
  
  --readInput :: IO (Maybe Float)
  readInput keyIsPressed = do
      Just t <- GLFW.getTime
      GLFW.setTime 0
  
      k <- keyIsPressed GLFW.Key'Escape
      return $ if k then Nothing else Just (realToFrac t)
  
  sceneRender :: Exp Obj (FrameBuffer 1 (Float, V4F))
  sceneRender = Accumulate accCtx PassAll reflectFrag (Rasterize rastCtx reflectPrims) directRender
    where
      directRender = Accumulate accCtx PassAll frag (Rasterize rastCtx directPrims) clearBuf
      cubeMapRender = Accumulate accCtx PassAll frag (Rasterize rastCtx cubePrims) clearBuf6
      
      accCtx = AccumulationContext Nothing (DepthOp Less True :. ColorOp NoBlending (one' :: V4B) :. ZT)
      rastCtx = triangleCtx { ctxCullMode = CullFront CCW }
      clearBuf = FrameBuffer (DepthImage n1 1000 :. ColorImage n1 (V4 0.1 0.2 0.6 1) :. ZT)
      clearBuf6 = FrameBuffer (DepthImage n6 1000 :. ColorImage n6 (V4 0.05 0.1 0.3 1) :. ZT)
      worldInput = Fetch "geometrySlot" Triangles (IV3F "position", IV3F "normal")
      reflectInput = Fetch "reflectSlot" Triangles (IV3F "position", IV3F "normal")
      directPrims = Transform directVert worldInput
      cubePrims = Reassemble geom (Transform cubeMapVert worldInput)
      reflectPrims = Transform directVert reflectInput
  
      lightPosition = Uni (IV3F "lightPosition")
      viewCameraMatrix = Uni (IM44F "viewCameraMatrix")
      viewCameraPosition = Uni (IV3F "viewCameraPosition")
      cubeCameraMatrix i = Uni (IM44F (cubeMatrixName i))
      cubeCameraPosition = Uni (IV3F "cubeCameraPosition")
      modelMatrix = Uni (IM44F "modelMatrix")
      
      transformGeometry :: Exp f V4F -> Exp f V3F -> Exp f M44F -> (Exp f V4F, Exp f V4F, Exp f V3F)
      transformGeometry localPos localNormal viewMatrix = (viewPos, worldPos, worldNormal)
        where
          worldPos = modelMatrix @*. localPos
          viewPos = viewMatrix @*. worldPos
          worldNormal = normalize' (v4v3 (modelMatrix @*. n3v4 localNormal))
  
      directVert :: Exp V (V3F, V3F) -> VertexOut () (V3F, V3F, V3F)
      directVert attr = VertexOut viewPos (floatV 1) ZT (Smooth (v4v3 worldPos) :. Smooth worldNormal :. Flat viewCameraPosition :. ZT)
        where
          (localPos, localNormal) = untup2 attr
          (viewPos, worldPos, worldNormal) = transformGeometry (v3v4 localPos) localNormal viewCameraMatrix
  
      cubeMapVert :: Exp V (V3F, V3F) -> VertexOut () V3F
      cubeMapVert attr = VertexOut (v3v4 localPos) (floatV 1) ZT (Smooth localNormal :. ZT)
        where
          (localPos, localNormal) = untup2 attr
  
      geom :: GeometryShader Triangle Triangle () () 6 V3F (V3F, V3F, V3F) 
      geom = GeometryShader n6 TrianglesOutput 18 init prim vert
        where
          init attr = tup2 (primInit, intG 6)
            where
              primInit = tup2 (intG 0, attr)
          prim primState = tup5 (layer, layer, primState', vertInit, intG 3) 
            where
              (layer, attr) = untup2 primState
              primState' = tup2 (layer @+ intG 1, attr)
              vertInit = tup3 (intG 0, viewMatrix, attr)
              viewMatrix = indexG (map cubeCameraMatrix [1..6]) layer
          vert vertState = GeometryOut vertState' viewPos pointSize ZT (Smooth (v4v3 worldPos) :. Smooth worldNormal :. Flat cubeCameraPosition :. ZT)
            where
              (index, viewMatrix, attr) = untup3 vertState
              vertState' = tup3 (index @+ intG 1, viewMatrix, attr)
              (attr0, attr1, attr2) = untup3 attr
              (localPos, pointSize, _, localNormal) = untup4 (indexG [attr0, attr1, attr2] index)
              (viewPos, worldPos, worldNormal) = transformGeometry localPos localNormal viewMatrix
  
      frag :: Exp F (V3F, V3F, V3F) -> FragmentOut (Depth Float :+: Color V4F :+: ZZ)
      frag attr = FragmentOutRastDepth (luminance :. ZT)
        where
          lambert = max' (floatF 0) (worldNormal @. normalize' (lightPosition @- worldPos))
          reflectedRay = normalize' (reflect' (worldPos @- (cameraPosition :: Exp F V3F)) worldNormal)
          directLight = normalize' (lightPosition @- worldPos)
          phong = max' (floatF 0) (reflectedRay @. directLight)
          colour = pack' (V3 (floatF 0.7) (floatF 0.05) (floatF 0))
          luminance = v3v4 (colour @* lambert @+ pow' phong (floatF 10))
          (worldPos, worldNormal, cameraPosition) = untup3 attr
  
      reflectFrag :: Exp F (V3F, V3F, V3F) -> FragmentOut (Depth Float :+: Color V4F :+: ZZ)
      reflectFrag attr = FragmentOutRastDepth (luminance :. ZT)
        where
          reflectedRay = reflect' (worldPos @- (cameraPosition :: Exp F V3F)) worldNormal
          luminance = reflectionSample reflectedRay
          (worldPos, worldNormal, cameraPosition) = untup3 attr
          reflectionSample dir = texture' (Sampler LinearFilter ClampToEdge reflectionMap) dir
          reflectionMap = Texture (TextureCube (Float RGBA)) (V2 256 256) NoMip [PrjFrameBuffer "" tix0 cubeMapRender]
  
  indexG :: GPU a => [Exp G a] -> Exp G Int32 -> Exp G a
  indexG xs index = go xs 0
    where
      go [x] _ = x
      go (x:xs) i = Cond (index @== intG i) x (go xs (i+1))
  
  cubeMatrixName :: Int -> SB.ByteString
  cubeMatrixName i = SB.pack (printf "cubeCameraMatrix%d" i)
  
  