Source

pywiimote / wiimote.py

#!/usr/bin/python
# -*- coding: utf-8 -*-
#-------------------------------python----------------------wiimote.py--
#                                                                       
#                             wiimote                                   
#                                                                       
# Copyright (c) 2010 J. Daniele Zambelli
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 3 as
# published by the Free Software Foundation
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
#

from __future__ import print_function, division
import cwiid, time
import numpy as np
import pdb  # for debug

class Wiimote(object):
    """A class for wiimote controll."""
    def __init__(self, acceleration=False, buttons=False, infrared=False):
        """
        Examples:
        ----------
        ::
            >>> import wiimote  # doctest: +SKIP
            >>> wiim = wiimote.Wiimote(acceleration=True,
            ...                        buttons=True) # doctest: +SKIP
            Press 1 & 2 on the Wiimote simultaneously to find it
            Wiimote activiated...       
        """
        self.conv = 9.81/25       # or 9.81/24
        print("Press 1 & 2 on the Wiimote simultaneously to find it")
        try:
            self.wiimote = cwiid.Wiimote()
        except RuntimeError:
            print("""
To find remote, type the following in a Shell:
$> hcitool scan
Then enter the address in the previus command.""")
#            raise RuntimeError('Device not detected') 
            exit(0)
        self.rpt_m = 0
        self.leds = 0
        self.led_on(1)
        self.rumble(0.5)
        self.setrpt(acceleration, buttons, infrared)
        self.acc_c = np.zeros(3)
        self.ir_c = np.array((0,0))
        
        print("Wiimote activiated...")
        #self.acc_c = self.acc_calibrate()

    def setrpt(self, acceleration=False, buttons=False, infrared=False):
        """Set which data will be in the state message,
        it light leds according to the report requiredset leds to

        """
        if acceleration:
            self.rpt_m |= cwiid.RPT_ACC
#            self.led_on(1)
#        else:
#            self.led_off(1)
        if buttons:
            self.rpt_m |= cwiid.RPT_BTN
#            self.led_on(2)
#        else:
#            self.led_off(2)
        if infrared:
            self.rpt_m |= cwiid.RPT_IR
#            self.led_on(3)
#        else:
#            self.led_off(3)
        self.wiimote.rpt_mode = self.rpt_m
#        self.wiimote.led = self.leds

    def led_on(self, led):
        """Turn on the led #led.
        Examples:
        ----------
        ::
            >>> wiim.led_on(2)
            
        """
        self.leds |= 1<<(led-1)
        self.wiimote.led = self.leds
        
    def led_off(self, led):
        """Turn off the led #led.
        Examples:
        ----------
        ::
            >>> wiim.led_off(2)
            
        """
        self.leds &= ~(1<<(led-1))
        self.wiimote.led = self.leds

    def rumble(self, secs=-1):
        """Rumble wiimote.
        Examples:
        ----------
        ::
            >>> wiim.rumble(1)
            
        """
        if secs == -1:
            self.wiimote.rumble = 1
        elif secs == 0:
            self.wiimote.rumble = 0
        else:
            self.wiimote.rumble = 1
            time.sleep(secs)
            self.wiimote.rumble = 0

    def acc_calibrate(self):
        """Accelerometer calibration.
        Examples:
        ----------
        ::
            >>> print(wiim.acc())       # doctest: +ELLIPSIS      
            [... ... ...]
            >>> wiim.acc_calibrate()    # doctest: +ELLIPSIS
            acc calibration ...
            >>> print(wiim.acc())       # doctest: +ELLIPSIS
            [... ... ...]
        """
        print("acc calibration ", end='')    # to eliminate
        rpt_m = self.wiimote.state['rpt_mode']
        rpt_m |= cwiid.RPT_ACC
        self.wiimote.rpt_mode = rpt_m
        acc = []
        for cnt in range(50):
            print('.', end='')          # to eliminate
            time.sleep(0.01)
            acc.append(self.wiimote.state['acc'])
        #pdb.set_trace()
        acc = np.array(acc)
        acc_c = acc.mean(axis=0)
        print("\nAcceleration calibrated: axc={0:5.2f}; ayc={1:5.2f}; "
              "azc={2:5.2f}".format(*acc_c))
        self.acc_c = acc_c
                
    def acc(self):
        """Return acceleration converted in m/s²

        ::
            >>> ax, ay, az =  wiim.acc() 
            >>> print("ax={0:f}; ay={0:f}; az={0:f}".format(ax, ay, az))
            ...       # doctest: +ELLIPSIS
            ax=...; ay=...; az=..."""
        #pdb.set_trace()
        acc = np.array(self.wiimote.state['acc'])
        return (acc-self.acc_c)*self.conv

    def ir_calibrate(self):
        """Infrared calibration.
        TODO:
            calibrare fornendo in sequenza il punto 0 e il punto 1
        Examples:
        ----------
        ::
            >>> print(wiim.ir()) # doctest: +ELLIPSIS
            [...]
            >>> wiim.ir_calibrate()   # doctest: +ELLIPSIS
            ir calibration ...
            >>> print(wiim.ir()) # doctest: +ELLIPSIS
            [...]
        """
        print("ir calibration   TODO", end='')    # to eliminate
#        self.wiimote.rpt_mode = self.wiimote.state['rpt_mode'] | cwiid.RPT_IR
#        ir = []
#        for cnt in range(50):
#            print('.', end='')          # to eliminate
#            time.sleep(0.01)
#            irs = self.wiimote.state['ir_src'][0]
#            if irs:
#                ir.append(irs['pos'])
##        pdb.set_trace()
#        if ir:
#            ir = np.array(ir)
#            ir_c =(np.array(ir.mean(axis=0), dtype='int'))
#            print("\nInfrared calibrated: irxc={0:d}; "
#                  "iryc={1:d}".format(*ir_c))
#            self.ir_c = ir_c
                
    def ir(self):
        """Return infrared list of points

        ::
            >>> print(wiim.ir()) # doctest: +ELLIPSIS
            [...]
        """
#        pdb.set_trace()
        irlist = self.wiimote.state['ir_src']
#        print('wiimote: ', irlist)
        return irlist
#        return [np.array(point['pos'])-self.ir_c for point in irlist if point]

    def close(self):
        self.wiimote.close()

if __name__ == '__main__':
    import doctest
    doctest.testmod(extraglobs={'wiim': Wiimote(acceleration=True,
                                                buttons=True,
                                                infrared=True)},
                    verbose = False)