import cv2
import numpy as np


CAMERA_INDEX = 0
HOMOGRAPHY_FILE = "homography_matrix.npy"

OUTPUT_WIDTH = 800
OUTPUT_HEIGHT = 600

MINIMUM_OBJECT_AREA = 500


def create_red_mask(hsv_image):
    """
    Detect red using two HSV ranges.

    Red wraps around the HSV hue boundary, so two ranges
    are required.
    """

    lower_red_1 = np.array([0, 100, 80], dtype=np.uint8)
    upper_red_1 = np.array([10, 255, 255], dtype=np.uint8)

    lower_red_2 = np.array([170, 100, 80], dtype=np.uint8)
    upper_red_2 = np.array([179, 255, 255], dtype=np.uint8)

    mask_1 = cv2.inRange(hsv_image, lower_red_1, upper_red_1)
    mask_2 = cv2.inRange(hsv_image, lower_red_2, upper_red_2)

    red_mask = cv2.bitwise_or(mask_1, mask_2)

    kernel = np.ones((5, 5), dtype=np.uint8)

    red_mask = cv2.morphologyEx(
        red_mask,
        cv2.MORPH_OPEN,
        kernel,
        iterations=2,
    )

    red_mask = cv2.morphologyEx(
        red_mask,
        cv2.MORPH_CLOSE,
        kernel,
        iterations=2,
    )

    return red_mask


def detect_red_objects(image, mask):
    """Find red objects and draw their contours and centres."""

    contours, _ = cv2.findContours(
        mask,
        cv2.RETR_EXTERNAL,
        cv2.CHAIN_APPROX_SIMPLE,
    )

    detected_objects = []

    for contour in contours:
        area = cv2.contourArea(contour)

        if area < MINIMUM_OBJECT_AREA:
            continue

        x, y, width, height = cv2.boundingRect(contour)

        moments = cv2.moments(contour)

        if moments["m00"] == 0:
            continue

        centre_x = int(moments["m10"] / moments["m00"])
        centre_y = int(moments["m01"] / moments["m00"])

        detected_objects.append(
            {
                "centre": (centre_x, centre_y),
                "area": area,
                "bounding_box": (x, y, width, height),
            }
        )

        cv2.drawContours(image, [contour], -1, (0, 255, 0), 2)

        cv2.rectangle(
            image,
            (x, y),
            (x + width, y + height),
            (255, 0, 0),
            2,
        )

        cv2.circle(
            image,
            (centre_x, centre_y),
            6,
            (0, 255, 255),
            -1,
        )

        label = f"Red object: ({centre_x}, {centre_y})"

        cv2.putText(
            image,
            label,
            (x, max(y - 10, 20)),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.6,
            (0, 255, 0),
            2,
        )

    return detected_objects


def main():
    try:
        homography_matrix = np.load(HOMOGRAPHY_FILE)
    except FileNotFoundError:
        raise FileNotFoundError(
            f"'{HOMOGRAPHY_FILE}' was not found. "
            "Run homography_calibration.py first."
        )

    camera = cv2.VideoCapture(CAMERA_INDEX)

    if not camera.isOpened():
        raise RuntimeError("Could not open the camera.")

    print("Red-object detection started.")
    print("Press Q to quit.")

    while True:
        success, frame = camera.read()

        if not success:
            print("Could not read a frame from the camera.")
            break

        # Use this if the camera is mounted sideways.
        frame = cv2.rotate(frame, cv2.ROTATE_90_CLOCKWISE)

        top_down_view = cv2.warpPerspective(
            frame,
            homography_matrix,
            (OUTPUT_WIDTH, OUTPUT_HEIGHT),
        )

        blurred = cv2.GaussianBlur(
            top_down_view,
            (7, 7),
            0,
        )

        hsv_image = cv2.cvtColor(
            blurred,
            cv2.COLOR_BGR2HSV,
        )

        red_mask = create_red_mask(hsv_image)

        detected_objects = detect_red_objects(
            top_down_view,
            red_mask,
        )

        if detected_objects:
            largest_object = max(
                detected_objects,
                key=lambda item: item["area"],
            )

            centre_x, centre_y = largest_object["centre"]

            print(
                f"Largest red object centre: "
                f"x={centre_x}, y={centre_y}",
                end="\r",
            )

        cv2.imshow("Red Object Detection", top_down_view)
        cv2.imshow("Red Mask", red_mask)

        key = cv2.waitKey(1) & 0xFF

        if key == ord("q"):
            break

    camera.release()
    cv2.destroyAllWindows()


if __name__ == "__main__":
    main()