import cv2
import cv2.aruco as aruco
# import dronekit
# from dronekit import connect
# from dronekit import VehicleMode
# from dronekit import mavutil
import logging
import time

logging.basicConfig(level='INFO')

class VisionSystem:
    def __init__(self) -> None:
        self.capture = cv2.VideoCapture(0)
        self.capture.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)
        self.capture.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)
        _, frame = self.capture.read()
        logging.info("Initialised camera")
        logging.info('Resolution: ' + str(frame.shape[0]) + ' x ' + str(frame.shape[1]))

    def find_marker(self, img, markerSize = 6, totalMarkers=250):
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        key = getattr(aruco, f'DICT_{markerSize}X{markerSize}_{totalMarkers}')
        aruco_dict = aruco.Dictionary_get(key)
        aruco_param = aruco.DetectorParameters_create()
        bboxs, ids, rejected = aruco.detectMarkers(gray, aruco_dict, parameters=aruco_param)
        if ids:
            tl = bboxs[0][0][0][0], bboxs[0][0][0][1]
            tr = bboxs[0][0][1][0], bboxs[0][0][1][1]
            br = bboxs[0][0][2][0], bboxs[0][0][2][1]
            bl = bboxs[0][0][3][0], bboxs[0][0][3][1]
            # logging.info(f"{tl} | {tr}")
            # logging.info(f"{bl} | {br}")
            # logging.info(f"-----------")

            # Get centre point of marker
            total_x = [p[0] for p in [tr, tl, bl, br]]
            total_y = [p[1] for p in [tr, tl, bl, br]]
            centroid = (sum(total_x) / 4, sum(total_y) / 4)

            # Get offset from centre
            offset_y = (centroid[0] - (1280/2)) / 1280
            offset_x = (centroid[1] - (720/2)) / 720
            clean_offset_y = float("{0:.4f}".format(offset_y))
            clean_offset_x = float("{0:.4f}".format(offset_x))

            return (clean_offset_x, clean_offset_y)
        return None

    def calculate_angular_offset(self, pixel_offset):
        '''
        -0.5 -> 0.5
        '''
        logging.info("Calculating offset PWM signal")
        x = pixel_offset[0]
        y = float("{0:.1f}".format((pixel_offset[1] + 1.5) * 1000))
        logging.info(f"Calculated offset: {y}")
        return (x, y)

    def get_pallet_angle(self):
        while True:
            success, img = self.capture.read()
            logging.warning(f"Capture success: {success}")
            logging.info("Looking for marker...")
            raw_cargo_offset = self.find_marker(img)
            logging.info(f"Data from find_marker(img): {raw_cargo_offset}")
            # cv2.imshow('img',img)
            if raw_cargo_offset:
                logging.info(f"Found marker in get_pallet_angle {raw_cargo_offset}")
                # cargo_offset = self.calculate_angular_offset(raw_cargo_offset)
                # logging.info(f"PWM signal: {cargo_offset[1]}")
                # return cargo_offset
            # return None

v = VisionSystem()
v.get_pallet_angle()