from time import sleep

import cv2
import numpy as np

from robot import Robot

POWER = 70

TURN_T = 0.079 # 10 degrees
DRIVE_T = 0.22 # 10 centimeter

RIGHT_WHEEL_OFFSET = 4

CLOCKWISE_OFFSET = 0.96

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]) for i, box in enumerate(corners)]

def main():
    arlo = Robot()
    while True:
        arucos = find_aruco(arlo.take_photo())
        if arucos != []:
            break
        arlo.go_diff(POWER, POWER, 1, 0)
        sleep(3 * TURN_T)
        arlo.stop()

    position = cv2.aruco.estimatePoseSingleMarkers(
        np.array([arucos[0][0]]), 14.5, CAMERA_MATRIX, DIST_COEF
    )[1][0][0]
    print(position)

    angle = np.rad2deg(np.arctan(position[0]/position[2]))
    print(angle)

    # drive_distance = np.sqrt(position[0]**2 + position[2]**2)


if __name__ == "__main__":
    main()