-- GearBox -- zooming rotating fractal gearbox demo
-- (GPL3+) 2012 Claude Heiland-Allen <claude@mathr.co.uk>
-- usage: gearbox <quality>
-- default quality is 9, useful range 6-14
-- 'f' toggles full screen, any other key to quit
-- tested with: ghc-7.0.4, ghc-7.4-rc1, ghc-7.6.1

{-# LANGUAGE TypeOperators #-}
import Control.Monad (forM, forM_, replicateM)
import Data.Fixed (mod')
import Data.IORef (IORef, newIORef, readIORef, modifyIORef)
import Data.Vec hiding (last, map, scale, take)
import Data.Vec.OpenGLRaw ()
import Foreign (Ptr, Storable(..), alloca, allocaArray, castPtr, nullPtr, with, withArray)
import Foreign.C (withCString, peekCStringLen)
import Graphics.UI.GLUT
  ( Size(Size)
  , DisplayMode(RGBMode, DoubleBuffered)
  , Key(Char)
  , KeyState(Down)
  , KeyboardMouseCallback
  , ($=)
  , addTimerCallback
  , createWindow
  , displayCallback
  , getArgsAndInitialize
  , initialDisplayMode
  , initialWindowSize
  , keyboardMouseCallback
  , mainLoop
  , postRedisplay
  , reportErrors
  , reshapeCallback
  , swapBuffers
  , fullScreenToggle
  )
import Graphics.Rendering.OpenGL.Raw
import System.Exit (exitSuccess)

-- 4x4 matrix type
type R = GLdouble
type V = R :. R :. R :. R :. ()
type M = V :. V :. V :. V :. ()

-- length of animation loop in frames
looplength :: Int
looplength = 60 * 25

-- number of passes to iterate function system
passes :: Int
passes = 6

-- state for callbacks
data IFS = IFS
  -- immutable state
  { fbo :: GLuint
  , t_source, t_ping, t_pong :: GLuint
  , p_shift :: GLuint, p_shift_u_t, p_shift_u_k :: GLint
  , w_tsize, w_width, w_height :: Int
  -- mutable state
  , w_frame :: Int, w_phase :: GLdouble
  }

-- shader to generate a single blue gear wheel
gear_frag_src :: String
gear_frag_src = unlines
  [ "void main() {"
  , "  vec2 p = (gl_TexCoord[0].xy - vec2(0.5)) * 2.0;"
  , "  float r = length(p);"
  , "  float a = atan(p.y, p.x);"
  , "  float ro = 0.94 + 0.03 * sin(30.0 * a);"
  , "  float ri = 0.82 - 0.01 * sin(90.0 * a);"
  , "  if (r < ro) {"
  , "    gl_FragColor = ri < r ? vec4(0.5,0.5,1.0,1.0) : vec4(0.25,0.25,0.5,1.0);"
  , "  } else {"
  , "    discard;"
  , "  }"
  , "}"
  ]

-- shader to rotate colours in YUV space
shift_frag_src :: String
shift_frag_src = unlines
  [ "uniform sampler2D t;"
  , "uniform float k;"
  , ""
  , "void main() {"
  , "  float c = cos(6.283185307179586 * k);"
  , "  float s = sin(6.283185307179586 * k);"
  , "  vec4 p = texture2D(t, gl_TexCoord[0].xy).rgba;"
  , "  mat3 yuv2rgb = mat3(1.0, 0.0, 1.28033, 1.0, -0.21482, -0.38059, 1.0, 2.12798, 0.0);" -- rec.709
  , "  mat3 rgb2yuv = mat3(0.2126, 0.7152, 0.0722, -0.09991, -0.33609, 0.43600, 0.615, -0.5586, -0.05639);"
  , "  p.rgb = clamp(p.rgb * rgb2yuv * mat3(1, 0, 0, 0, c, s, 0, -s, c) * yuv2rgb, 0.0, 1.0);"
  , "  gl_FragColor = p;"
  , "}"
  ]

-- draw a textured quad
unitQuad :: IO ()
unitQuad = do
  glBegin gl_QUADS
  glTexCoord2f 0 0
  glVertex2f (-1) (-1)
  glTexCoord2f 1 0
  glVertex2f 1 (-1)
  glTexCoord2f 1 1
  glVertex2f 1 1
  glTexCoord2f 0 1
  glVertex2f (-1) 1
  glEnd

-- load a matrix into GL
loadMatrix :: M -> IO ()
loadMatrix ((a:.b:.c:.d:.()):.(e:.f:.g:.h:.()):.(i:.j:.k:.l:.()):.(m:.n:.o:.p:.()):.()) =
  withArray [a,b,c,d,e,f,g,h,i,j,k,l,m,n,o,p] $ glLoadMatrixd

-- helper for ping-pong technique
swap :: (a, b) -> (b, a)
swap ~(a,b) = (b,a)

-- simple transformations
scale :: R -> M
scale s = diagonal (s :. s :. s :. 1 :. ())
mirror :: M
mirror = diagonal (1 :. (-1) :. 1 :. 1 :. ())
rotateXY :: R -> M
rotateXY t = (c:.s:.0:.0:.()):.((-s):.c:.0:.0:.()):.(0:.0:.1:.0:.()):.(0:.0:.0:.1:.()):.()
  where c = cos t ; s = sin t
translateX :: R -> M
translateX d = (1:.0:.0:.0:.()):.(0:.1:.0:.0:.()):.(0:.0:.1:.0:.()):.(d:.0:.0:.1:.()):.()
(<>) :: M -> M -> M
(<>) = multmm

-- display callback
display :: IORef IFS -> IO ()
display ifsR = do
  ifs <- readIORef ifsR
  -- compute transforms for iterated function system fractal
  let p0 = w_phase ifs
      f0 = w_frame ifs
      q = f0 `mod` looplength
      k = q < (looplength `div` 2)
      p = if k then p0 else -p0
      s = 0.28
      d = 0.54
      transforms =
        [ identity
        , mirror <> scale s <> rotateXY ( 4 * p)
        ] ++ 
        [ mirror <> scale s <> rotateXY (-3 * p) <> translateX d <>
          rotateXY (2 * pi * fromIntegral (if k then i else -i) / 5 + p)
        | i <- [0 .. 4 :: Int]
        ]
  -- iterated function system from gear to gears
  glViewport 0 0 (fromIntegral $ w_tsize ifs) (fromIntegral $ w_tsize ifs)
  glMatrixMode gl_PROJECTION
  glLoadIdentity
  glOrtho (-1) 1 (-1) 1 (-1) 1
  glMatrixMode gl_MODELVIEW
  glLoadIdentity
  glBindFramebuffer gl_FRAMEBUFFER (fbo ifs)
  glUseProgramObject (fromIntegral $ p_shift ifs)
  glUniform1i (p_shift_u_t ifs) 0
  to <- fmap last $ forM (take passes . ((True, (t_source ifs, t_ping ifs)) :) .
      zip (repeat False) $ (iterate swap (t_ping ifs, t_pong ifs))) $
      \(pass1, (ping, pong)) -> do
    glBindTexture gl_TEXTURE_2D ping
    glGenerateMipmap gl_TEXTURE_2D
    glFramebufferTexture2D gl_FRAMEBUFFER gl_COLOR_ATTACHMENT0 gl_TEXTURE_2D pong 0
    glClear (fromIntegral gl_COLOR_BUFFER_BIT)
    forM_ ([0 .. 6 :: Int] `zip` transforms) $ \(w, m) -> do
      let sk = case w of
            0 -> 0
            1 -> 0
            _ -> fromIntegral (4 * w - 7) / 10
      glUniform1f (p_shift_u_k ifs) sk
      loadMatrix $ (if pass1 && not k then rotateXY (pi / 30) else identity) <> m
      unitQuad
    glBindTexture gl_TEXTURE_2D 0
    glFramebufferTexture2D gl_FRAMEBUFFER gl_COLOR_ATTACHMENT0 gl_TEXTURE_2D 0 0
    return pong
  glUseProgramObject 0
  glBindFramebuffer gl_FRAMEBUFFER 0
  -- display the output scaled to the window
  let w = fromIntegral $ w_width  ifs
      h = fromIntegral $ w_height ifs
  glViewport 0 0 w h
  glClear (fromIntegral gl_COLOR_BUFFER_BIT)
  glMatrixMode gl_PROJECTION
  glLoadIdentity
  let (ax, ay)
        | w >= h = (1, fromIntegral h / fromIntegral w)
        | otherwise = (fromIntegral w / fromIntegral h, 1)
  glOrtho (-ax) ax (-ay) ay (-1) 1
  glMatrixMode gl_MODELVIEW
  glLoadIdentity
  glBindTexture gl_TEXTURE_2D to
  glGenerateMipmap gl_TEXTURE_2D
  loadMatrix $ scale 1.5 <>
    scale (s ** negate ((p0 / (2 * pi)) `mod'` 1)) <> rotateXY (- 2 * p)
  unitQuad
  glBindTexture gl_TEXTURE_2D 0
  -- animate
  modifyIORef ifsR $ \ifs' -> 
     let w_frame' = w_frame ifs' + 1
     in  ifs'
           { w_frame = w_frame'
           , w_phase = fromIntegral w_frame' * 4 * pi / fromIntegral looplength
           }
  swapBuffers
  reportErrors

-- timer callback (25fps)
timer :: IORef IFS -> IO ()
timer ifsR = do
  addTimerCallback 40 (timer ifsR)
  postRedisplay Nothing

-- window reshape callback
reshape :: IORef IFS -> Size -> IO ()
reshape ifsR (Size w h) = modifyIORef ifsR $ \ifs -> ifs
  { w_width = fromIntegral w, w_height = fromIntegral h }

-- compile a shader
shader :: String -> IO (Either String GLuint)
shader frag = do
  p <- glCreateProgramObject
  f <- glCreateShaderObject gl_FRAGMENT_SHADER
  withCString frag $ \s -> with s $ \src -> do
    glShaderSource (fromIntegral f) 1 (castPtr src) nullPtr
    glCompileShader (fromIntegral f)
    glAttachObject p f
    glLinkProgram (fromIntegral p)
    alloca $ \sp -> do
      glGetObjectParameteriv (fromIntegral p) gl_OBJECT_LINK_STATUS sp
      success <- peek sp
      if (success == 0)
        then
          alloca $ \lp -> do
            glGetProgramiv (fromIntegral p) gl_INFO_LOG_LENGTH lp
            l <- peek lp
            poke lp 0
            allocaArray (fromIntegral l) $ \lg -> do
              glGetProgramInfoLog (fromIntegral p) (fromIntegral l) (castPtr lp) lg
              l' <- peek lp
              ss <- peekCStringLen (castPtr lg, fromIntegral l')
              return (Left ss)
        else return (Right (fromIntegral p))

-- generate a texture
newTexture :: Int -> IO GLuint
newTexture size = alloca $ \pt -> do
  glGenTextures 1 pt
  t <- peek pt
  glBindTexture gl_TEXTURE_2D t
  let s = fromIntegral size
  glTexImage2D gl_TEXTURE_2D 0 (fromIntegral gl_RGBA) s s 0
    gl_RGBA gl_UNSIGNED_BYTE (nullPtr :: Ptr GLubyte)
  glTexParameteri gl_TEXTURE_2D gl_TEXTURE_MIN_FILTER (fromIntegral gl_LINEAR_MIPMAP_LINEAR)
  glTexParameteri gl_TEXTURE_2D gl_TEXTURE_MAG_FILTER (fromIntegral gl_LINEAR)
  glTexParameteri gl_TEXTURE_2D gl_TEXTURE_WRAP_S (fromIntegral gl_CLAMP_TO_EDGE)
  glTexParameteri gl_TEXTURE_2D gl_TEXTURE_WRAP_T (fromIntegral gl_CLAMP_TO_EDGE)
  glBindTexture gl_TEXTURE_2D 0
  return t

-- keyboard callback
keyboard :: IORef IFS -> KeyboardMouseCallback
keyboard _  (Char 'f') Down _ _ = fullScreenToggle
keyboard _  (Char _) Down _ _ = exitSuccess
keyboard _ _ _ _ _ = return ()

-- main program
main :: IO ()
main = do
  let w = 512
      h = 288
  initialWindowSize $= Size (fromIntegral w) (fromIntegral h)
  initialDisplayMode $= [ RGBMode, DoubleBuffered ]
  (_, args) <- getArgsAndInitialize
  quality <- case map reads args of
    [[(q,"")]] | 5 < q && q < 15-> return q
    _ -> return (9 :: Int)
  let tsize' = 2 ^ quality
  _ <- createWindow "GearBox"
  -- compile shaders (FIXME: print log to stderr on failure and exit)
  Right p_shift' <- shader shift_frag_src
  p_shift_u_t' <- withCString "t" $ glGetUniformLocation p_shift' . castPtr
  p_shift_u_k' <- withCString "k" $ glGetUniformLocation p_shift' . castPtr
  Right p_gear' <- shader gear_frag_src
  -- misc setup
  glHint gl_GENERATE_MIPMAP_HINT gl_NICEST
  glEnable gl_TEXTURE_2D
  glEnable gl_BLEND
  glBlendFunc gl_SRC_ALPHA gl_ONE_MINUS_SRC_ALPHA
  glClearColor 0.25 0.25 0.5 0
  -- allocate textures
  [t_source', t_ping', t_pong'] <- replicateM 3 (newTexture tsize')
  -- allocate frame buffer
  fbo' <- alloca $ \p -> do
    glGenFramebuffers 1 p
    peek p
  -- initialize state
  ifsR <- newIORef IFS
    { fbo = fbo', t_source = t_source', t_ping = t_ping', t_pong = t_pong'
    , p_shift = p_shift', p_shift_u_t = p_shift_u_t', p_shift_u_k = p_shift_u_k'
    , w_width = w, w_height = h, w_tsize = tsize'
    , w_frame = 0, w_phase = 0
    }
  -- render the gear texture
  glViewport 0 0 (fromIntegral tsize') (fromIntegral tsize')
  glMatrixMode gl_PROJECTION
  glLoadIdentity
  glOrtho (-1) 1 (-1) 1 (-1) 1
  glMatrixMode gl_MODELVIEW
  glLoadIdentity
  glBindFramebuffer gl_FRAMEBUFFER fbo'
  glFramebufferTexture2D gl_FRAMEBUFFER gl_COLOR_ATTACHMENT0 gl_TEXTURE_2D t_source' 0
  glClear (fromIntegral gl_COLOR_BUFFER_BIT)
  glUseProgramObject (fromIntegral p_gear')
  unitQuad
  glUseProgramObject 0
  glFramebufferTexture2D gl_FRAMEBUFFER gl_COLOR_ATTACHMENT0 gl_TEXTURE_2D 0 0
  glBindFramebuffer gl_FRAMEBUFFER 0
  glBindTexture gl_TEXTURE_2D t_source'
  glGenerateMipmap gl_TEXTURE_2D
  glBindTexture gl_TEXTURE_2D 0
  -- set callbacks
  reshapeCallback $= Just (reshape ifsR)
  displayCallback $= display ifsR
  addTimerCallback 40 (timer ifsR)
  keyboardMouseCallback $= Just (keyboard ifsR)
  reportErrors
  mainLoop