{-| Draw and update waterfall plots with OpenGL. Useful for spectrograms.
-}
module Graphics.DynamicGraph.Waterfall (
    jet,
    jet_mod,
    hot,
    bw,
    wb,
    graph
    ) where

import Control.Monad
import Graphics.UI.GLFW as G
import Graphics.Rendering.OpenGL
import Graphics.GLUtil

import Control.Monad.Trans.Class
import Control.Monad.Trans.Either
import Foreign.Storable
import Foreign.Marshal.Array

import Pipes

import Paths_dynamic_graph

-- | The matlab / octave \"jet\" color map
jet :: [GLfloat]
jet =  [0, 0, 0.5,  0, 0, 1,  0, 0.5, 1,   0, 1, 1,  0.5, 1, 0.5,  1, 1, 0,  1, 0.5, 0,  1, 0, 0,  0.5, 0, 0]

-- | \"jet\" modified so that low values are a darker blue
jet_mod :: [GLfloat]
jet_mod =  [0, 0, 0.1,  0, 0, 1,  0, 0.5, 1,   0, 1, 1,  0.5, 1, 0.5,  1, 1, 0,  1, 0.5, 0,  1, 0, 0,  0.5, 0, 0]

-- | The matlab / octave \"hot\" color map
hot :: [GLfloat]
hot =  [0, 0, 0,  1, 0, 0,  1, 1, 0,  1, 1, 1]

-- | Ranges from black to white
bw :: [GLfloat]
bw =  [0, 0, 0, 1, 1, 1]

-- | Ranges from white to black
wb :: [GLfloat]
wb =  [1, 1, 1, 0, 0, 0]

{-| @(graph windowWidth windowHeight width height colormap)@ creates
    a window of width @windowWidth@ and height @windowHeight@ for displaying
    a waterfall plot. A Consumer is returned for updating the waterfall
    plot. Feeding an instance of IsPixelData of length @width@ shifts all
    rows of the waterfall down and updates the top row with the data. The
    waterfall is @height@ rows high. @colorMap@ is used to map values to
    display color.
-}
graph :: IsPixelData a => Int -> Int -> Int -> Int -> [GLfloat] -> EitherT String IO (Consumer a IO ())
graph windowWidth windowHeight width height colorMap = do
    res' <- lift $ createWindow windowWidth windowHeight "" Nothing Nothing
    win <- maybe (left "error creating window") return res'

    lift $ do
        makeContextCurrent (Just win)

        --Load the shaders
        vertFN <- getDataFileName "shaders/waterfall.vert"
        fragFN <- getDataFileName "shaders/waterfall.frag"
        vs <- loadShader VertexShader   vertFN
        fs <- loadShader FragmentShader fragFN
        p  <- linkShaderProgram [vs, fs]

        --Set stuff
        currentProgram $= Just p

        ab <- genObjectName 

        loc <- get $ attribLocation p "coord"

        let stride = fromIntegral $ sizeOf (undefined::GLfloat) * 2
            vad    = VertexArrayDescriptor 2 Float stride offset0

        bindBuffer ArrayBuffer  $= Just ab
        vertexAttribArray   loc $= Enabled
        vertexAttribPointer loc $= (ToFloat, vad)

        let xCoords :: [GLfloat]
            xCoords = [-1, -1, 1, -1, 1, 1, -1, 1]
        withArray xCoords $ \ptr -> 
            bufferData ArrayBuffer $= (fromIntegral $ sizeOf(undefined::GLfloat) * 8, ptr, StaticDraw)

        let yCoords :: [GLfloat]
            yCoords = take (width * height) $ repeat 0

        activeTexture $= TextureUnit 0
        texture Texture2D $= Enabled
        to <- loadTexture (TexInfo (fromIntegral width) (fromIntegral height) TexMono yCoords)
        
        loc <- get $ uniformLocation p "texture"
        asUniform (0 :: GLint) loc 

        textureFilter Texture2D $= ((Linear', Nothing), Linear')
        textureWrapMode Texture2D S $= (Repeated, ClampToEdge)
        textureWrapMode Texture2D T $= (Repeated, Repeat)

        activeTexture $= TextureUnit 1
        texture Texture2D $= Enabled
        to <- loadTexture (TexInfo (fromIntegral $ length colorMap `quot` 3) 1 TexRGB colorMap)
        textureFilter Texture2D $= ((Linear', Nothing), Linear')
        textureWrapMode Texture2D S $= (Repeated, ClampToEdge)
        textureWrapMode Texture2D T $= (Repeated, ClampToEdge)

        loc <- get $ uniformLocation p "colorMap"
        asUniform (1 :: GLint) loc 

        let lcm :: GLfloat
            lcm = fromIntegral $ length colorMap `quot` 3
        loc <- get $ uniformLocation p "scale"
        asUniform ((lcm - 1) / lcm) loc 

        loc <- get $ uniformLocation p "offset"
        asUniform (0.5 / lcm) loc 

        --No idea why this is needed
        activeTexture $= TextureUnit 0

        loc <- get $ uniformLocation p "voffset"

        let pipe yoffset = do
                dat <- await

                lift $ do
                    makeContextCurrent (Just win)

                    let textureOffset = (yoffset + height - 1) `mod` height

                    withPixels dat $ \ptr -> texSubImage2D Texture2D 0 (TexturePosition2D 0 (fromIntegral textureOffset)) (TextureSize2D (fromIntegral width) 1) (PixelData Red Float ptr)

                    asUniform (fromIntegral yoffset / fromIntegral height :: GLfloat) loc

                    drawArrays Quads 0 4
                    swapBuffers win

                pipe $ if yoffset + 1 >= height then 0 else yoffset + 1

        return $ pipe 0