import cv2
import pygame
import time

SPEED = 50
FPS = 120

class FrontEnd(object):
    def __init__(self):
        pygame.init()
        pygame.display.set_caption("Tello video stream")
        self.screen = pygame.display.set_mode([960, 720])
        self.joystick = pygame.joystick.Joystick(0)

        self.for_back_velocity = 0
        self.left_right_velocity = 0
        self.up_down_velocity = 0
        self.yaw_velocity = 0
        self.speed = 10

        self.send_rc_control = False

        pygame.time.set_timer(pygame.USEREVENT + 1, 1000 // FPS)

    def run(self):
        should_stop = False
        while not should_stop:
            for event in pygame.event.get():
                if event.type == pygame.USEREVENT + 1:
                    self.update()
                elif event.type == pygame.QUIT:
                    should_stop = True
                elif event.type == pygame.KEYDOWN:
                    if event.key == pygame.K_ESCAPE:
                        should_stop = True
                    else:
                        self.keydown(event.key)
                elif event.type == pygame.KEYUP:
                    self.keyup(event.key)
                elif event.type == pygame.JOYAXISMOTION:
                    self.stickmotion(event)
                elif event.type == pygame.JOYBUTTONDOWN:
                    self.buttonup(event.button)

            self.screen.fill([0, 0, 0])

            time.sleep(1 / FPS)

    def keydown(self, key):
        if key == pygame.K_UP:  # set forward velocity
            self.for_back_velocity = SPEED
        elif key == pygame.K_DOWN:  # set backward velocity
            self.for_back_velocity = -SPEED
        elif key == pygame.K_LEFT:  # set left velocity
            self.left_right_velocity = -SPEED
        elif key == pygame.K_RIGHT:  # set right velocity
            self.left_right_velocity = SPEED
        elif key == pygame.K_w:  # set up velocity
            self.up_down_velocity = SPEED
        elif key == pygame.K_s:  # set down velocity
            self.up_down_velocity = -SPEED
        elif key == pygame.K_a:  # set yaw counter clockwise velocity
            self.yaw_velocity = -SPEED
        elif key == pygame.K_d:  # set yaw clockwise velocity
            self.yaw_velocity = SPEED

    def keyup(self, key):
        if key == pygame.K_UP or key == pygame.K_DOWN:  # set zero forward/backward velocity
            self.for_back_velocity = 0
        elif key == pygame.K_LEFT or key == pygame.K_RIGHT:  # set zero left/right velocity
            self.left_right_velocity = 0
        elif key == pygame.K_w or key == pygame.K_s:  # set zero up/down velocity
            self.up_down_velocity = 0
        elif key == pygame.K_a or key == pygame.K_d:  # set zero yaw velocity
            self.yaw_velocity = 0
        elif key == pygame.K_t:  # takeoff
            print("Taking off")
            self.send_rc_control = True
        elif key == pygame.K_l:  # land
            print("Landing")
            self.send_rc_control = False

    def stickmotion(self, event):
        print(f"Motion: {event.axis}")
        if event.axis == 0: # left stick left/right
            self.yaw_velocity = round(self.joystick.get_axis(event.axis), 3) * SPEED
        elif event.axis == 3: # right stick left/right
            self.left_right_velocity = round(self.joystick.get_axis(event.axis), 3) * SPEED
        elif event.axis == 4: # right stick up/down
            self.for_back_velocity = -round(self.joystick.get_axis(event.axis), 3) * SPEED
        elif event.axis == 2: # left trigger
            self.up_down_velocity = -min(round(self.joystick.get_axis(event.axis), 3) + 1, 1) * SPEED
        elif event.axis == 5: # right trigger
            self.up_down_velocity = min(round(self.joystick.get_axis(event.axis), 3) + 1, 1) * SPEED

    def buttonup(self, button):
        if button == 2:
            print("Landing")
            self.send_rc_control = False
        elif button == 0:
            print("Taking off")
            self.send_rc_control = True

    def update(self):
        if self.send_rc_control:
            print(f"Left/Right: {self.left_right_velocity}, Front/Back: {self.for_back_velocity}, Up/Down: {self.up_down_velocity}, Yaw: {self.yaw_velocity}")

if __name__ == '__main__':
    FrontEnd().run()