import cv2
import numpy as np


IMAGE_PATH = "calibration_image.jpg"
HOMOGRAPHY_FILE = "homography_matrix.npy"

OUTPUT_WIDTH = 800
OUTPUT_HEIGHT = 600

clicked_points = []


def mouse_callback(event, x, y, flags, parameter):
    """Store workspace corners selected by the user."""

    if event == cv2.EVENT_LBUTTONDOWN and len(clicked_points) < 4:
        clicked_points.append([x, y])
        print(f"Point {len(clicked_points)}: ({x}, {y})")


def draw_selected_points(image):
    """Draw selected calibration points on the image."""

    display = image.copy()

    for index, point in enumerate(clicked_points):
        x, y = point

        cv2.circle(display, (x, y), 7, (0, 0, 255), -1)

        cv2.putText(
            display,
            str(index + 1),
            (x + 10, y - 10),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.7,
            (0, 255, 0),
            2,
        )

    if len(clicked_points) > 1:
        for index in range(len(clicked_points) - 1):
            point_1 = tuple(clicked_points[index])
            point_2 = tuple(clicked_points[index + 1])

            cv2.line(display, point_1, point_2, (255, 0, 0), 2)

    if len(clicked_points) == 4:
        cv2.line(
            display,
            tuple(clicked_points[3]),
            tuple(clicked_points[0]),
            (255, 0, 0),
            2,
        )

    return display


def calculate_homography(image):
    """Calculate and save the perspective-transformation matrix."""

    source_points = np.float32(clicked_points)

    destination_points = np.float32(
        [
            [0, 0],
            [OUTPUT_WIDTH - 1, 0],
            [OUTPUT_WIDTH - 1, OUTPUT_HEIGHT - 1],
            [0, OUTPUT_HEIGHT - 1],
        ]
    )

    homography_matrix = cv2.getPerspectiveTransform(
        source_points,
        destination_points,
    )

    np.save(HOMOGRAPHY_FILE, homography_matrix)

    corrected_image = cv2.warpPerspective(
        image,
        homography_matrix,
        (OUTPUT_WIDTH, OUTPUT_HEIGHT),
    )

    print("\nHomography matrix:")
    print(homography_matrix)
    print(f"\nSaved as: {HOMOGRAPHY_FILE}")

    return corrected_image


def main():
    image = cv2.imread(IMAGE_PATH)

    if image is None:
        raise FileNotFoundError(
            f"Could not load '{IMAGE_PATH}'. "
            "Check that the image exists in the script directory."
        )

    # Rotate the image 90 degrees clockwise.
    image = cv2.rotate(image, cv2.ROTATE_90_CLOCKWISE)

    cv2.namedWindow("Homography Calibration")
    cv2.setMouseCallback("Homography Calibration", mouse_callback)

    print("Select the four workspace corners in this order:")
    print("1. Top-left")
    print("2. Top-right")
    print("3. Bottom-right")
    print("4. Bottom-left")
    print("\nPress C to calculate the homography.")
    print("Press R to reset the points.")
    print("Press Q to quit.")

    while True:
        display = draw_selected_points(image)

        cv2.imshow("Homography Calibration", display)

        key = cv2.waitKey(20) & 0xFF

        if key == ord("r"):
            clicked_points.clear()
            print("\nPoints reset.")

        elif key == ord("c"):
            if len(clicked_points) != 4:
                print("Please select exactly four points.")
                continue

            corrected_image = calculate_homography(image)
            cv2.imshow("Corrected Top-Down View", corrected_image)

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

    cv2.destroyAllWindows()


if __name__ == "__main__":
    main()