#!/usr/bin/env python

# BB-8 Python driver by Alistair Buxton <a.j.buxton@gmail.com>

from bluepy import btle

import time
import pygame

class BB8(btle.DefaultDelegate):
    def __init__(self, deviceAddress):
        btle.DefaultDelegate.__init__(self)

        # Address type must be "random" or it won't connect.
        self.peripheral = btle.Peripheral(deviceAddress, btle.ADDR_TYPE_RANDOM)
        self.peripheral.setDelegate(self)

        self.seq = 0

        # Attribute UUIDs are identical to Ollie.
        self.antidos = self.getSpheroCharacteristic('2bbd')
        self.wakecpu = self.getSpheroCharacteristic('2bbf')
        self.txpower = self.getSpheroCharacteristic('2bb2')
        self.roll = self.getSpheroCharacteristic('2ba1')
        self.notify = self.getSpheroCharacteristic('2ba6')

        # This startup sequence is also identical to the one for Ollie.
        # It even uses the same unlock code.
        print('Sending antidos')
        self.antidos.write('011i3', withResponse=True)
        print('Sending txpower')
        self.txpower.write('\x0007', withResponse=True)
        print('Sending wakecpu')
        self.wakecpu.write('\x01', withResponse=True)

    def getSpheroCharacteristic(self, fragment):
        return self.peripheral.getCharacteristics(uuid='22bb746f'+fragment+'75542d6f726568705327')[0]

    def dumpCharacteristics(self):
        for s in self.peripheral.getServices():
            print(s)
            for c in s.getCharacteristics():
                print( c, hex(c.handle))

    def cmd(self, did, cid, data=[], answer=True, resetTimeout=True):
        # Commands are as specified in Sphero API 1.50 PDF.
        # https://github.com/orbotix/DeveloperResources/
        seq = (self.seq&255)
        self.seq += 1
        sop2 = 0xfc
        sop2 |= 1 if answer else 0
        sop2 |= 2 if resetTimeout else 0
        dlen = len(data)+1
        chk = (sum(data)+did+cid+seq+dlen)&255
        chk ^= 255

        msg = [0xff, sop2, did, cid, seq, dlen] + data + [chk]
        print('cmd:', ' '.join([chr(c).encode('hex') for c in msg]))
        # Note: withResponse is very important. Most commands won't work without it.
        self.roll.write(''.join([chr(c) for c in msg]), withResponse=True)

    def handleNotification(self, cHandle, data):
        print( 'Notification:', cHandle, data.encode('hex'))
        # print( 'Notification:', cHandle)

    def waitForNotifications(self, time):
        self.peripheral.waitForNotifications(time)

    def disconnect(self):
        self.peripheral.disconnect()



if __name__ == '__main__':

    pygame.display.set_mode((320, 240))
    c = pygame.time.Clock()

    # Connect by address. Use "sudo hcitool lescan" to find address.
    bb = BB8('C0:D0:AC:38:8B:F4')

    # Dump all GATT stuff.
    #bb.dumpCharacteristics()

    # Request some sensor stream.
    #bb.cmd(0x02, 0x11, [0, 80, 0, 1, 0x80, 0, 0, 0,   0])

    #for i in range(255):
        # Set RGB LED colour.
    #    bb.cmd(0x02, 0x20, [254, i, 2, 0])
        # Wait for streamed data.
    #    bb.waitForNotifications(1.0)

    # Ask for battery level
    #bb.cmd(0x00, 0x20, [])
        # Wait for streamed data.
    #bb.waitForNotifications(1.0)

    # constant stream, so repeat command
    several_times = 25
    # move straight
    h = 0 

    v = 20 # speed up
    bb.cmd(0x02, 0x30, [v, (h&0xff00)>>8, h&0xff, 1]) 
    v = 40 # speed up
    bb.cmd(0x02, 0x30, [v, (h&0xff00)>>8, h&0xff, 1]) 
    v = 60 # speed up
    bb.cmd(0x02, 0x30, [v, (h&0xff00)>>8, h&0xff, 1]) 
    v = 80 # speed up
    bb.cmd(0x02, 0x30, [v, (h&0xff00)>>8, h&0xff, 1]) 

    for t in range(several_times): # timesteps
        v = 100
        h = 0 
        bb.cmd(0x02, 0x30, [v, (h&0xff00)>>8, h&0xff, 1])

    v = 80 # speed down
    bb.cmd(0x02, 0x30, [v, (h&0xff00)>>8, h&0xff, 1]) 
    v = 60 # speed down
    bb.cmd(0x02, 0x30, [v, (h&0xff00)>>8, h&0xff, 1]) 
    v = 20 # speed down
    h += 30 # start preparing next turn
    bb.cmd(0x02, 0x30, [v, (h&0xff00)>>8, h&0xff, 1]) 
    v = 0 # speed down
    h +=30  # start preparing turn
    bb.cmd(0x02, 0x30, [v, (h&0xff00)>>8, h&0xff, 1]) 

    # turn right
    h = 90 

    v = 20 # speed up
    bb.cmd(0x02, 0x30, [v, (h&0xff00)>>8, h&0xff, 1])
    v = 40 # speed up
    bb.cmd(0x02, 0x30, [v, (h&0xff00)>>8, h&0xff, 1])
    v = 60 # speed up
    bb.cmd(0x02, 0x30, [v, (h&0xff00)>>8, h&0xff, 1]) 
    v = 80 # speed up
    bb.cmd(0x02, 0x30, [v, (h&0xff00)>>8, h&0xff, 1]) 

    for t in range(several_times): # timesteps
        h = 90 
        v = 100
        bb.cmd(0x02, 0x30, [v, (h&0xff00)>>8, h&0xff, 1]) 

    v = 80 # speed down
    bb.cmd(0x02, 0x30, [v, (h&0xff00)>>8, h&0xff, 1]) 
    v = 60 # speed down
    bb.cmd(0x02, 0x30, [v, (h&0xff00)>>8, h&0xff, 1]) 
    v = 20 # speed down
    h +=30  # start preparing turn
    bb.cmd(0x02, 0x30, [v, (h&0xff00)>>8, h&0xff, 1]) 
    v = 0 # speed down
    h +=30  # start preparing turn
    bb.cmd(0x02, 0x30, [v, (h&0xff00)>>8, h&0xff, 1]) 


    # turn back
    h = 180 

    v = 20 # speed up
    bb.cmd(0x02, 0x30, [v, (h&0xff00)>>8, h&0xff, 1])
    v = 40 # speed up
    bb.cmd(0x02, 0x30, [v, (h&0xff00)>>8, h&0xff, 1])
    v = 60 # speed up
    bb.cmd(0x02, 0x30, [v, (h&0xff00)>>8, h&0xff, 1]) 
    v = 80 # speed up
    bb.cmd(0x02, 0x30, [v, (h&0xff00)>>8, h&0xff, 1]) 

    for t in range(several_times): # timesteps
        v = 100
        h = 180 
        bb.cmd(0x02, 0x30, [v, (h&0xff00)>>8, h&0xff, 1])

    v = 80 # speed down
    bb.cmd(0x02, 0x30, [v, (h&0xff00)>>8, h&0xff, 1]) 
    v = 60 # speed down
    bb.cmd(0x02, 0x30, [v, (h&0xff00)>>8, h&0xff, 1]) 
    v = 20 # speed down
    h +=30  # start preparing turn
    bb.cmd(0x02, 0x30, [v, (h&0xff00)>>8, h&0xff, 1]) 
    v = 0 # speed down
    h +=30  # start preparing turn
    bb.cmd(0x02, 0x30, [v, (h&0xff00)>>8, h&0xff, 1]) 


    # turn left
    h = 270 

    v = 20 # speed up
    bb.cmd(0x02, 0x30, [v, (h&0xff00)>>8, h&0xff, 1])
    v = 40 # speed up
    bb.cmd(0x02, 0x30, [v, (h&0xff00)>>8, h&0xff, 1])
    v = 60 # speed up
    bb.cmd(0x02, 0x30, [v, (h&0xff00)>>8, h&0xff, 1]) 
    v = 80 # speed up
    bb.cmd(0x02, 0x30, [v, (h&0xff00)>>8, h&0xff, 1]) 

    for t in range(several_times): # timesteps
        v = 100
        h = 270 
        bb.cmd(0x02, 0x30, [v, (h&0xff00)>>8, h&0xff, 1])

    v = 80 # speed down
    bb.cmd(0x02, 0x30, [v, (h&0xff00)>>8, h&0xff, 1]) 
    v = 60 # speed down
    bb.cmd(0x02, 0x30, [v, (h&0xff00)>>8, h&0xff, 1]) 
    v = 20 # speed down
    h +=30  # start preparing turn
    bb.cmd(0x02, 0x30, [v, (h&0xff00)>>8, h&0xff, 1]) 
    v = 0 # speed down
    h +=30  # start preparing turn
    bb.cmd(0x02, 0x30, [v, (h&0xff00)>>8, h&0xff, 1]) 
    h = 0  # back to start orientation
    bb.cmd(0x02, 0x30, [v, (h&0xff00)>>8, h&0xff, 1]) 


    keys = [False] * 1024
    print("Continue?")
    proceed = False

    while proceed == False:
        events = pygame.event.get()
        for event in events:
           if event.type == pygame.KEYDOWN:
               keys[event.key] = True
           elif event.type == pygame.KEYUP:
               keys[event.key] = False

        if keys[pygame.K_q] or keys[pygame.K_n]:
            bb.disconnect()
            quit()

        if keys[pygame.K_c] or keys[pygame.K_y]:
            proceed = True 

    print("Lets continue")

    # start with WASD commands
    h = 0

    while True:
        c.tick(10)
        events = pygame.event.get()
        for event in events:
            if event.type == pygame.KEYDOWN:
                keys[event.key] = True
            elif event.type == pygame.KEYUP:
                keys[event.key] = False

        if keys[pygame.K_w]:
            v = 100
        elif keys[pygame.K_s]:
            v = 255
        else:
            v = 0

        if keys[pygame.K_a]:
            h -= 15
        elif keys[pygame.K_d]:
            h += 15
        while h<0:
            h += 360
        while h>359:
            h -= 360
        print h

        bb.cmd(0x02, 0x30, [v, (h&0xff00)>>8, h&0xff, 1])

    # Must manually disconnect or you won't be able to reconnect.
    bb.disconnect()