-- SnowGlobe -- randomized fractal snowflakes demo
-- (GPL3+) 2012,2014,2015,2016 Claude Heiland-Allen <claude@mathr.co.uk>
-- tested with: ghc-7.6.3, ghc-7.8.2, ghc-7.10.1~rc2, ghc-8.0.1

import Vector hiding (i)
import Graphics.UI.GLUT hiding (scale, FramebufferObject)
import qualified Graphics.UI.GLUT as GL
import Graphics.GL
  ( glTexImage2D
  , glClampColor
  , glGenerateMipmap
  , glGenFramebuffers
  , glBindFramebuffer
  , glFramebufferTexture2D
  , glUniformMatrix3fv
  , glGetTexImage
  )
import Graphics.GL.Tokens
{-
  ( GL_TEXTURE_2D
  , GL_R32F
  , GL_RED
  , GL_RGBA
  , GL_UNSIGNED_BYTE
  , GL_FALSE
  , GL_CLAMP_VERTEX_COLOR
  , GL_CLAMP_READ_COLOR
  , GL_CLAMP_FRAGMENT_COLOR
  , GL_ALPHA
  , GL_FRAMEBUFFER
  , GL_COLOR_ATTACHMENT0
  )
-}
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, stdout, hPutStr, withBinaryFile, IOMode(WriteMode), hPutBuf)
import System.Random (randomRIO)
import Foreign (alloca, peek, nullPtr, withArray)
import Graphics.Rendering.OpenGL.Capture (capturePPM)
import Data.ByteString as BS (hPut, pack)
import Foreign (Ptr, allocaArray)
import Data.Word (Word8)

shader :: Maybe String -> Maybe String -> IO Program
shader mV mF = do
  p <- createProgram
  vs <- case mV of
    Nothing -> return []
    Just v -> do
      vert <- createShader VertexShader
      shaderSourceBS vert $= BS.pack (map (toEnum . fromEnum) 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 <- createShader FragmentShader
      shaderSourceBS frag $= BS.pack (map (toEnum . fromEnum) f)
      compileShader frag
      msg <- get (shaderInfoLog frag)
      when (not (null msg)) $ hPutStrLn stderr msg
      return [frag]
  attachedShaders p $= concat [vs, fs]
  linkProgram p
  msg <- get (programInfoLog p)
  when (not (null msg)) $ hPutStrLn stderr msg
  return p

type N = Int

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

size :: N
size = 2048

er :: R
er = 4

accuracy :: R
accuracy = 1 / fromIntegral size

tSize :: N
tSize = 1024

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, pColour'colour :: UniformLocation
  , tPing, tPong :: TextureObject, fBuffer :: FramebufferObject
  , sFlakes :: [Flake], sTextures :: Map N TextureObject
  , wSize :: Size, wFullScreen :: Maybe Size
  , sRenderFlake :: IO Render, sNextName :: N
  , sRecord, sRecordFlakes :: Bool
  }

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

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

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

flakeForce :: Flake -> V2 -> V2
flakeForce f p =
  let d@(V2 x y) = p ^-^ flakePosition f
      m = let d2 = x * x + y * y in d2 * d2
  in  (flakeMass f / m) *^ d

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

mapVector :: (R -> R) -> V2 -> V2
mapVector f (V2 x y) = V2 (f x) (f y)

flakeField :: [Flake] -> V2 -> V2
flakeField fs p = foldl' (^+^) o [ flakeForce f (p ^+^ d) | f <- fs, p /= flakePosition f, dx <- [-2,0,2], dy <- [-2,0,2], let d = V2 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 V2 xs ys
      p = fst . minimumBy (comparing snd) . map (\p' -> (p', norm2 (flakeField fs p'))) $ ps
  return Flake{ flakeName = name, flakeTime = 0, flakePosition = p, flakeVelocity = o }
  where
    norm2 (V2 x y) = x * x + y * y

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)

getUniformLocation :: UniformLocation -> GLint
getUniformLocation (UniformLocation u) = u

-- multi step passes
flakeRenderPass :: IORef SnowGlobe -> [M3] -> 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)) . concatMap matrixToList $ ts) $ glUniformMatrix3fv (getUniformLocation $ 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)

matrixToList :: M3 -> [R]
matrixToList (M3 a b c d e f g h i) = [a,b,c,d,e,f,g,h,i]

-- 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)
  uniform (pColour'colour s0) $= TexCoord3 1 1 (1 :: GLfloat)
  fullQuad
  currentProgram $= Nothing
  unbindFBO
  textureBinding Texture2D $= Just t
  glGenerateMipmap GL_TEXTURE_2D
  textureBinding Texture2D $= Nothing
  return $ Done t

transformRST :: R -> R -> V2 -> M3
transformRST a l (V2 x y) = M3 c s x  (-s) c y  0 0 1
  where c = l * cos a
        s = l * sin a

flakeTransforms :: [R] -> [(R, M3)]
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   = V2 0 0
    v   = V2 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;"
    , "uniform vec3 colour;"
    , ""
    , "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(colour, n);"
    , "}"
    ])
  pColour'src'  <- get $ uniformLocation pColour' "src"
  pColour'speed'<- get $ uniformLocation pColour' "speed"
  pColour'colour'<- get $ uniformLocation pColour' "colour"
  tPing'        <- newTex size
  tPong'        <- newTex size
  fBuffer'      <- newFBO
  glClampColor GL_CLAMP_VERTEX_COLOR   $ fromIntegral GL_FALSE
  glClampColor GL_CLAMP_READ_COLOR     $ fromIntegral GL_FALSE
  glClampColor GL_CLAMP_FRAGMENT_COLOR $ fromIntegral 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', pColour'colour = pColour'colour'
    , tPing = tPing', tPong = tPong', fBuffer = fBuffer'
    , sFlakes = [], sTextures = M.empty, wSize = wSize', wFullScreen = Nothing
    , sNextName = 0, sRenderFlake = return undefined
    , sRecord = False, sRecordFlakes = False
    }
  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 sR (Char 'R') Down _ _ = modifyIORef sR $ \s -> s{ sRecord = not (sRecord s) }
keyboard sR (Char 'S') Down _ _ = modifyIORef sR $ \s -> s{ sRecordFlakes = not (sRecordFlakes s) }
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
        }
      when (sRecordFlakes s) $
        saveTexture t ("snowglobe-" ++ show (flakeName f) ++ ".pgm")
    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
  when (sRecord s) $ hPut stdout =<< capturePPM
  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
      V2 x y = flakePosition f
  case flakeName f `M.lookup` sTextures s of
    Nothing -> return ()
    t -> do
      textureBinding Texture2D $= t
      unsafePreservingMatrix $ do
        translate $ Vector3 (realToFrac x :: GLdouble) (realToFrac y :: GLdouble) 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_R32F) (fromIntegral s) (fromIntegral s) 0 GL_RED 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

saveTexture :: TextureObject -> FilePath -> IO ()
saveTexture t f = withBinaryFile f WriteMode $ \h -> do
  let header = "P5\n" ++ show tSize ++ " " ++ show tSize ++ "\n255\n"
      n = tSize * tSize
  hPutStr h header
  allocaArray n $ \p -> do
    textureBinding Texture2D $= Just t
    glGetTexImage GL_TEXTURE_2D 0 GL_ALPHA GL_UNSIGNED_BYTE (p :: Ptr Word8)
    textureBinding Texture2D $= Nothing
    hPutBuf h p n

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