from time import sleep
from math import cos, sin

import cv2
import numpy as np

from robot import Robot

POWER = 70

TURN_T = 7.9 # 1 degree
DRIVE_T = 22 # 1 centimeter

RIGHT_WHEEL_OFFSET = 4

CLOCKWISE_OFFSET = 0.94

FOCAL_LENGTH = 1691

CAMERA_MATRIX = np.array(
    [[FOCAL_LENGTH,0,512],[0,FOCAL_LENGTH,360],[0,0,1]],
    dtype=np.float32
)
DIST_COEF = np.array([0,0,0,0,0], dtype=np.float32)

def find_aruco(image):
    aruco_dict = cv2.aruco.Dictionary_get(cv2.aruco.DICT_6X6_250)
    aruco_params = cv2.aruco.DetectorParameters_create()
    corners, ids, _ = cv2.aruco.detectMarkers(
        image,
        aruco_dict,
        parameters=aruco_params
    )

    if corners is None:
        return []

    return [(box[0], ids[i][0]) for i, box in enumerate(corners)]

def find_arucos(arlo):
    aruco_dict = {}
    theta = np.deg2rad(-20)
    rot = np.array([[cos(theta), -sin(theta)], [sin(theta), cos(theta)]])

    while True:
        arucos = find_aruco(arlo.take_photo())
        if arucos != []:
            for aruco, id_ in arucos:
                if id_ in [1,6] and id_ not in aruco_dict:
                    print(f"found box {id_}")
                    position = cv2.aruco.estimatePoseSingleMarkers(
                        np.array([aruco]), 14.5, CAMERA_MATRIX, DIST_COEF
                    )[1][0][0]
                    position = np.array([position[2], position[0]])
                    position[0] += 22.5
                    aruco_dict[id_] = position

            if len(aruco_dict) >= 2:
                break
        arlo.go_diff(POWER, POWER, 1, 0)
        sleep((20 * TURN_T * CLOCKWISE_OFFSET)/1000)
        arlo.stop()
        for key, value in aruco_dict.items():
            aruco_dict[key] = np.dot(rot, value)

    return np.array(list(aruco_dict.values())[:2])

def main():
    arlo = Robot()
    aruco_positions = find_arucos(arlo)
    print(aruco_positions)

    position = [
        np.average(aruco_positions[:,0]),
        np.average(aruco_positions[:,1])
    ]
    print(position)

    angle = np.rad2deg(np.arctan(position[1]/position[0]))
    drive_distance = np.sqrt(position[0]**2 + position[1]**2)
    if angle < 0:
        arlo.go_diff(POWER, POWER, 0, 1)
        sleep((abs(angle) * TURN_T)/1000)
        arlo.stop()
    else:
        arlo.go_diff(POWER, POWER, 1, 0)
        sleep((abs(angle) * TURN_T * CLOCKWISE_OFFSET)/1000)
        arlo.stop()

    arlo.go_diff(POWER, POWER + RIGHT_WHEEL_OFFSET, 1, 1)
    sleep((drive_distance * DRIVE_T)/1000)
    arlo.stop()


if __name__ == "__main__":
    main()