from djitellopy import Tello
import cv2
import pygame
import numpy as np
import time

MOTION = 50
FPS = 120

class TelloController():
    def __init__(self):
        pygame.init()
        pygame.display.set_caption("Tello controller")
        self.screen = pygame.display.set_mode([960, 720])
        self.has_controller = pygame.joystick.get_count() > 0
        if self.has_controller:
            self.joystick = pygame.joystick.Joystick(0)
        self.tello = Tello()

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

        self.send_rc_control = False

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

    def run(self):
        self.tello.connect()

        # Reset stream in case it was left on by accident
        self.tello.streamoff()
        self.tello.streamon()

        frame_read = self.tello.get_frame_read()

        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 and self.has_controller:
                    self.stickmotion(event)
                elif event.type == pygame.JOYBUTTONDOWN and self.has_controller:
                    self.buttonup(event.button)

            if frame_read.stopped:
                break

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

            frame = frame_read.frame

            # Display battery status
            text = "{}%".format(self.tello.get_battery())
            cv2.putText(frame, text, (5, 720 - 5),
                cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            frame = np.rot90(frame)
            frame = np.flipud(frame)

            frame = pygame.surfarray.make_surface(frame)
            self.screen.blit(frame, (0, 0))
            pygame.display.update()

            time.sleep(1 / FPS)

        self.tello.end()

    def keydown(self, key):
        if key == pygame.K_UP:  # Forward
            self.for_back_velocity = MOTION
        elif key == pygame.K_DOWN:  # Backward
            self.for_back_velocity = -MOTION
        elif key == pygame.K_LEFT:  # Left
            self.left_right_velocity = -MOTION
        elif key == pygame.K_RIGHT:  # Right
            self.left_right_velocity = MOTION
        elif key == pygame.K_w:  # Up
            self.up_down_velocity = MOTION
        elif key == pygame.K_s:  # Down
            self.up_down_velocity = -MOTION
        elif key == pygame.K_a:  # Rotate counter clockwise
            self.yaw_velocity = -MOTION
        elif key == pygame.K_d:  # Rotate clockwise
            self.yaw_velocity = MOTION

    def keyup(self, key):
        if key == pygame.K_UP or key == pygame.K_DOWN:  # Reset forward/backward
            self.for_back_velocity = 0
        elif key == pygame.K_LEFT or key == pygame.K_RIGHT:  # Reset left/right
            self.left_right_velocity = 0
        elif key == pygame.K_w or key == pygame.K_s:  # Reset up/down
            self.up_down_velocity = 0
        elif key == pygame.K_a or key == pygame.K_d:  # Reset yaw
            self.yaw_velocity = 0
        elif key == pygame.K_t:  # Take off
            self.tello.takeoff()
            self.send_rc_control = True
        elif key == pygame.K_l:  # Land
            not self.tello.land()
            self.send_rc_control = False
        elif key == pygame.K_h:
            self.tello.flip_forward()
        elif key == pygame.K_j:
            self.tello.flip_back()

    def stickmotion(self, event):
        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 == 2: # 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 == 0:
            self.tello.land()
            self.send_rc_control = False
        elif button == 3:
            self.tello.takeoff()
            self.send_rc_control = True

    def update(self):
        if self.send_rc_control:
            self.tello.send_rc_control(
                # Whilst this is technically less accurate, I'm not sure how well
                # these drones handle repeating numbers like 0.333333333 and I doubt
                # that level of precision is necessary at all
                round(self.left_right_velocity, 3),
                round(self.for_back_velocity, 3),
                round(self.up_down_velocity, 3),
                round(self.yaw_velocity, 3)
            )

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