#!/usr/bin/python3
# coding=utf8
#第6章 拓展课程\2.拓展课程之智能搬运/第4课 标签跟随(6.Advanced Lesson\2.AI Transport Lesson\Lesson 4 Tag Tracking)
import sys
import cv2
import time
import math
import threading
import numpy as np
from common import apriltag
from common import yaml_handle
from common import kinematics
from common.ros_robot_controller_sdk import Board
from calibration.camera import Camera 


board = Board()
ik = kinematics.IK(board)



APRILTAG_SIZE = 36 # mm
FACTOR = 147*100/APRILTAG_SIZE
DISTANCE_TO_CAMERA = 120 #/mm

if sys.version_info.major == 2:
    print('Please run this program with python3!')
    sys.exit(0)

# #加载参数(load parameter)
# param_data = np.load(calibration_param_path + '.npz')

# #获取参数(obtain parameter)
# mtx = param_data['mtx_array']
# dist = param_data['dist_array']
# newcameramtx, _ = cv2.getOptimalNewCameraMatrix(mtx, dist, (640, 480), 0, (640, 480))
# mapx, mapy = cv2.initUndistortRectifyMap(mtx, dist, None, newcameramtx, (640, 480), 5)

# fx = mtx[0, 0]
# fy = mtx[1, 1]
# cx = mtx[0, 2]
# cy = mtx[1, 2]

lab_data = None
servo_data = None
def load_config():
    global lab_data, servo_data
    
    lab_data = yaml_handle.get_yaml_data(yaml_handle.lab_file_path)
    servo_data = yaml_handle.get_yaml_data(yaml_handle.servo_file_path)

load_config()

# 初始位置(initial position)
def init_move():

    board.pwm_servo_set_position(0.5, [[1, 1500] , [2, servo_data['servo2']]])

    ik.stand(ik.initial_pos)

ditance, angle = -1, 0
# 变量重置(reset variables)
def reset():
    global __target_color
    global distance, angle

    __target_color = ()
    distance, angle = -1, 0

# app初始化调用(call app initialization)
def init():
    print("Follow Init")
    init_move()

__isRunning = False
# app开始玩法调用(call app start game)
def start():
    global __isRunning
    reset()
    __isRunning = True
    print("Follow Start")

# app停止玩法调用(call app stop game)
def stop():
    global __isRunning
    __isRunning = False
    print("Follow Stop")

# app退出玩法调用(call app exit game)
def exit():
    global __isRunning
    __isRunning = False
    print("Follow Exit")

#执行动作组(execute action group)
def move():

    while True:
        if __isRunning:
            if distance > 0:
                if angle > 5:
                    if 40 > angle:
                        ik.turn_right(ik.initial_pos, 2, angle, 50, 1)
                    else:
                        time.sleep(0.01)
                elif angle < -5:
                    if -40 < angle:
                        ik.turn_left(ik.initial_pos, 2, -angle, 50, 1)
                    else:
                        time.sleep(0.01)
                elif 250 < distance:
                    d_d = int(distance - 250)
                    if d_d > 150:
                        ik.go_forward(ik.initial_pos, 2, 150, 60, 1)
                    elif d_d > 10:
                        ik.go_forward(ik.initial_pos, 2, d_d, 50, 1)
                    else:
                        time.sleep(0.01)
                elif 0 < distance < 150:
                    ik.back(ik.initial_pos, 2, 150, 100, 1)
                else:
                    time.sleep(0.01)
            else:
                time.sleep(0.01)

#启动动作的线程(enable the thread of action)
th = threading.Thread(target=move)
th.setDaemon(True)
th.start()

# 检测apriltag(detect apriltag)
detector = apriltag.Detector(searchpath=apriltag._get_demo_searchpath())
def apriltagDetect(img):
    gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    detections = detector.detect(gray, return_image=False)
    tag = [-1, -1, 0]
    
    if len(detections) != 0:
        for detection in detections:
            corners = detection.corners
            corners = np.rint(corners)  # 获取四个角点(obtain the four corner points)
            w = math.sqrt(pow(corners[0][0] - corners[1][0], 2) + pow(corners[0][1] - corners[1][1], 2))
            h = math.sqrt(pow(corners[0][0] - corners[3][0], 2) + pow(corners[0][1] - corners[3][1], 2))
            
            cv2.drawContours(img, [np.array(corners, np.intp)], -1, (0, 255, 255), 2)
            
            tag_family = str(detection.tag_family, encoding='utf-8')  # 获取tag_family(obtain tag_family)
            tag_id = str(detection.tag_id)  # 获取tag_id(obtain tag_id)

            object_center_x, object_center_y = int(detection.center[0]), int(detection.center[1])  # 中心点(center point)
            cv2.circle(img, (object_center_x, object_center_y), 5, (0, 255, 255), -1)

            if tag_family == 'tag36h11':
                tag = [object_center_x, object_center_y, w]

    return tag

def calculate_rotation(center_x, img_width, apriltag_width):
    x = (center_x - img_width/2)*(APRILTAG_SIZE/apriltag_width)
    distance = int(APRILTAG_SIZE*FACTOR/apriltag_width)
    angle = int(math.degrees(math.atan2(x, DISTANCE_TO_CAMERA + distance)))
    return angle, distance

def run(img):
    global distance, angle
    
    img_copy = img.copy()
    img_h, img_w = img.shape[:2]
    
    if not __isRunning:
        return img

    tag = apriltagDetect(img)
    if tag[0] != -1:
        centerX, centerY, apriltag_width = tag[0], tag[1], tag[2]
        angle, distance = calculate_rotation(centerX, img_w, apriltag_width)
        #print(angle, distance)
    else:
        angle, distance  = 0, 0

    return img

if __name__ == '__main__':
    init()
    start()
    
    camera = Camera()
    camera.camera_open()
    while True:
        img = camera.frame
        if img is not None:
            frame = img.copy()
            Frame = run(frame)
            cv2.imshow('Frame', Frame)
            key = cv2.waitKey(1)
            if key == 27:
                break
        else:
            time.sleep(0.01)
    camera.camera_close()
    cv2.destroyAllWindows()