hw/examples/example02/Main.hs

{-# LANGUAGE OverloadedStrings, RecordWildCards #-}

module Main where

import Affection

import SDL (($=))
import qualified SDL

import qualified Graphics.Rendering.OpenGL as GL
import qualified Graphics.GLUtil as GLU

import Physics.Bullet.Raw

import Control.Concurrent.STM (atomically)
import Control.Concurrent.STM.TVar

import Control.Monad (when)
import Control.Monad.IO.Class (liftIO)

import Linear as L

import System.Random (randomRIO)

import SpatialMath

import Init
import Types

import Debug.Trace as T

main :: IO ()
main =
  withAffection (AffectionConfig
    { initComponents = All
    , windowTitle    = "hw - example 02"
    , windowConfigs  =
      [ ( 0
        , SDL.defaultWindow
          { SDL.windowInitialSize = SDL.V2 1920 1080
          , SDL.windowGraphicsContext = SDL.OpenGLContext SDL.defaultOpenGL
            { SDL.glProfile = SDL.Core SDL.Normal 3 2
            }
          }
        , SDL.FullscreenDesktop
        )
      ]
    } :: AffectionConfig StateData)

instance Affectionate StateData where
  preLoop = const (return ())
  handleEvents sd = mapM_(handle sd)
  update = Main.update
  draw = Main.draw
  loadState = load
  cleanUp = const (return ())
  hasNextStep = liftIO . readTVarIO . nextStep

quit = liftIO . atomically . flip writeTVar False . nextStep

update :: StateData -> Double -> Affection ()
update sd dt = do
  let g = 0.0667300
  (phys, physos) <- liftIO $ do
    p <- readTVarIO (physics sd)
    po <- readTVarIO (physicsObjects sd)
    return (p, po)
  mapM_ (\smallBall -> do
    ms1 <- liftIO $ getMotionState (bodyRigidBody smallBall)
    ms2 <- liftIO $ getMotionState (bodyRigidBody $ poBigBall physos)
    r1 <- liftIO $ return . fmap realToFrac =<< getPosition ms1
    r2 <- liftIO $ return . fmap realToFrac =<< getPosition ms2
    let m1 = bodyMass smallBall
        -- m2 = bodyMass (poBigBall physos)
        -- m2 = 1000000000000000
        m2 = 1000000
        eta_sq = 0.1 ^ 2
        force = (g * m2 * m1 *^ (r2 - r1)) ^/
          ((sqrt (((r2 - r1) `dot` (r2 - r1)) + eta_sq)) ^ 3)
    liftIO $ applyCentralForce (bodyRigidBody smallBall) force
    ) (poSmallBalls physos ++ poBigBalls physos)

  mapM_ (\(bb1, bb2) -> do
    ms1 <- liftIO $ getMotionState (bodyRigidBody bb1)
    ms2 <- liftIO $ getMotionState (bodyRigidBody bb2)
    r1 <- liftIO $ return . fmap realToFrac =<< getPosition ms1
    r2 <- liftIO $ return . fmap realToFrac =<< getPosition ms2
    let m1 = bodyMass bb1
        -- m2 = bodyMass (poBigBall physos)
        m2 = bodyMass bb2
        eta_sq = 0.1 ^ 2
        force = (g * m2 * m1 *^ (r2 - r1)) ^/
          ((sqrt (((r2 - r1) `dot` (r2 - r1)) + eta_sq)) ^ 3)
    liftIO $ applyCentralForce (bodyRigidBody bb1) force
    ) ((,) <$> (poBigBalls physos) <*> (poBigBalls physos))

  mapM_ (\(bb1, bb2) -> do
    ms1 <- liftIO $ getMotionState (bodyRigidBody bb1)
    ms2 <- liftIO $ getMotionState (bodyRigidBody bb2)
    r1 <- liftIO $ return . fmap realToFrac =<< getPosition ms1
    r2 <- liftIO $ return . fmap realToFrac =<< getPosition ms2
    let m1 = bodyMass bb1
        -- m2 = bodyMass (poBigBall physos)
        m2 = bodyMass bb2
        eta_sq = 0.1 ^ 2
        force = (g * m2 * m1 *^ (r2 - r1)) ^/
          ((sqrt (((r2 - r1) `dot` (r2 - r1)) + eta_sq)) ^ 3)
    liftIO $ applyCentralForce (bodyRigidBody bb1) force
    ) ((,) <$> (poSmallBalls physos) <*> (poBigBalls physos))

  liftIO $ stepSimulation (pWorld phys) dt 10 Nothing
  posrots <- mapM ((\ball -> do
    ms <- liftIO $ getMotionState ball
    npos <- liftIO $ return . fmap realToFrac =<< getPosition ms
    nrot <- liftIO $ return . fmap realToFrac =<< getRotation ms
    return (npos, nrot))
    . bodyRigidBody) (poSmallBalls physos)
  posrots2 <- mapM ((\ball -> do
    ms <- liftIO $ getMotionState ball
    npos <- liftIO $ return . fmap realToFrac =<< getPosition ms
    nrot <- liftIO $ return . fmap realToFrac =<< getRotation ms
    return (npos, nrot))
    . bodyRigidBody) (poBigBalls physos)
  liftIO $ atomically $ do
    modifyTVar (ships sd) $ \ships ->
      map (\(ship, (pos, rot)) ->
        ship
          { shipRot = rot
          , shipPos = pos
          }
        ) (zip ships posrots)
    modifyTVar (oplanets sd) $ \oplanets ->
      map (\(ball, (pos, rot)) ->
        ball
          { shipRot = rot
          , shipPos = pos
          }
        ) (zip oplanets posrots2)
  liftIO $ atomically $ do
    ind <- readTVar (focusIndex sd)
    nplanets <- readTVar (oplanets sd)
    planet <- readTVar (planet sd)
    cam <- readTVar (camera sd)
    writeTVar (camera sd) cam
      { cameraFocus = shipPos ((planet : nplanets) !! ind)
      }

draw :: StateData -> Affection ()
draw sd = do
  GL.viewport $= (GL.Position 0 0, GL.Size 1920 1080)
  (planet, oplanets, ships, program, program2) <- liftIO $ do
    p <- readTVarIO $ planet sd
    o <- readTVarIO $ oplanets sd
    s <- readTVarIO $ ships sd
    pr <- readTVarIO $ program sd
    pr2 <- readTVarIO $ program2 sd
    return (p, o, s, pr, pr2)
  drawThings program (planet : ships)
  -- drawThings program (ships)
  drawThings program2 oplanets
  where
  drawThings prog ts = do
    (camera, proj, program) <- liftIO $ do
      cam <- readTVarIO (camera sd)
      p   <- readTVarIO (proj sd)
      program <- readTVarIO (program sd)
      return (cam, p, program)
    GL.currentProgram $= (Just . GLU.program $ prog)
    mapM_ (\Ship{..} -> do
      let view = lookAt
            (cameraFocus camera +
              rotVecByEulerB2A
              (cameraRot camera)
              (V3 0 0 (-cameraDist camera)))
            (cameraFocus camera)
            (V3 0 1 0)
          model = mkTransformation shipRot shipPos
          pvm = proj !*! view !*! model
      liftIO $ GLU.setUniform program "mvp" pvm
      GL.bindVertexArrayObject $= Just shipVao
      liftIO $ GL.drawArrays GL.Triangles 0 (fromIntegral shipVaoLen)
      ) ts

handle :: StateData -> SDL.EventPayload -> Affection ()
handle sd (SDL.WindowClosedEvent _) = quit sd

handle sd (SDL.KeyboardEvent dat) = do
  let key = SDL.keysymKeycode (SDL.keyboardEventKeysym dat)
  when (SDL.keyboardEventKeyMotion dat == SDL.Pressed) $
    handleKey sd key
handle sd (SDL.MouseMotionEvent dat) = do
  let (V2 rx ry) = fromIntegral <$> SDL.mouseMotionEventRelMotion dat
  liftIO $ atomically $ modifyTVar (camera sd) $ \c ->
    case SDL.mouseMotionEventState dat of
      -- [SDL.ButtonRight] ->
      --   let (V3 sx sy sz) = rotVecByEuler (cameraRot c) (V3 (rx / 10) 0 (ry / 10))
      --   in  c {cameraFocus = cameraFocus c + V3 sx 0 sy}
      [] ->
        let dphi = pi / 4 / 45 / 10
            (Euler yaw pitch roll) = cameraRot c
            nangle
              | nangle' >= qc = qc - mu
              | nangle' <= -qc = -qc + mu
              | otherwise     = nangle'
              where
                nangle' = (dphi * ry) + roll
                qc = pi / 2
                mu = 0.01
            nrot =
              Euler
                yaw
                (pitch + (rx * dphi))
                nangle
        in  c
          { cameraRot = nrot
          }
      _ ->
        c

handle _ _ = return ()

handleKey :: StateData -> SDL.Keycode -> Affection ()
handleKey sd code
  | code == SDL.KeycodeTab = do
      ps <- liftIO ((:) <$> readTVarIO (planet sd) <*> readTVarIO (oplanets sd))
      liftIO $ atomically $ modifyTVar (focusIndex sd) $ \ind ->
        if ind + 1 < length ps
        then ind + 1
        else 0
  | code == SDL.KeycodeR =
    GL.clearColor $= GL.Color4 1 0 0 1
  | code == SDL.KeycodeG =
    GL.clearColor $= GL.Color4 0 1 0 1
  | code == SDL.KeycodeB =
    GL.clearColor $= GL.Color4 0 0 1 1
  | code == SDL.KeycodeP = do
    r <- liftIO $ randomRIO (0, 1)
    g <- liftIO $ randomRIO (0, 1)
    b <- liftIO $ randomRIO (0, 1)
    a <- liftIO $ randomRIO (0, 1)
    GL.clearColor $= GL.Color4 r g b a
  | code == SDL.KeycodeEscape =
    quit sd
  | code == SDL.KeycodeF = do
    dt <- deltaTime <$> get
    liftIO $ putStrLn $ show (1 / dt) ++ " FPS"
  | code == SDL.KeycodeT =
    toggleScreen 0
  | code `elem`
    [ SDL.KeycodeW
    , SDL.KeycodeS
    , SDL.KeycodeA
    , SDL.KeycodeD
    , SDL.KeycodeQ
    , SDL.KeycodeE
    ]
    = do
      ship <- head <$> liftIO (readTVarIO $ ships sd)
      let rot = shipRot ship
          dphi = pi / 2 / 45
          nquat = case code of
            SDL.KeycodeW -> rot * axisAngle (V3 1 0 0) (-dphi)
            SDL.KeycodeS -> rot * axisAngle (V3 1 0 0) dphi
            SDL.KeycodeA -> rot * axisAngle (V3 0 1 0) (-dphi)
            SDL.KeycodeD -> rot * axisAngle (V3 0 1 0) dphi
            SDL.KeycodeE -> rot * axisAngle (V3 0 0 1) (-dphi)
            SDL.KeycodeQ -> rot * axisAngle (V3 0 0 1) dphi
            _            -> rot
      liftIO $ atomically $ modifyTVar (ships sd) $ \ships ->
        ship
          { shipRot = nquat
          } : tail ships
  | otherwise =
    return ()