import os.path
from util import *
from yoloDetections_to_npArray import tran
import yolov5


def getFiles(path):
    ans = []
    for root, dirs, files in os.walk(path):
        for i in files:
            ans.append(root + '/' + i)
    return ans


class simulation_drone:
    def __init__(self):
        self.img_files = getFiles("imgs")
        self.img_files_index = 0
        self.gear = 0

    def drone_increase_gear(self):
        if self.gear < 3:
            self.gear += 1

    def drone_decrease_gear(self):
        if self.gear > 0:
            self.gear -= 1

    def drone_get_gear(self):
        return self.gear + 1

    def drone_get_battery(self):
        return 100

    @fn_timer
    def drone_move(self, a, b, c, d, speed_magnification=1):
        log("drone_move", a, b, c, d, speed_magnification)
        time.sleep(0.1)

    @fn_timer
    def drone_flip(self, direction):
        log("drone_flip", direction)
        time.sleep(0.1)

    # 获取当前飞行速度
    def drone_get_speed(self):
        return 0

    # 起飞
    @fn_timer
    def drone_takeoff(self):
        log("drone_takeoff")
        time.sleep(2)

    # 停止
    @fn_timer
    def drone_stop(self):
        log("drone_stop")
        time.sleep(0.1)

    # 降落
    @fn_timer
    def drone_land(self):
        log("drone_land")
        time.sleep(2)

    def drone_init_camera(self):
        log("drone_init_camera")

    def drone_get_img(self):
        img = cv2.imread(self.img_files[self.img_files_index])
        self.img_files_index += 1
        self.img_files_index %= len(self.img_files)
        time.sleep(0.05)
        return img

    def drone_close_camera(self):
        log("drone_close_camera")

    def close(self):
        log("drone_close")

#
# def main():
#     drone = simulation_drone()
#     img = drone.drone_get_img()
#     model = yolov5.load('balloon_recognition_model.pt')
#     while True:
#         img = drone.drone_get_img()
#         cv2.waitKey(1)
#         image_recognition_results = model(img)
#         img = tran(image_recognition_results)
#         cv2.imshow("Drone", img)
#         .drone_autonomous_operation(image_recognition_results)
#
#
# if __name__ == '__main__':
#     main()