-- SnowGlobe -- randomized fractal snowflakes demo
-- (GPL3+) 2012 Claude Heiland-Allen <claudiusmaximus@goto10.org>
-- usage: snowglobe
-- 'f' toggles full screen, any other key to quit
-- tested with: ghc-7.0.4, ghc-7.4-rc1

import Numeric.LinearAlgebra hiding (i, reshape)
import qualified Numeric.LinearAlgebra as L
import Graphics.UI.GLUT hiding (scale)
import qualified Graphics.UI.GLUT as GL
import Graphics.Rendering.OpenGL.Raw
  ( glTexImage2D, gl_TEXTURE_2D, gl_LUMINANCE32F, gl_LUMINANCE, gl_RGBA
  , gl_UNSIGNED_BYTE, gl_FALSE, glClampColor, glGenerateMipmap
  , gl_CLAMP_VERTEX_COLOR, gl_CLAMP_READ_COLOR, gl_CLAMP_FRAGMENT_COLOR
  , glGenFramebuffers, glBindFramebuffer, glFramebufferTexture2D
  , gl_FRAMEBUFFER, gl_COLOR_ATTACHMENT0, glUniformMatrix3fv
  )
import Control.Exception (evaluate)
import Control.Monad (forM_, replicateM, when)
import Data.IORef (IORef, modifyIORef, newIORef, readIORef, writeIORef)
import Data.List (foldl', minimumBy)
import Data.Ord (comparing)
import Data.Map (Map)
import qualified Data.Map as M
import qualified Data.Set as S
import System.Exit (exitSuccess)
import System.IO (hPutStrLn, stderr)
import System.Random (randomRIO)
import Unsafe.Coerce (unsafeCoerce) -- for OpenGL newtypes
import Foreign (alloca, peek, nullPtr, withArray)

shader :: Maybe String -> Maybe String -> IO Program
shader mV mF = do
  [p] <- genObjectNames 1
  vs <- case mV of
    Nothing -> return []
    Just v -> do
      [vert] <- genObjectNames 1
      shaderSource vert $= [ v ]
      compileShader vert
      msg <- get (shaderInfoLog vert)
      when (not (null msg)) $ hPutStrLn stderr msg
      return [vert]
  fs <- case mF of
    Nothing -> return []
    Just f -> do
      [frag] <- genObjectNames 1
      shaderSource frag $= [ f ]
      compileShader frag
      msg <- get (shaderInfoLog frag)
      when (not (null msg)) $ hPutStrLn stderr msg
      return [frag]
  attachedShaders p $= (vs, fs)
  linkProgram p
  msg <- get (programInfoLog p)
  when (not (null msg)) $ hPutStrLn stderr msg
  return p

type N = Int
type R = Double
type V = L.Vector R
type M = L.Matrix R

sum' :: Num a => [a] -> a
sum' = foldl' (+) 0

{-# RULES "realToFrac/glDouble" realToFrac = unsafeCoerceGLdouble #-}
unsafeCoerceGLdouble :: Double -> GLdouble
unsafeCoerceGLdouble = unsafeCoerce

{-# RULES "realToFrac/glFloat" realToFrac = unsafeCoerceGLfloat #-}
unsafeCoerceGLfloat :: Float -> GLfloat
unsafeCoerceGLfloat = unsafeCoerce

size :: N
size = 1024

er :: R
er = 4

accuracy :: R
accuracy = 1 / fromIntegral size

tSize :: N
tSize = 512

maxSpeed :: R
maxSpeed = 10

data SnowGlobe = SnowGlobe
  { pInitial :: Program, pInitial'er, pInitial'rho :: UniformLocation
  , pStep :: Program, pStep'er, pStep'ts, pStep'src :: UniformLocation
  , pColour :: Program, pColour'src, pColour'speed :: UniformLocation
  , tPing, tPong :: TextureObject, fBuffer :: FramebufferObject
  , sFlakes :: [Flake], sTextures :: Map N TextureObject
  , wSize :: Size, wFullScreen :: Maybe Size
  , sRenderFlake :: IO Render, sNextName :: N
  }

data Flake = Flake
  { flakeName :: !N
  , flakeTime :: !R
  , flakePosition :: !V
  , flakeVelocity :: !V
  }

flakeMass :: Flake -> R
flakeMass f = let s = flakeSize f in 1 + s * s

flakeSize :: Flake -> R
flakeSize f = sin (pi * flakeTime f)

flakeForce :: Flake -> V -> V
flakeForce f p =
  let d = p - flakePosition f
      [x,y] = toList d
      m = let d2 = x * x + y * y in d2 * d2
  in  scale (flakeMass f / m) d

flakeUpdate :: R -> [Flake] -> Flake -> Flake
flakeUpdate dt fs f = f
  { flakeTime = flakeTime f + dt / 2
  , flakePosition = flakePosition f + scale dt (flakeVelocity f)
  , flakeVelocity = mapVector ((* maxSpeed) . tanh . (/ maxSpeed)) $ 0.999 * (flakeVelocity f + scale (dt / flakeMass f) (flakeField fs (flakePosition f)))
  }

flakeField :: [Flake] -> V -> V
flakeField fs p = sum' [ flakeForce f (p + d) | f <- fs, p /= flakePosition f, dx <- [-2,0,2], dy <- [-2,0,2], let d = 2|>[dx,dy] ]

flakeSpawn :: [Flake] -> N -> IO Flake
flakeSpawn fs name = do
  xs <- replicateM 32 $ randomRIO (-1, 1)
  ys <- replicateM 32 $ randomRIO (-1, 1)
  let ps = zipWith (\x y -> 2|>[x,y]) xs ys
      p = fst . minimumBy (comparing snd) . map (\p' -> (p', norm2 (flakeField fs p'))) $ ps
  return Flake{ flakeName = name, flakeTime = 0, flakePosition = p, flakeVelocity = 0 }

flakesUpdate :: R -> [Flake] -> [Flake]
flakesUpdate dt fs =
  let gs = map (flakeUpdate dt fs) fs
      alive = filter ((< 1) . flakeTime) gs
  in  map flakeWrap alive

flakeWrap :: Flake -> Flake
flakeWrap f = f{ flakePosition = mapVector wrap (flakePosition f) }
  where
    wrap x =
      let y = (x + 1) / 2
          z = y - fromIntegral (floor y :: N)
      in  z * 2 - 1

-- amortized rendering over several frames
data Render = Done TextureObject | Step (IO Render)

-- initial pass
flakeRenderStart :: IORef SnowGlobe -> IO Render
flakeRenderStart sR = do
  s0 <- readIORef sR
  p <- replicateM 4 $ randomRIO (0.02, 0.98)
  let rts = flakeTransforms p
      rho = maximum (map fst rts)
      ts = map snd rts
      passes = clamp 4 256 . round . logBase rho $ accuracy
  loadIdentity
  ortho2D 0 1 0 1
  viewport $= (Position 0 0, Size (fromIntegral size) (fromIntegral size))
  currentProgram $= Just (pInitial s0)
  uniform (pInitial'er s0) $= TexCoord1 (realToFrac er :: GLfloat)
  uniform (pInitial'rho s0) $=  TexCoord1 (realToFrac rho :: GLfloat)
  bindFBO (fBuffer s0) (tPing s0)
  unitQuad
  unbindFBO
  currentProgram $= Nothing
  return $ Step (flakeRenderPass sR ts passes passes)

-- multi step passes
flakeRenderPass :: IORef SnowGlobe -> [M] -> N -> N -> IO Render
flakeRenderPass sR _  passes 0 = flakeRenderFinish sR passes
flakeRenderPass sR ts passes n = do
  s0 <- readIORef sR
  loadIdentity
  ortho2D 0 1 0 1
  viewport $= (Position 0 0, Size (fromIntegral size) (fromIntegral size))
  currentProgram $= Just (pStep s0)
  uniform (pStep'er s0) $= TexCoord1 (realToFrac er :: GLfloat)
  uniform (pStep'src s0) $= TexCoord1 (0 :: GLint)
  withArray (map ((realToFrac :: Float -> GLfloat) . (realToFrac :: Double -> Float)) . concat . concatMap L.toLists $ ts) $ glUniformMatrix3fv (unsafeCoerce $ pStep'ts s0) 6 1
  bindFBO (fBuffer s0) (tPong s0)
  textureBinding Texture2D $= Just (tPing s0)
  unitQuad
  textureBinding Texture2D $= Nothing
  unbindFBO
  writeIORef sR s0{ tPing = tPong s0, tPong = tPing s0 }
  currentProgram $= Nothing
  if even n
    then return $ Step (flakeRenderPass sR ts passes (n - 1))
    else                flakeRenderPass sR ts passes (n - 1)

-- colourize pass
flakeRenderFinish :: IORef SnowGlobe -> N -> IO Render
flakeRenderFinish sR passes = do
  s0 <- readIORef sR
  t <- newTexRGBA tSize
  bindFBO (fBuffer s0) t
  loadIdentity
  ortho2D (-1) 1 (-1) 1
  viewport $= (Position 0 0, Size (fromIntegral tSize) (fromIntegral tSize))
  textureBinding Texture2D $= Just (tPong s0)
  currentProgram $= Just (pColour s0)
  uniform (pColour'src s0) $= TexCoord1 (0 :: GLint)
  uniform (pColour'speed s0) $= TexCoord1 (1 / fromIntegral passes :: GLfloat)
  fullQuad
  currentProgram $= Nothing
  unbindFBO
  textureBinding Texture2D $= Just t
  glGenerateMipmap gl_TEXTURE_2D
  textureBinding Texture2D $= Nothing
  return $ Done t

transformRST :: R -> R -> V -> M
transformRST a l p = (3><3)[ c, s, x, -s, c, y, 0, 0, 1 ]
  where x:y:_ = toList p
        c = l * cos a
        s = l * sin a

flakeTransforms :: [R] -> [(R, M)]
flakeTransforms [a,b,c,d] = [(la,inv ta),(lb,inv tb),(lc,inv tc1),(ld,inv td1),(lc,inv tc2),(ld,inv td2)]
  where
    lx = a + b
    ly = 2 * (c + d)
    la  = a / lx
    lb  = b / lx
    lc  = c / ly
    ld  = d / ly
    u   = 2|> [0,  0]
    v   = 2|> [0, la]
    l   =     pi / 3
    ll  = 2 * pi / 3
    r   = -l
    rr  = -ll
    ta  = transformRST  0 la u
    tb  = transformRST  0 lb v
    tc1 = transformRST  l lc v
    tc2 = transformRST  r lc v
    td1 = transformRST ll ld v
    td2 = transformRST rr ld v
flakeTransforms _ = error "flakeTransforms"

main :: IO ()
main = do
  _ <- getArgsAndInitialize
  let wSize' = Size 1280 720
  initialWindowSize $= wSize'
  initialDisplayMode $= [DoubleBuffered]
  _ <- createWindow "SnowGlobe"
  pInitial'     <- shader Nothing (Just $ unlines
    [ "uniform float er;"
    , "uniform float rho;"
    , ""
    , "void main() {"
    , "  vec2 p = er * (gl_TexCoord[0].xy * 2.0 - vec2(1.0));"
    , "  float l = length(p);"
    , "  float n;"
    , "  if (l >= er) {"
    , "    n = 0.0;"
    , "  } else if (er > l && l >= rho * er) {"
    , "    n = (log(er) - log(l)) / -log(rho);"
    , "  } else {"
    , "    n = -1.0;"
    , "  }"
    , "  gl_FragData[0] = vec4(n);"
    , "}"
    ])
  pInitial'er'  <- get $ uniformLocation pInitial' "er"
  pInitial'rho' <- get $ uniformLocation pInitial' "rho"
  pStep'        <- shader Nothing (Just $ unlines
    [ "uniform float er;"
    , "uniform mat3 ts[6];"
    , ""
    , "uniform sampler2D src;"
    , ""
    , "void main() {"
    , "  vec2 p0 = er * (gl_TexCoord[0].xy * 2.0 - vec2(1.0));"
    , "  float m = -1.0;"
    , "  for (int i = 0; i < 6; ++i) {"
    , "    vec3 p = ts[i] * vec3(p0, 1.0);"
    , "    vec2 q = p.xy / p.z;"
    , "    float l = length(q);"
    , "    if (l < er) {"
    , "      m = max(m, texture2D(src, (q / er + vec2(1.0)) / 2.0).x);"
    , "    }"
    , "  }"
    , "  if (m >= 0.0) {"
    , "    m += 1.0;"
    , "  }"
    , "  m = max(m, texture2D(src, (p0 / er + vec2(1.0)) / 2.0).x);"
    , "  gl_FragData[0] = vec4(m);"
    , "}"
    ])
  pStep'er'     <- get $ uniformLocation pStep' "er"
  pStep'ts'     <- get $ uniformLocation pStep' "ts"
  pStep'src'    <- get $ uniformLocation pStep' "src"
  pColour'      <- shader Nothing (Just $ unlines
    [ "uniform sampler2D src;"
    , "uniform float speed;"
    , ""
    , "void main() {"
    , "  vec2 p = gl_TexCoord[0].xy;"
    , "  float n = texture2D(src, p).x;"
    , "  p -= vec2(0.5);"
    , "  const mat2 r = mat2(0.5, 0.8660254037844386, -0.8660254037844386, 0.5);"
    , "  for (int i = 1; i < 6; ++i) {"
    , "    p = r * p;"
    , "    n = max(n, texture2D(src, p + vec2(0.5)).x);"
    , "  }"
    , "  if (n > 0.0) {"
    , "    n *= speed;"
    , "  } else {"
    , "    n = 0.0;"
    , "  }"
    , "  n *= n;"
    , "  n *= n;"
    , "  gl_FragData[0] = vec4(vec3(1.0), n);"
    , "}"
    ])
  pColour'src'  <- get $ uniformLocation pColour' "src"
  pColour'speed'<- get $ uniformLocation pColour' "speed"
  tPing'        <- newTex size
  tPong'        <- newTex size
  fBuffer'      <- newFBO
  glClampColor gl_CLAMP_VERTEX_COLOR   gl_FALSE
  glClampColor gl_CLAMP_READ_COLOR     gl_FALSE
  glClampColor gl_CLAMP_FRAGMENT_COLOR gl_FALSE
  sR <- newIORef SnowGlobe
    { pInitial = pInitial', pInitial'er = pInitial'er', pInitial'rho = pInitial'rho'
    , pStep = pStep', pStep'er = pStep'er', pStep'ts = pStep'ts', pStep'src = pStep'src'
    , pColour = pColour', pColour'src = pColour'src', pColour'speed = pColour'speed'
    , tPing = tPing', tPong = tPong', fBuffer = fBuffer'
    , sFlakes = [], sTextures = M.empty, wSize = wSize', wFullScreen = Nothing
    , sNextName = 0, sRenderFlake = return undefined
    }
  modifyIORef sR $ \s'->s'{ sRenderFlake = flakeRenderStart sR }
  addTimerCallback 40 timer
  displayCallback $= display sR
  reshapeCallback $= Just (reshape sR)
  keyboardMouseCallback $= Just (keyboard sR)
  mainLoop

keyboard :: IORef SnowGlobe -> Key -> KeyState -> Modifiers -> Position -> IO ()
keyboard sR (SpecialKey KeyF11) Down _ _ = toggleFullScreen sR
keyboard sR (Char 'f') Down _ _ = toggleFullScreen sR
keyboard _ (Char _) Down _ _ = exitSuccess
keyboard _ _ _ _ _ = return ()

toggleFullScreen :: IORef SnowGlobe -> IO ()
toggleFullScreen sR = do
  s <- readIORef sR
  case wFullScreen s of
    Nothing -> do
      writeIORef sR s{ wFullScreen = Just (wSize s) }
      cursor $= None
      fullScreen
    Just ws -> do
      writeIORef sR s{ wFullScreen = Nothing }
      cursor $= Inherit
      windowSize $= ws

reshape :: IORef SnowGlobe -> Size -> IO ()
reshape sR sz = do
  s <- readIORef sR
  writeIORef sR s{ wSize = sz }

timer :: IO ()
timer = do
  addTimerCallback 40 timer
  postRedisplay Nothing

display' :: IORef SnowGlobe -> IO ()
display' sR = do
  update sR
  s <- readIORef sR
  _ <- evaluate (sum' . map flakeName . sFlakes $ s)
  let names = S.fromList $ map flakeName (sFlakes s)
      expired = S.filter (`S.notMember` names) (M.keysSet (sTextures s))
      sTextures' = foldr M.delete (sTextures s) (S.toList expired)
  deleteObjectNames [sTextures s M.! n | n <- S.toList expired]
  modifyIORef sR $ \s'->s'{ sTextures = sTextures' }
  r <- sRenderFlake s
  case r of
    Done t -> do
      f <- flakeSpawn (sFlakes s) (sNextName s)
      modifyIORef sR $ \s'->s'
        { sRenderFlake = flakeRenderStart sR
        , sFlakes = f : sFlakes s'
        , sTextures = M.insert (flakeName f) t (sTextures s')
        , sNextName = sNextName s' + 1
        }
    Step sRenderFlake' ->
      modifyIORef sR $ \s'->s'{ sRenderFlake = sRenderFlake' }

update :: IORef SnowGlobe -> IO ()
update sR = do
  s <- readIORef sR
  let sFlakes' = flakesUpdate (1 / 256) (sFlakes s)
  writeIORef sR (s{ sFlakes = sFlakes' })

display :: IORef SnowGlobe -> IO ()
display sR = do
  s <- readIORef sR
  loadIdentity
  let Size w h = wSize s
      r = 0.7
      (x, y)
        | h < w     = (r, r * fromIntegral h / fromIntegral w)
        | otherwise = (r * fromIntegral w / fromIntegral h, r)
  ortho2D (-x) x (-y) y
  viewport $= (Position 0 0, wSize s)
  clearColor $= Color4 0 0 0.25 1
  clear [ColorBuffer]
  texture Texture2D $= Enabled
  blend $= Enabled
  blendFunc $= (SrcAlpha, OneMinusSrcAlpha)
  forM_ (sFlakes s) (flakeDraw s)
  blend $= Disabled
  texture Texture2D $= Disabled
  swapBuffers
  reportErrors
  display' sR

flakeDraw :: SnowGlobe -> Flake -> IO ()
flakeDraw s f = do
  let d :: GLdouble
      d = realToFrac $ flakeSize f / 4
      a :: GLdouble
      a = realToFrac $ 360 * flakeTime f * sin (fromIntegral (flakeName f))
      p :: GLdouble -> GLdouble -> IO ()
      p u v = do
        texCoord $ TexCoord2 ((1+u)/2) ((1+v)/2)
        vertex $ Vertex2 u v
      x :: GLdouble
      y :: GLdouble
      [x, y] = map realToFrac . toList . flakePosition $ f
  case flakeName f `M.lookup` sTextures s of
    Nothing -> return ()
    t -> do
      textureBinding Texture2D $= t
      unsafePreservingMatrix $ do
        translate $ Vector3 x y 0
        rotate a (Vector3 0 0 1)
        GL.scale d d d
        renderPrimitive Quads $ p (-1) (-1) >> p 1 (-1) >> p 1 1 >> p (-1) 1
      textureBinding Texture2D $= Nothing

newTex :: N -> IO TextureObject
newTex s = do
  [t] <- genObjectNames 1
  textureBinding Texture2D $= Just t
  glTexImage2D gl_TEXTURE_2D 0 (fromIntegral gl_LUMINANCE32F) (fromIntegral s) (fromIntegral s) 0 gl_LUMINANCE gl_UNSIGNED_BYTE nullPtr
  textureFilter Texture2D $= ((Linear', Nothing), Linear')
  textureWrapMode Texture2D S $= (Repeated, ClampToEdge)
  textureWrapMode Texture2D T $= (Repeated, ClampToEdge)
  textureBinding Texture2D $= Nothing
  return t

newTexRGBA :: N -> IO TextureObject
newTexRGBA s = do
  [t] <- genObjectNames 1
  textureBinding Texture2D $= Just t
  glTexImage2D gl_TEXTURE_2D 0 (fromIntegral gl_RGBA) (fromIntegral s) (fromIntegral s) 0 gl_RGBA gl_UNSIGNED_BYTE nullPtr
  textureFilter Texture2D $= ((Linear', Just Linear'), Linear')
  textureWrapMode Texture2D S $= (Repeated, ClampToEdge)
  textureWrapMode Texture2D T $= (Repeated, ClampToEdge)
  textureBinding Texture2D $= Nothing
  return t

newtype FramebufferObject = FramebufferObject GLuint

newFBO :: IO FramebufferObject
newFBO = fmap FramebufferObject (alloca $ \p -> glGenFramebuffers 1 p >> peek p)

bindFBO :: FramebufferObject -> TextureObject -> IO ()
bindFBO (FramebufferObject f) (TextureObject t) = do
  glBindFramebuffer gl_FRAMEBUFFER f
  glFramebufferTexture2D gl_FRAMEBUFFER gl_COLOR_ATTACHMENT0 gl_TEXTURE_2D t 0

unbindFBO :: IO ()
unbindFBO = do
  glFramebufferTexture2D gl_FRAMEBUFFER gl_COLOR_ATTACHMENT0 gl_TEXTURE_2D 0 0
  glBindFramebuffer gl_FRAMEBUFFER 0

fullQuad :: IO ()
fullQuad = do
  renderPrimitive Quads $ do
    t (0.5-r) (0.5+r) >> v (-r2) ( r2)
    t (0.5-r) (0.5-r) >> v (-r2) (-r2)
    t (0.5+r) (0.5-r) >> v ( r2) (-r2)
    t (0.5+r) (0.5+r) >> v ( r2) ( r2)
  where
    r = sqrt 0.5 / realToFrac er
    r2 = 1
    t, v :: GLdouble -> GLdouble -> IO ()
    t x y = texCoord (TexCoord2 x y)
    v x y = vertex (Vertex2 x y)

unitQuad :: IO ()
unitQuad = renderPrimitive Quads $ do
    t 0 1 >> v 0 1
    t 0 0 >> v 0 0
    t 1 0 >> v 1 0
    t 1 1 >> v 1 1
  where
    t, v :: GLdouble -> GLdouble -> IO ()
    t x y = texCoord (TexCoord2 x y)
    v x y = vertex (Vertex2 x y)

clamp :: Ord a => a -> a -> a -> a
clamp mi ma x = mi `max` x `min` ma