#!/usr/bin/env python3
"""Interactive metric ground-plane calibration for the line follower.

Four tape crosses with known positions relative to the vehicle front centre are
clicked in the live *raw* camera image.  Their matching virtual ground-plane
positions are used to calculate the homography consumed by line_follow_debug.
The camera remains forward-facing; the bird's-eye view is the result of the
calibration, not an assumption about where the camera is mounted.
"""

import threading

import cv2
import numpy as np
import rospy
from cv_bridge import CvBridge, CvBridgeError
from sensor_msgs.msg import Image


class CameraOffsetCalibrator:
    """Subscriber-only four-point perspective calibration tool."""

    DISPLAY_SCALE = 2
    MARKER_NAMES = ("near-left", "near-right", "far-right", "far-left")

    def __init__(self):
        rospy.init_node("camera_offset_calibrator")
        self.image_topic = rospy.get_param("~image_topic", "/usb_cam/image_raw")
        self.reference_size = self._read_size("~perspective_reference_size", [640, 480])
        self.marker_positions_m = self._read_world_points(
            "~calibration_markers_m",
            [[0.30, 0.20], [0.30, -0.20], [0.90, -0.20], [0.90, 0.20]],
        )
        self.metric_scale = max(1.0, float(rospy.get_param("~metric_scale_px_per_m", 400.0)))
        self.metric_origin = self._read_point("~metric_origin_px", [320, 480])

        self.bridge = CvBridge()
        self.lock = threading.Lock()
        self.source_image = None
        self.clicked_points = []
        self.last_result = None

        self.subscriber = rospy.Subscriber(
            self.image_topic, Image, self.image_callback, queue_size=1
        )
        rospy.loginfo(
            "Metric perspective calibration subscribes to %s only; it never publishes /cmd_vel.",
            self.image_topic,
        )
        rospy.loginfo(
            "Click four marker centres in this order: near-left, near-right, far-right, far-left."
        )

    @staticmethod
    def _read_size(name, default):
        values = rospy.get_param(name, default)
        if not isinstance(values, (list, tuple)) or len(values) != 2:
            rospy.logwarn("%s must contain [width, height]; using %s", name, default)
            values = default
        return max(1, int(values[0])), max(1, int(values[1]))

    @staticmethod
    def _read_point(name, default):
        values = rospy.get_param(name, default)
        if not isinstance(values, (list, tuple)) or len(values) != 2:
            rospy.logwarn("%s must contain [x, y]; using %s", name, default)
            values = default
        return np.array(values, dtype=np.float32)

    @staticmethod
    def _read_world_points(name, default):
        values = rospy.get_param(name, default)
        valid = isinstance(values, (list, tuple)) and len(values) == 4
        if valid:
            valid = all(isinstance(point, (list, tuple)) and len(point) == 2 for point in values)
        if not valid:
            rospy.logwarn("%s must contain four [forward_m, left_m] points; using defaults", name)
            values = default
        return np.array(values, dtype=np.float32)

    def image_callback(self, message):
        try:
            raw = self.bridge.imgmsg_to_cv2(message, desired_encoding="bgr8")
        except CvBridgeError as error:
            rospy.logwarn_throttle(2.0, "Cannot convert camera image: %s", error)
            return

        # Calibration values are deliberately recorded in the same 640x480
        # reference coordinate system as line_follow_debug.yaml.
        source = cv2.resize(raw, self.reference_size, interpolation=cv2.INTER_AREA)
        with self.lock:
            self.source_image = source

    def destination_points(self):
        """Map vehicle-ground coordinates (forward, left) to a virtual image."""
        forward = self.marker_positions_m[:, 0]
        left = self.marker_positions_m[:, 1]
        return np.column_stack(
            (
                self.metric_origin[0] - left * self.metric_scale,
                self.metric_origin[1] - forward * self.metric_scale,
            )
        ).astype(np.float32)

    def result_matrix(self):
        if len(self.clicked_points) != 4:
            return None
        return cv2.getPerspectiveTransform(
            np.array(self.clicked_points, dtype=np.float32), self.destination_points()
        )

    def log_result(self):
        matrix = self.result_matrix()
        if matrix is None:
            return
        source = [[int(x), int(y)] for x, y in self.clicked_points]
        destination = [[round(float(x), 1), round(float(y), 1)] for x, y in self.destination_points()]
        self.last_result = (source, destination)
        rospy.loginfo("=" * 64)
        rospy.loginfo("METRIC PERSPECTIVE CALIBRATION RESULT")
        rospy.loginfo("src_pts: %s", source)
        rospy.loginfo("dst_pts: %s", destination)
        rospy.loginfo("metric_scale_px_per_m: %.1f", self.metric_scale)
        rospy.loginfo("metric_origin_px: [%d, %d]", int(self.metric_origin[0]), int(self.metric_origin[1]))
        rospy.loginfo("Copy only src_pts and dst_pts into line_follow_debug.yaml.")
        rospy.loginfo("=" * 64)

    def mouse_callback(self, event, x, y, _flags, _userdata):
        if event == cv2.EVENT_RBUTTONDOWN:
            self.clicked_points = []
            self.last_result = None
            rospy.loginfo("Calibration points cleared.")
            return
        if event != cv2.EVENT_LBUTTONDOWN:
            return

        point = (int(x / self.DISPLAY_SCALE), int(y / self.DISPLAY_SCALE))
        with self.lock:
            source = self.source_image
        if source is None or not (0 <= point[0] < source.shape[1] and 0 <= point[1] < source.shape[0]):
            return
        if len(self.clicked_points) >= 4:
            rospy.logwarn("Four points are already selected. Press R or right-click to start again.")
            return

        self.clicked_points.append(point)
        index = len(self.clicked_points) - 1
        marker = self.marker_positions_m[index]
        rospy.loginfo(
            "Point %d/4 (%s): image=[%d, %d], ground=[forward %.3f m, left %.3f m]",
            index + 1,
            self.MARKER_NAMES[index],
            point[0], point[1], marker[0], marker[1],
        )
        if len(self.clicked_points) == 4:
            self.log_result()

    def source_display(self, source):
        view = source.copy()
        for index, point in enumerate(self.clicked_points):
            color = (0, 0, 255) if index < 2 else (0, 255, 255)
            cv2.circle(view, point, 5, color, -1)
            cv2.putText(
                view,
                "%d %s" % (index + 1, self.MARKER_NAMES[index]),
                (point[0] + 8, point[1] - 8),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.42,
                color,
                1,
            )
        if len(self.clicked_points) == 4:
            cv2.polylines(view, [np.array(self.clicked_points, dtype=np.int32)], True, (0, 255, 255), 1)
        cv2.putText(view, "RAW CAMERA: CLICK 4 GROUND MARKERS", (10, 28),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 255), 2)
        cv2.putText(view, "1 near-L  2 near-R  3 far-R  4 far-L", (10, 53),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.52, (0, 255, 255), 1)
        return cv2.resize(
            view, (view.shape[1] * self.DISPLAY_SCALE, view.shape[0] * self.DISPLAY_SCALE),
            interpolation=cv2.INTER_LINEAR,
        )

    def draw_ground_grid(self, image):
        """Draw a 10 cm vehicle-coordinate grid in the virtual view."""
        view = image.copy()
        height, width = view.shape[:2]
        max_forward = int((self.metric_origin[1] + 1) / self.metric_scale) + 1
        half_width = int(max(self.metric_origin[0], width - self.metric_origin[0]) / self.metric_scale) + 1
        for forward_cm in range(0, max_forward * 10 + 1, 10):
            forward = forward_cm / 10.0
            y = int(round(self.metric_origin[1] - forward * self.metric_scale))
            if 0 <= y < height:
                color = (55, 55, 55) if forward_cm % 50 else (100, 100, 100)
                cv2.line(view, (0, y), (width - 1, y), color, 1)
        for left_cm in range(-half_width * 10, half_width * 10 + 1, 10):
            left = left_cm / 10.0
            x = int(round(self.metric_origin[0] - left * self.metric_scale))
            if 0 <= x < width:
                color = (55, 55, 55) if left_cm % 50 else (100, 100, 100)
                cv2.line(view, (x, 0), (x, height - 1), color, 1)
        if 0 <= int(self.metric_origin[0]) < width and 0 <= int(self.metric_origin[1]) < height:
            cv2.drawMarker(
                view, tuple(self.metric_origin.astype(int)), (0, 255, 0), cv2.MARKER_CROSS, 14, 2
            )
        return view

    def bird_display(self, source):
        matrix = self.result_matrix()
        if matrix is None:
            bird = np.zeros_like(source)
            text = "Select %d more point(s)" % (4 - len(self.clicked_points))
            cv2.putText(bird, text, (20, 42), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 255), 2)
        else:
            bird = cv2.warpPerspective(source, matrix, self.reference_size, flags=cv2.INTER_LINEAR)
            bird = self.draw_ground_grid(bird)
            cv2.putText(bird, "METRIC BIRD VIEW (origin = vehicle front centre)", (10, 28),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.48, (0, 255, 255), 1)
            cv2.putText(bird, "S: print result  R/right-click: reset  Q/Esc: quit", (10, 50),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.43, (0, 255, 255), 1)
        return cv2.resize(
            bird, (bird.shape[1] * self.DISPLAY_SCALE, bird.shape[0] * self.DISPLAY_SCALE),
            interpolation=cv2.INTER_LINEAR,
        )

    def run(self):
        source_window = "Metric perspective calibration: raw camera"
        bird_window = "Metric perspective calibration: calibrated ground view"
        cv2.namedWindow(source_window, cv2.WINDOW_AUTOSIZE)
        cv2.namedWindow(bird_window, cv2.WINDOW_AUTOSIZE)
        cv2.setMouseCallback(source_window, self.mouse_callback)
        rate = rospy.Rate(15)
        while not rospy.is_shutdown():
            with self.lock:
                source = None if self.source_image is None else self.source_image.copy()
            if source is not None:
                cv2.imshow(source_window, self.source_display(source))
                cv2.imshow(bird_window, self.bird_display(source))
            key = cv2.waitKey(1) & 0xFF
            if key in (ord("r"), ord("R")):
                self.clicked_points = []
                self.last_result = None
                rospy.loginfo("Calibration points cleared.")
            elif key in (ord("s"), ord("S")):
                self.log_result()
            elif key in (ord("q"), ord("Q"), 27):
                break
            rate.sleep()
        cv2.destroyAllWindows()


if __name__ == "__main__":
    try:
        CameraOffsetCalibrator().run()
    except rospy.ROSInterruptException:
        pass