TrackIR.hs

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}

import Network.Socket hiding (recv, sendAll, send)
import Network.Socket.ByteString (recv, sendAll, send)
import qualified Data.ByteString as B
import Control.Applicative

import qualified Graphics.Win32

import Control.Exception (SomeException, catch)
import System.Exit (ExitCode(ExitSuccess), exitWith)
import Control.Monad
import Control.Concurrent
import Control.Concurrent.STM
import Control.Concurrent.Async

import System.CPUTime


import Foreign
import Data.IORef
import Data.Word
import Data.Serialize
import Linear
import WinHook

data PoseInfo = PoseInfo { deviceId :: Word8
                         , bitflag :: Word8
                         , accelr :: V3 Float
                         , gyro :: V3 Float
                         , magneto :: V3 Float
                         , pose :: V3 Float
                           } deriving Show

getV3Float :: Get (V3 Float)
getV3Float = V3 <$> getFloat32le <*> getFloat32le <*> getFloat32le

putV3Float :: (V3 Float) -> Put
putV3Float (V3 a b c) = do
  putFloat32le a
  putFloat32le b
  putFloat32le c

getPoseInfo :: Get PoseInfo
getPoseInfo = PoseInfo
              <$> getWord8
              <*> getWord8
              <*> getV3Float
              <*> getV3Float
              <*> getV3Float
              <*> getV3Float

putPoseInfo (PoseInfo a b c d e f) = do
  putWord8 a
  putWord8 b
  putV3Float c
  putV3Float d
  putV3Float e
  putV3Float f


modifyReadTVar var f = do
  x <- readTVar var
  let x' = f x
  writeTVar var x'
  return x'


eulerToQ (V3 p r y) = Quaternion w (V3 x y' z)
                      where
                        w  = c1 * c2 * c3 + s1 * s2 * s3
                        x  = s1 * s2 * c3 + c1 * c2 * s3
                        y' = s1 * c2 * c3 + c1 * s2 * s3
                        z  = c1 * s2 * c3 - s1 * c2 * s3
                        c1 = cos $ y / 2
                        c2 = cos $ p / 2
                        c3 = cos $ r / 2
                        s1 = sin $ y / 2
                        s2 = sin $ p / 2
                        s3 = sin $ r / 2
                        

qToEuler (Quaternion w (V3 x y z)) = V3 (y') (p) (r)
                                     where
                                       y' = atan2 ( 2 * y * w - 2 * x * z) (1 - 2 *y^2 - 2* z^2)
                                       p = asin (2 * x * y + 2 * z * w)
                                       r = atan2 ( 2* x * w - 2 * y * z) ( 1 -2 * x^2 - 2 * z^2)

serverLoop sockS valQ = do
  msg <- recv sockS 512
  case runGet getPoseInfo msg of
   Right poseI -> do
     atomically . writeTQueue valQ . gyro $ poseI
   Left _ -> print "fail"
  serverLoop sockS valQ

sendLoop sock var = do
  t <- async (threadDelay 8333)
  val <- atomically $ readTVar var
  sendAll sock . runPut . putPoseInfo $ PoseInfo 0 3 (V3 0 0 0) (V3 0 0 0) (V3 0 0 0) (V3 2 2 0.2 * qToEuler val)
  wait t
  sendLoop sock var


procLoop valQ state offset ts@(t1: t2s) = do
  t2 <- getCPUTime
  let td =  (1/57.2958)  :: Float
  a <- atomically $ readTQueue valQ
  atomically $ modifyTVar' state (\q -> q * eulerToQ ((a-offset)^*td))
  procLoop valQ state offset (t2s ++ [t2])

  where at = (fromIntegral $ sum ts) / (fromIntegral $ length ts) / (10e11)

setupNetServer = do
  addrsinfo <- getAddrInfo
               (Just (defaultHints {addrFlags = [AI_PASSIVE]}))
               Nothing
               (Just "5555")
  let addr = head addrsinfo
  sock <- socket (addrFamily addr) Datagram defaultProtocol

  bindSocket sock (addrAddress addr)
  return sock


setupNetClient = do
  addrsinfo <- getAddrInfo
               Nothing
               (Just "127.0.0.1")
               (Just "5556")
  let addr = head addrsinfo
  sock <- socket (addrFamily addr) Datagram defaultProtocol
  connect sock (addrAddress addr)
  return sock


msgPump :: IO ()
msgPump = Graphics.Win32.allocaMessage $ \ msg ->
  let pump = do
        Graphics.Win32.getMessage msg Nothing
          `catch` \(_::SomeException) -> exitWith ExitSuccess
        Graphics.Win32.translateMessage msg
        Graphics.Win32.dispatchMessage msg
        pump
  in pump

keyHook tvar ncode wp lp = do
  KBDLLHOOKSTRUCT vk scan flags time _ <- peek lp

  case wp of
   0x0100 -> return ()
   0x0101 -> atomically $ writeTQueue tvar vk
   
  callNextHookEx ncode wp lp >>= return

main :: IO ()
main = withSocketsDo $ do
  sockS <- setupNetServer
  sockC <- setupNetClient
  
  valuesQ <- atomically $ newTQueue
  forkIO $ serverLoop sockS valuesQ
  
  state <- atomically $ newTVar (axisAngle (V3 0 0 0) 0)
  forkIO $ sendLoop sockC state

  {-keys <- atomically $ newTQueue
  setWindowsHookEx wH_KEYBOARD_LL (keyHook keys) Nothing 0
  --msgPump-}
  
  
  --initstate <- mapM (\x -> atomically $ readTQueue valuesQ) [0..9999]

  --let offset = ( (sum initstate) / 10000)

  let offset = V3 0 0 0
  print offset

  putStrLn "calibration complete"
                          
  
  {-forkIO . forever $ do
    x <- atomically $ readTQueue keys
    print $ x-}
  
  forkIO $ procLoop valuesQ state offset [0, 0, 0, 0]
  --msgPump
  forever $ do
    _ <- getLine
    atomically $ writeTVar state (axisAngle (V3 0 0 0) 0)
  

WinHook.hs

{-# LINE 1 "WinHook.hsc" #-}
module WinHook where
{-# LINE 2 "WinHook.hsc" #-}
import System.Win32.Types (maybePtr, DWORD, WPARAM, LPARAM, LRESULT, HANDLE, HINSTANCE, MbHINSTANCE)
import Foreign


{-# LINE 6 "WinHook.hsc" #-}

type KBHookProc = Int -> WPARAM -> Ptr KBDLLHOOKSTRUCT -> IO LRESULT

data KBDLLHOOKSTRUCT = KBDLLHOOKSTRUCT { vkCode :: DWORD
                                       , scanCode :: DWORD
                                       , flags :: DWORD
                                       , time :: DWORD
                                       , dwExtraInfo :: DWORD } -- don't use this.


instance Storable KBDLLHOOKSTRUCT where
  sizeOf _    = ((24))
{-# LINE 18 "WinHook.hsc" #-}
  alignment _ = alignment (undefined :: DWORD)
  peek ptr = do
    vkcode <- ((\hsc_ptr -> peekByteOff hsc_ptr 0)) ptr
{-# LINE 21 "WinHook.hsc" #-}
    scancode <- ((\hsc_ptr -> peekByteOff hsc_ptr 4)) ptr
{-# LINE 22 "WinHook.hsc" #-}
    flags <- ((\hsc_ptr -> peekByteOff hsc_ptr 8)) ptr
{-# LINE 23 "WinHook.hsc" #-}
    time <- ((\hsc_ptr -> peekByteOff hsc_ptr 12)) ptr
{-# LINE 24 "WinHook.hsc" #-}
    extra <- ((\hsc_ptr -> peekByteOff hsc_ptr 16)) ptr
{-# LINE 25 "WinHook.hsc" #-}
    return KBDLLHOOKSTRUCT { vkCode = vkcode
                           , scanCode = scancode
                           , flags = flags
                           , time = time
                           , dwExtraInfo = extra }
  poke ptr (KBDLLHOOKSTRUCT vk scan flags time extra) = do
    ((\hsc_ptr -> pokeByteOff hsc_ptr 0)) ptr vk
{-# LINE 32 "WinHook.hsc" #-}
    ((\hsc_ptr -> pokeByteOff hsc_ptr 4)) ptr scan
{-# LINE 33 "WinHook.hsc" #-}
    ((\hsc_ptr -> pokeByteOff hsc_ptr 8)) ptr flags
{-# LINE 34 "WinHook.hsc" #-}
    ((\hsc_ptr -> pokeByteOff hsc_ptr 12)) ptr time
{-# LINE 35 "WinHook.hsc" #-}
    ((\hsc_ptr -> pokeByteOff hsc_ptr 16)) ptr extra
{-# LINE 36 "WinHook.hsc" #-}
                       
wH_MOUSE_LL = 14 :: Int
wH_KEYBOARD_LL = 13 :: Int

foreign import ccall "wrapper"
  mkKBHookClosure :: KBHookProc -> IO (FunPtr KBHookProc)

setWindowsHookEx :: Int -> KBHookProc -> MbHINSTANCE -> DWORD -> IO ()
setWindowsHookEx nCode hookFunc hMod t = do
  fn <- mkKBHookClosure hookFunc
  c_setWindowsHookEx nCode fn (maybePtr hMod) t

foreign import ccall unsafe "user32.h SetWindowsHookExW"
  c_setWindowsHookEx :: Int -> FunPtr KBHookProc -> HINSTANCE -> DWORD -> IO ()


callNextHookEx :: Int -> WPARAM -> Ptr KBDLLHOOKSTRUCT -> IO LRESULT
callNextHookEx nCode wp lp =
  c_callNextHookEx 0 nCode wp lp

foreign import ccall unsafe "user32.h CallNextHookEx"
  c_callNextHookEx :: Int -> Int -> WPARAM -> Ptr KBDLLHOOKSTRUCT -> IO LRESULT