Quadrocopter viewer

Main.hs

{-# LANGUAGE NamedFieldPuns, RecordWildCards #-}
--------------------------------------------------------------------------------
-- OpenGL quadrocopter viewer
--------------------------------------------------------------------------------
module Main where


import Data.IORef
import Control.Monad
import Control.Applicative
import Graphics.UI.GLUT

-- |Quadrocopter properties
data Quadrocopter
  = Quadrocopter { qPosition :: ( GLfloat, GLfloat, GLfloat )
                 , qRotation :: ( GLfloat, GLfloat, GLfloat )
                 , qSpeed :: ( GLfloat, GLfloat, GLfloat, GLfloat )
                 }
  deriving ( Eq, Ord, Show )

-- |Rotates around y
qRotateY :: GLfloat -> Quadrocopter -> Quadrocopter
qRotateY d q@Quadrocopter{..} = 
  let ( rx, ry, rz ) = qRotation
  in q{ qRotation = ( rx, ry + d, rz ) }

-- |Increases speed
qThrottle :: GLfloat -> Quadrocopter -> Quadrocopter
qThrottle d q@Quadrocopter{..} = 
  let ( s00, s01, s10, s11 ) = qSpeed
  in q{ qSpeed = ( s00 + d, s01 + d, s10 + d, s11 + d ) }


-- |Application state
data State 
  = State { sMouseX :: IORef ( GLint, GLint )
          , sSize :: IORef ( GLint, GLint )
          , sRotX :: IORef GLfloat
          , sRotY :: IORef GLfloat
          , sZoom :: IORef GLfloat
          , sPan :: IORef ( GLfloat, GLfloat, GLfloat )
          , sQuadro :: IORef Quadrocopter
          }

-- |Sets the vertex color
color3 :: GLfloat -> GLfloat -> GLfloat -> IO ()
color3 x y z
  = color $ Color4 x y z 1.0


-- |Sets the vertex position
vertex3 :: GLfloat -> GLfloat -> GLfloat -> IO ()
vertex3 x y z
  = vertex $ Vertex3 x y z


-- |Called when stuff needs to be drawn
display :: State -> DisplayCallback
display State{..} = do

  ( width, height ) <- get sSize
  rx                <- get sRotX
  ry                <- get sRotY
  z                 <- get sZoom
  ( tx, ty, tz )    <- get sPan
  quadro            <- get sQuadro

  clear [ ColorBuffer, DepthBuffer ]


  matrixMode $= Projection
  loadIdentity
  perspective 45.0 (fromIntegral width / fromIntegral height) 0.1 500.0


  matrixMode $= Modelview 0
  loadIdentity 
  translate $ Vector3 0 0 (-z * 10.0)
  rotate rx $ Vector3 1 0 0
  rotate ry $ Vector3 0 1 0
  translate $ Vector3 (-tx) (-ty) (-tz)
  
  -- |Draw reference system
  renderPrimitive Lines $ do
    color3 1.0 0.0 0.0
    vertex3 0.0 0.0 0.0
    vertex3 20.0 0.0 0.0

    color3 0.0 1.0 0.0
    vertex3 0.0 0.0 0.0
    vertex3 0.0 20.0 0.0

    color3 0.0 0.0 1.0
    vertex3 0.0 0.0 0.0
    vertex3 0.0 0.0 20.0

  preservingMatrix $ displayQuadrocopter quadro

  flush

-- |Renders the quadrocopter
displayQuadrocopter :: Quadrocopter -> IO ()
displayQuadrocopter Quadrocopter{..} = do
  let ( x, y, z )            = qPosition
      ( rx, ry, rz )         = qRotation
      ( s00, s01, s10, s11 ) = qSpeed

  translate $ Vector3 x y z
  rotate rx $ Vector3 1 0 0
  rotate ry $ Vector3 0 1 0
  rotate rz $ Vector3 0 0 1

  displayQuad 2 0.5 1.5

  preservingMatrix $ do
    translate $ Vector3 ( 2.0) 0.5 ( 1.5 :: GLfloat)
    displayQuad 0.4 0.4 0.4
    renderPrimitive Lines $ do
      color3 1.0 1.0 0.0
      vertex3 0.0 0.0 0.0
      vertex3 0.0 s00 0.0

  preservingMatrix $ do
    translate $ Vector3 ( 2.0) 0.5 (-1.5 :: GLfloat)
    displayQuad 0.4 0.4 0.4
    renderPrimitive Lines $ do
      color3 1.0 1.0 0.0
      vertex3 0.0 0.0 0.0
      vertex3 0.0 s01 0.0

  preservingMatrix $ do
    translate $ Vector3 (-2.0) 0.5 ( 1.5 :: GLfloat)
    displayQuad 0.4 0.4 0.4
    renderPrimitive Lines $ do
      color3 1.0 1.0 0.0
      vertex3 0.0 0.0 0.0
      vertex3 0.0 s10 0.0

  preservingMatrix $ do
    translate $ Vector3 (-2.0) 0.5 (-1.5 :: GLfloat)
    displayQuad 0.4 0.4 0.4
    renderPrimitive Lines $ do
      color3 1.0 1.0 0.0
      vertex3 0.0 0.0 0.0
      vertex3 0.0 s11 0.0


-- |Displays a quad
displayQuad :: GLfloat -> GLfloat -> GLfloat -> IO ()
displayQuad w h d = preservingMatrix $ do
  scale w h d
  renderPrimitive Quads $ do
    color3 1.0 0.0 0.0
    vertex3 (-1.0) ( 1.0) ( 1.0)
    vertex3 (-1.0) (-1.0) ( 1.0)
    vertex3 ( 1.0) (-1.0) ( 1.0)
    vertex3 ( 1.0) ( 1.0) ( 1.0)

    color3 1.0 0.0 0.0
    vertex3 (-1.0) (-1.0) (-1.0)
    vertex3 (-1.0) ( 1.0) (-1.0)
    vertex3 ( 1.0) ( 1.0) (-1.0)
    vertex3 ( 1.0) (-1.0) (-1.0)

    color3 0.0 1.0 0.0
    vertex3 ( 1.0) (-1.0) ( 1.0)
    vertex3 ( 1.0) (-1.0) (-1.0)
    vertex3 ( 1.0) ( 1.0) (-1.0)
    vertex3 ( 1.0) ( 1.0) ( 1.0)

    color3 0.0 1.0 0.0
    vertex3 (-1.0) (-1.0) (-1.0)
    vertex3 (-1.0) (-1.0) ( 1.0)
    vertex3 (-1.0) ( 1.0) ( 1.0)
    vertex3 (-1.0) ( 1.0) (-1.0)

    color3 0.0 0.0 1.0
    vertex3 (-1.0) (-1.0) ( 1.0)
    vertex3 (-1.0) (-1.0) (-1.0)
    vertex3 ( 1.0) (-1.0) (-1.0)
    vertex3 ( 1.0) (-1.0) ( 1.0)

    color3 0.0 0.0 1.0
    vertex3 (-1.0) ( 1.0) (-1.0)
    vertex3 (-1.0) ( 1.0) ( 1.0)
    vertex3 ( 1.0) ( 1.0) ( 1.0)
    vertex3 ( 1.0) ( 1.0) (-1.0)

-- |Called when the sSize of the viewport changes
reshape :: State -> ReshapeCallback
reshape State{..} (Size width height) = do
  sSize $= ( width, height )
  viewport $= (Position 0 0, Size width height)
  postRedisplay Nothing


-- |Animation
idle :: State -> IdleCallback
idle State{..} = do
  postRedisplay Nothing


-- |Mouse motion
motion :: State -> Position -> IO ()
motion State{..} (Position x y) = do
  ( mx, my ) <- get sMouseX

  sRotY $~! (+ fromIntegral ( fromIntegral x - mx ) )
  sRotX $~! (+ fromIntegral ( fromIntegral y - my ) )

  sMouseX $= ( x, y )


-- |Button input
input :: State -> Key -> KeyState -> Modifiers -> Position -> IO ()
input State{..} (MouseButton LeftButton) Down _ (Position x y)
  = sMouseX $= ( x, y )
input state (MouseButton WheelDown) Down _ pos
  = wheel state 0 120 pos
input state (MouseButton WheelUp) Down _ pos
  = wheel state 0 (-120) pos
input State{..} (SpecialKey key) Down _ _ =
  modifyIORef sQuadro $ case key of
    KeyLeft  -> qRotateY 3.0
    KeyRight -> qRotateY (-3.0)
    KeyUp    -> qThrottle 1.0
    KeyDown  -> qThrottle (-1.0)
input mxy _ _ _ _
  = return ()


-- |Mouse wheel movement (sZoom)
wheel :: State -> WheelNumber -> WheelDirection -> Position -> IO ()
wheel State{..} _num dir _pos
  | dir > 0   = get sZoom >>= (\x -> sZoom $= clamp (x + 0.1))
  | otherwise = get sZoom >>= (\x -> sZoom $= clamp (x - 0.1))
  where 
    clamp x = 0.5 `max` (30.0 `min` x)


-- |Main
main :: IO ()
main = do
  void $ getArgsAndInitialize >> createWindow "LSystems 3D viewer"

  -- Create a new state 
  state <- State <$> newIORef ( 0, 0 )
                 <*> newIORef ( 0, 1 )
                 <*> newIORef 0.0
                 <*> newIORef 0.0
                 <*> newIORef 5.0
                 <*> newIORef ( 0, 0, 0 )
                 <*> newIORef Quadrocopter { qPosition = ( 0.0, 0.0, 0.0 )
                                           , qRotation = ( 0.0, 0.0, 0.0)
                                           , qSpeed = ( 0.0, 0.0, 0.0, 0.0 )
                                           }

  -- OpenGL
  clearColor  $= Color4 0 0 0 1
  shadeModel  $= Smooth
  depthMask   $= Enabled
  depthFunc   $= Just Lequal
  lineWidth   $= 3.0

  -- Callbacks
  displayCallback       $= display state
  reshapeCallback       $= Just (reshape state)
  idleCallback          $= Just (idle state)
  mouseWheelCallback    $= Just (wheel state)
  motionCallback        $= Just (motion state)
  keyboardMouseCallback $= Just (input state)


  -- Let's get started
  mainLoop