wiistate.py

Go to the documentation of this file.

# ###############################################################################
#
# File:         wiistate.py
# RCS:          $Header: $
# Description:  Object representation of Wiimote's state
# Author:       Andreas Paepcke
# Created:      Thu Aug 13 11:01:34 2009 (Andreas Paepcke) paepcke@anw.willowgarage.com
# Modified:     Thu Jan 13 13:50:15 2011 (Andreas Paepcke) paepcke@bhb.willowgarage.com
# Language:     Python
# Package:      N/A
# Status:       Experimental (Do Not Distribute)
#
# ###############################################################################
#
# Revisions:
#
# Thu Mar 18 10:56:09 2010 (David Lu) davidlu@wustl.edu
#  Added nunchuk state variables
# Fri Oct 29 08:58:21 2010 (Miguel Angel Julian Aguilar, QBO Project) miguel.angel@thecorpora.com
#  Added classic controller state variables
# Mon Nov 08 11:43:23 2010 (David Lu) davidlu@wustl.edu
#  Added calibration for nunchuk
# ###############################################################################

from __future__ import absolute_import
from .wiimoteConstants import *
from .wiiutils import *
import numpy as np


class WIIState(object):
    """
    Holds the state of a WIIRemote-plus.

    The state is passed in and is as communicated
    by one message from the WII+ device. We unpack
    the information and place it into individual
    dictionaries for callers to grab.

    Public instance variables:
      o time             Time in fractional seconds since beginning of Epoch of when
                           state was measured (Float).
      o ascTime          Time when state was measured (Human-readable)
      o rumble           True/False if wiimote vibration is on/off
      o angleRate        A GyroReading instance containing gyro (a.k.a. angular rate) measurement
      o acc              A WIIReading instance containing accelerometer measurement corrected by
                           the calibration information that is stored in the Wiimote
      o accRaw           A WIIReading instance containing accelerometer measurement uncorrected
      o buttons          A dictionary for which buttons are being held down. That could be
                           multiple buttons. Keys are:
                                 BTN_1, BTN_2, BTN_PLUS, BTN_MINUS, BTN_A, BTN_B,
                                 BTN_UP, BTN_DOWN, BTN_LEFT, BTN_RIGHT, BTN_HOME
                           Values are 1/0
      o IRSources        Dictionary with on/off values for which IR lights are
                         being sensed. Keys are:
                                 IR1, IR2, IR3, IR4
                         Values are 1/0
      o motionPlusPresent True if a gyro Motion+ is plugged into the Wiimote. Else False

      o nunchukPresent   True if nunchuk is plugged in. Else False
      o nunchukAccRaw    A WIIReading instance with acceleromoter measurement from the nunchuk (raw values)
      o nunchukAcc       The same, but zeroed using factory calibration
      o nunchukStickRaw  A tuple with the two axes of the joystick on the nunchuk, raw readings
      o nunchukStick     A tuple with the two axes of the joystick on the nunchuk, zeroed to be [-1, 1]
      o nunchukButtons   A dictionary for which nunchuk buttons are down. Keys are BTN_C and BTN_Z

    Public methods:
      o setAccelerometerCalibration   Bias setting for accelerometer. This triplet is used to
                                        turn raw accelerometer values into calibrated values.
      o setGyroCalibration            Bias setting for gyro. This triplet is used to
                                        turn raw gyro values into calibrated values.
    """

    _accCalibrationZero = None
    _gyroZeroReading = None
    _nunchukZeroReading = None
    _nunchukJoystickZero = None

    def __init__(self, state, theTime, theRumble, buttonStatus):
        """Unpack the given state, normalizing if normalizers are passed in."""
        self.time = theTime
        self.ascTime = repr(theTime)
        self.rumble = theRumble
        self.IRSources = [None, None, None, None]
        self.battery = None
        self.acc = None
        self.accRaw = None
        self.angleRate = None
        self.angleRate = None
        self.angleRageRaw = None
        self.motionPlusPresent = False
        self.buttons = {BTN_1: False, BTN_2: False, BTN_PLUS: False,
                        BTN_MINUS: False, BTN_A: False, BTN_B: False,
                        BTN_UP: False, BTN_DOWN: False, BTN_LEFT: False,
                        BTN_RIGHT: False, BTN_HOME: False}
        self.nunchukPresent = False
        self.nunchukAccRaw = None
        self.nunchukAcc = None
        self.nunchukStick = None
        self.nunchukStickRaw = None
        self.nunchukButtons = {BTN_C: False, BTN_Z: False}

        self.classicPresent = False
        self.classicStickLeft = None
        self.classicStickRight = None
        self.classicButtons = {CLASSIC_BTN_A: False, CLASSIC_BTN_B: False,
                               CLASSIC_BTN_L: False, CLASSIC_BTN_R: False,
                               CLASSIC_BTN_X: False, CLASSIC_BTN_Y: False,
                               CLASSIC_BTN_ZL: False, CLASSIC_BTN_ZR: False,
                               CLASSIC_BTN_PLUS: False, CLASSIC_BTN_MINUS: False,
                               CLASSIC_BTN_UP: False, CLASSIC_BTN_DOWN: False,
                               CLASSIC_BTN_LEFT: False, CLASSIC_BTN_RIGHT: False,
                               CLASSIC_BTN_HOME: False}

        # Handle buttons on the WII
        # A zero means no button is down.

        if buttonStatus == 0:
            for key in self.buttons.keys():
                self.buttons[key] = False
                continue

        self.buttons[BTN_1] = (buttonStatus & BTN_1) > 0
        self.buttons[BTN_2] = (buttonStatus & BTN_2) > 0
        self.buttons[BTN_PLUS] = (buttonStatus & BTN_PLUS) > 0
        self.buttons[BTN_MINUS] = (buttonStatus & BTN_MINUS) > 0
        self.buttons[BTN_A] = (buttonStatus & BTN_A) > 0
        self.buttons[BTN_B] = (buttonStatus & BTN_B) > 0
        self.buttons[BTN_UP] = (buttonStatus & BTN_UP) > 0
        self.buttons[BTN_DOWN] = (buttonStatus & BTN_DOWN) > 0
        self.buttons[BTN_LEFT] = (buttonStatus & BTN_LEFT) > 0
        self.buttons[BTN_RIGHT] = (buttonStatus & BTN_RIGHT) > 0
        self.buttons[BTN_HOME] = (buttonStatus & BTN_HOME) > 0

        for msgComp in state:
            # msgComp has: (1,2) for Status:Button one pushed, or
            #              (3, [None,None,None,None]) for LEDs
            #              (7, {'angle_rage': (123,456,789)})
            msgType = msgComp[0]

            if msgType == WII_MSG_TYPE_ACC:
                # Second list member is accelerator triplet of numbers:
                accStatus = msgComp[1]
                self.accRaw = WIIReading(accStatus, self.time)

                # If this class knows about accelerometer calibration
                # data, correct the raw reading:
                if self._accCalibrationZero is not None and \
                   self._accCalibrationOne is not None and \
                   np.linalg.norm(self._accCalibrationOne - self._accCalibrationZero) > 1E-5:
                    self.acc = WIIReading((self.accRaw - self._accCalibrationZero) /
                                          (self._accCalibrationOne - self._accCalibrationZero), self.time)
                else:
                    self.acc = WIIReading(self.accRaw, self.time)

                continue

            elif msgType == WII_MSG_TYPE_IR:
                # Second list member is a list of 4 dictionaries,
                # one for each of the Wii IR sources:

                IRStatus = msgComp[1]
                # IRStatus is an array of four Dicts. Ex: [{'pos': (317, 445)}, None, None, None]
                self.IRSources[0] = IRStatus[0]
                self.IRSources[1] = IRStatus[1]
                self.IRSources[2] = IRStatus[2]
                self.IRSources[3] = IRStatus[3]

                continue

            elif msgType == WII_MSG_TYPE_MOTIONPLUS:
                # Second list member is a dictionary with the single
                # key 'angle_rate', which yields as its value a gyro
                # readings triplet of numbers:

                gyroDict = msgComp[1]

                if gyroDict is not None:
                    # If this class knows about a zero-movement reading for the
                    # gyro, subtract that reading from the raw measurement:

                    self.angleRateRaw = GyroReading(gyroDict['angle_rate'], self.time)
                    if self._gyroZeroReading is not None:
                        self.angleRate = GyroReading(self.angleRateRaw - self._gyroZeroReading, self.time)
                    else:
                        self.angleRate = self.angleRateRaw

                    self.motionPlusPresent = True

                continue
            elif msgType == WII_MSG_TYPE_NUNCHUK:
                nunChuk = msgComp[1]
                if nunChuk is not None:
                    self.nunchukPresent = True
                    self.nunchukAccRaw = WIIReading(nunChuk['acc'], self.time)

                    # If this class knows about accelerometer calibration
                    # data, correct the raw reading:
                    if self._nunchukZeroReading is not None:
                        self.nunchukAcc = WIIReading((self.nunchukAccRaw - self._nunchukZeroReading) /
                                                     (self._nunchukOneReading - self._nunchukZeroReading), self.time)
                    else:
                        self.nunchukAcc = WIIReading(self.nunchukAccRaw, self.time)

                    self.nunchukStickRaw = nunChuk['stick']

                    # scale the joystick to roughly [-1, 1]
                    if (self._nunchukJoystickZero is None):
                        calibration = [127, 127]
                    else:
                        calibration = self._nunchukJoystickZero

                    [joyx, joyy] = self.nunchukStickRaw
                    joyx = -(joyx-calibration[0])/100.
                    joyy = (joyy-calibration[1])/100.
                    # create a deadzone in the middle
                    if abs(joyx) < .05:
                        joyx = 0
                    if abs(joyy) < .05:
                        joyy = 0
                    self.nunchukStick = [joyx, joyy]

                    nunButtons = nunChuk['buttons']
                    self.nunchukButtons[BTN_C] = (nunButtons & BTN_C) > 0
                    self.nunchukButtons[BTN_Z] = (nunButtons & BTN_Z) > 0
                continue
            elif msgType == WII_MSG_TYPE_CLASSIC:
                clasSic = msgComp[1]
                if clasSic is not None:
                    self.classicPresent = True
                    self.classicStickLeft = clasSic['l_stick']
                    self.classicStickRight = clasSic['r_stick']
                    clasButtons = clasSic['buttons']
                    self.classicButtons[CLASSIC_BTN_A] = (clasButtons & CLASSIC_BTN_A) > 0
                    self.classicButtons[CLASSIC_BTN_B] = (clasButtons & CLASSIC_BTN_B) > 0
                    self.classicButtons[CLASSIC_BTN_DOWN] = (clasButtons & CLASSIC_BTN_DOWN) > 0
                    self.classicButtons[CLASSIC_BTN_HOME] = (clasButtons & CLASSIC_BTN_HOME) > 0
                    self.classicButtons[CLASSIC_BTN_L] = (clasButtons & CLASSIC_BTN_L) > 0
                    self.classicButtons[CLASSIC_BTN_LEFT] = (clasButtons & CLASSIC_BTN_LEFT) > 0
                    self.classicButtons[CLASSIC_BTN_MINUS] = (clasButtons & CLASSIC_BTN_MINUS) > 0
                    self.classicButtons[CLASSIC_BTN_PLUS] = (clasButtons & CLASSIC_BTN_PLUS) > 0
                    self.classicButtons[CLASSIC_BTN_R] = (clasButtons & CLASSIC_BTN_R) > 0
                    self.classicButtons[CLASSIC_BTN_RIGHT] = (clasButtons & CLASSIC_BTN_RIGHT) > 0
                    self.classicButtons[CLASSIC_BTN_UP] = (clasButtons & CLASSIC_BTN_UP) > 0
                    self.classicButtons[CLASSIC_BTN_X] = (clasButtons & CLASSIC_BTN_X) > 0
                    self.classicButtons[CLASSIC_BTN_Y] = (clasButtons & CLASSIC_BTN_Y) > 0
                    self.classicButtons[CLASSIC_BTN_ZL] = (clasButtons & CLASSIC_BTN_ZL) > 0
                    self.classicButtons[CLASSIC_BTN_ZR] = (clasButtons & CLASSIC_BTN_ZR) > 0
                continue

    @classmethod
    def setAccelerometerCalibration(cls, zeroReading, oneReading):
        """Set the current accelerometer zeroing calibration."""
        cls._accCalibrationZero = WIIReading(zeroReading)
        cls._accCalibrationOne = WIIReading(oneReading)

    @classmethod
    def getAccelerometerCalibration(cls):
        """Return current accelerometer zeroing offset as two lists of x/y/z: the
        zero-reading, and the one-reading."""
        return (cls._accCalibrationZero.tuple(), cls._accCalibrationOne.tuple())

    @classmethod
    def setGyroCalibration(cls, zeroReading):
        """Set the x/y/z zeroing offsets for the gyro. Argument is a list"""

        cls._gyroZeroReading = GyroReading(zeroReading)

    @classmethod
    def getGyroCalibration(cls):
        """Return current gyro zeroing offset as a list of x/y/z. """
        return cls._gyroZeroReading.tuple()

    @classmethod
    def setNunchukAccelerometerCalibration(cls, zeroReading, oneReading):
        """Set the current nunchuk accelerometer zeroing calibration."""
        cls._nunchukZeroReading = WIIReading(zeroReading)
        cls._nunchukOneReading = WIIReading(oneReading)

    @classmethod
    def setNunchukJoystickCalibration(cls, readings):
        """Set the origin for the nunchuk joystick"""
        cls._nunchukJoystickZero = readings

    @classmethod
    def getNunchukAccelerometerCalibration(cls):
        """Return current nunchuk accelerometer zeroing offset as two lists of x/y/z: the
        zero-reading, and the one-reading."""
        return (cls._nunchukZeroReading.tuple(), cls._nunchukOneReading.tuple())

    def __str__(self):

        # Timestamp:
        res = 'Time: ' + self.ascTime + '\n'

        # Buttons that are pressed:
        butRes = ''

        if self.buttons is not None:
            if self.buttons[BTN_1]:
                butRes += ', 1'
            if self.buttons[BTN_2]:
                butRes += ', 2'
            if self.buttons[BTN_PLUS]:
                butRes += ', Plus'
            if self.buttons[BTN_MINUS]:
                butRes += ', Minus'
            if self.buttons[BTN_A]:
                butRes += ', A'
            if self.buttons[BTN_B]:
                butRes += ', B'
            if self.buttons[BTN_UP]:
                butRes += ', 4Way-up'
            if self.buttons[BTN_DOWN]:
                butRes += ', 4Way-down'
            if self.buttons[BTN_LEFT]:
                butRes += ', 4Way-left'
            if self.buttons[BTN_RIGHT]:
                butRes += ', 4Way-right'
            if self.buttons[BTN_HOME]:
                butRes += ', Home'

        # If none of the buttons is down, indicate that:
        if not butRes:
            res += 'Buttons: none.\n'
        else:
            res += 'Buttons: ' + butRes.lstrip(', ') + '\n'

        # Accelerator:
        if self.acc is not None:
            res += 'Accelerator: (' + \
                   repr(self.acc[X]) + ',' + \
                   repr(self.acc[Y]) + ',' + \
                   repr(self.acc[Z]) + ')\n'

        # Gyro (angular rate):
        if self.angleRate is not None:
            res += 'Gyro (angular rate): (' + \
                   repr(self.angleRate[X]) + ',' + \
                   repr(self.angleRate[Y]) + ',' + \
                   repr(self.angleRate[Z]) + ')\n'

        # Rumble status:
        if self.rumble:
            res += 'Rumble: On.\n'
        else:
            res += 'Rumble: Off.\n'

        # IR Sources:
        irRes = ''

        if self.IRSources is not None:
            if self.IRSources[IR1] is not None:
                irRes += 'IR source 1'

            if self.IRSources[IR2] is not None:
                irRes += 'IR source 2'

            if self.IRSources[IR3] is not None:
                irRes += 'IR source 3'

            if self.IRSources[IR4] is not None:
                irRes += 'IR source 4'

            if not irRes:
                res += irRes.lstrip(', ') + '\n'
            else:
                res += 'No IR sources detected.\n'

        return res

    def __repr__(self):
        return self.__str__()


class WIIReading(object):
    """Instances hold one 3-D reading.

    Methods:
      [X], [Y], [Z] to obtain respective axis paramters.
      tuple() to obtain x/y/z as a NumPy array.
      +,-,/ to add or subtract readings from each other
          as one vector operation (pairwise for each dimension).
    """

    # Private instance vars:
    #     o _measurement = np.array(3, dtype=numpy.float64)
    #     o time

    def __init__(self, xyz, theTime=None):
        """Create a (possibly) time stamped WII Reading.

        Parameter xyz is an array of x,y,z coordinates of the
        reading. WIIReading instances can be added, subtracted, and
        divided into each other. The operations are pairwise over
        x, y, and z. A numpy array of x,y,z is available by
        calling tuple(). The time stamp is available via time().

        """
        self.time = theTime
        self._measurement = np.array([xyz[X], xyz[Y], xyz[Z]], dtype=np.float64)

    def __getitem__(self, key):
        if key not in (X, Y, Z):
            raise AttributeError("Attempt to index into a 3-D measurement array with index " + repr(key) + ".")
        return self._measurement[key]

        def __str__(self):
            return '[x=' + repr(self._measurement[X]) + \
                   ', y=' + repr(self._measurement[Y]) + \
                   ' z=' + repr(self._measurement[Z]) + \
                   ']'

    def __repr__(self):
        return '[' + str(self._measurement[X]) + ', ' + \
               str(self._measurement[Y]) + ', ' + str(self._measurement[Z]) + ']'

    def tuple(self):
        return self._measurement

    def __add__(self, other):
        """Adding two readings returns a numpy tuple with readings added pairwise."""

        return self._measurement + other._measurement

    def __sub__(self, other):
        """Subtracting two readings returns a numpy tuple with components subtracted pairwise."""

        return self._measurement - other._measurement

    def __div__(self, other):
        """Dividing two readings returns a numpy tuple with components divided pairwise."""

        return self._measurement / other._measurement

    def scale(self, scaleFactor):
        """Return a numpy tuple that with X, Y, Z scaled by the given factor."""

        return self._measurement * scaleFactor


class GyroReading():
    """Instances hold one gyroscope reading.

        Methods:
          [PHI], [THETA], [PSI] to obtain respective axis paramters.
          tuple() to obtain phi/theta/psi as a NumPy array.
          +,-,/ to add or subtract readings from each other
          as one vector operation (pairwise for each dimension).
    """
    # Local instance vars:
    #   o _measurement = np.array(3, dtype=numpy.float64)
    #   o time

    def __init__(self, phiThetaPsi, theTime=None):
        """Create a (possibly) time stamped WII Reading.

        Parameter phiThetaPsi is an array of phi,theta,psi coordinates of the
        gyro reading. GyroReading instances can be added, subtracted, and
        divided into each other. The operations are pairwise over
        phi, theta, and psi. A numpy array of phi,theta,psi is available by
        calling tuple(). The time stamp is available via time().
        """

        self.time = theTime
        self._measurement = np.array([phiThetaPsi[PHI], phiThetaPsi[THETA], phiThetaPsi[PSI]], dtype=np.float64)

    def __getitem__(self, key):
        if key not in (PHI, THETA, PSI):
            raise AttributeError("Attempt to index into a 3-D measurement array with index " + repr(key) + ".")
        return self._measurement[key]

    def __str__(self):
        return '[PHI (roll)=' + repr(self._measurement[PHI]) + \
               ', THETA (pitch)=' + repr(self._measurement[THETA]) + \
               ', PSI (yaw)=' + repr(self._measurement[PSI]) + \
               ']'

    def __repr__(self):
        '[' + str(self._measurement[PHI]) + ', ' + \
          str(self._measurement[THETA]) + ', ' + \
          str(self._measurement[PSI]) + ']'

    def tuple(self):
        return self._measurement

    def __add__(self, other):
        """Adding two gyro readings returns a new reading with components added pairwise."""
        return self._measurement + other._measurement

    def __sub__(self, other):
        """Subtracting two gyro readings returns a new reading
        with components subtracted pairwise.

        """
        return self._measurement - other._measurement

    def __div__(self, other):
        """Dividing two readings returns a numpy tuple with components divided pairwise."""

        return self._measurement / other._measurement

    def scale(self, scaleFactor):
        """Return a numpy tuple that with X, Y, Z scaled by the given factor."""

        return self._measurement * scaleFactor