#!/usr/bin/env python3
"""Recognise LED direction signs with an RKNN YOLOv5 model.

The node deliberately consumes the shared camera topic instead of opening a
camera device.  Camera exposure, gain and white balance remain owned by the
single camera-driver node.
"""

from __future__ import annotations

import atexit
from collections import Counter, deque
import logging
from pathlib import Path
import subprocess
import threading
from typing import Iterable, Optional, Sequence, Tuple

import cv2
import numpy as np
import rospy
import yaml
from cv_bridge import CvBridge, CvBridgeError
from sensor_msgs.msg import Image
from std_msgs.msg import Float32, String
from std_srvs.srv import SetBool, SetBoolResponse

try:
    from rknnlite.api import RKNNLite
except ImportError as error:  # pragma: no cover - depends on target hardware
    raise RuntimeError(
        "Unable to import RKNNLite. Start this node through the "
        "traffic_sign_node wrapper so it uses /home/ucar/venv3.9."
    ) from error

# rknnlite 1.5.2 replaces Python's normal log level names (``DEBUG``,
# ``INFO``...) with one-character variants.  ROS Noetic's logging config uses
# the normal names, so restore them before rospy.init_node() configures logs.
for _level, _name in (
    (logging.CRITICAL, "CRITICAL"),
    (logging.ERROR, "ERROR"),
    (logging.WARNING, "WARNING"),
    (logging.INFO, "INFO"),
    (logging.DEBUG, "DEBUG"),
    (logging.NOTSET, "NOTSET"),
):
    logging.addLevelName(_level, _name)


DEFAULT_CLASSES = ("left", "right", "straight", "stop")
DEFAULT_ANCHORS = np.array(
    [[10, 13], [16, 30], [33, 23], [30, 61], [62, 45], [59, 119],
     [116, 90], [156, 198], [373, 326]],
    dtype=np.float32,
)
DEFAULT_MASKS = ((0, 1, 2), (3, 4, 5), (6, 7, 8))


def sigmoid(values: np.ndarray) -> np.ndarray:
    return 1.0 / (1.0 + np.exp(-values))


def letterbox(image: np.ndarray, size: int) -> Tuple[np.ndarray, float, Tuple[float, float]]:
    """Resize without stretching and return scale/padding for box restoration."""
    height, width = image.shape[:2]
    scale = min(float(size) / height, float(size) / width)
    resized_width, resized_height = int(round(width * scale)), int(round(height * scale))
    resized = cv2.resize(image, (resized_width, resized_height), interpolation=cv2.INTER_LINEAR)
    pad_x = (size - resized_width) / 2.0
    pad_y = (size - resized_height) / 2.0
    bordered = cv2.copyMakeBorder(
        resized,
        int(round(pad_y - 0.1)),
        int(round(pad_y + 0.1)),
        int(round(pad_x - 0.1)),
        int(round(pad_x + 0.1)),
        cv2.BORDER_CONSTANT,
        value=(114, 114, 114),
    )
    return bordered, scale, (pad_x, pad_y)


def nms_boxes(boxes: np.ndarray, scores: np.ndarray, threshold: float) -> np.ndarray:
    x1, y1, x2, y2 = boxes.T
    areas = (x2 - x1) * (y2 - y1)
    order = scores.argsort()[::-1]
    keep = []
    while order.size:
        current = order[0]
        keep.append(current)
        if order.size == 1:
            break
        remaining = order[1:]
        xx1 = np.maximum(x1[current], x1[remaining])
        yy1 = np.maximum(y1[current], y1[remaining])
        xx2 = np.minimum(x2[current], x2[remaining])
        yy2 = np.minimum(y2[current], y2[remaining])
        intersection = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
        union = areas[current] + areas[remaining] - intersection
        iou = intersection / np.maximum(union, 1e-6)
        order = remaining[iou <= threshold]
    return np.asarray(keep, dtype=np.int32)


class TrafficSignRecognizer:
    def __init__(self) -> None:
        rospy.init_node("traffic_sign_recognition")
        self._bridge = CvBridge()
        self._state_lock = threading.RLock()
        self._package_dir = Path(__file__).resolve().parent.parent
        default_calibration = self._package_dir / "config" / "head_camera.yaml"
        self._calibration_file = Path(
            rospy.get_param("~calibration_file", str(default_calibration))
        ).expanduser()
        self._camera_matrix, self._distortion_coefficients, self._calibration_size = (
            self._load_camera_calibration(self._calibration_file)
        )
        self._classes = tuple(rospy.get_param("~classes", list(DEFAULT_CLASSES)))
        self._input_size = int(rospy.get_param("~input_size", 640))
        self._object_threshold = float(rospy.get_param("~object_threshold", 0.70))
        self._nms_threshold = float(rospy.get_param("~nms_threshold", 0.45))
        self._flip_horizontal = bool(rospy.get_param("~flip_horizontal", True))
        self._brightness = float(rospy.get_param("~brightness", 0.0))
        self._contrast = float(rospy.get_param("~contrast", 1.0))
        self._saturation = float(rospy.get_param("~saturation", 1.0))
        self._lab_clahe = bool(rospy.get_param("~lab_clahe", True))
        self._clahe = cv2.createCLAHE(
            clipLimit=float(rospy.get_param("~clahe_clip_limit", 1.5)),
            tileGridSize=tuple(rospy.get_param("~clahe_tile_grid", [8, 8])),
        )
        self._gamma = float(rospy.get_param("~gamma", 1.0))
        self._stability_window = int(rospy.get_param("~stability_window", 5))
        self._stable_count = int(rospy.get_param("~stable_count", 4))
        if self._stable_count > self._stability_window:
            raise ValueError("stable_count must not exceed stability_window")
        self._history: deque[Optional[str]] = deque(maxlen=self._stability_window)
        self._last_direction = "NONE"

        default_model = self._package_dir / "models" / "traffic_sign_direction.rknn"
        configured_model = Path(rospy.get_param("~model_path", str(default_model))).expanduser()
        self._model_path = configured_model
        if not self._model_path.is_file():
            raise FileNotFoundError(
                "RKNN model is missing: %s. Train and convert the LED sign model, then place "
                "traffic_sign_direction.rknn in this package's models directory or set ~model_path."
                % self._model_path
            )

        self._rknn = RKNNLite()
        result = self._rknn.load_rknn(str(self._model_path))
        if result != 0:
            raise RuntimeError("RKNN model load failed with code %s: %s" % (result, self._model_path))
        result = self._rknn.init_runtime()
        if result != 0:
            raise RuntimeError("RKNN runtime initialisation failed with code %s" % result)
        self._camera_profile_enabled = bool(rospy.get_param("~camera_profile_enabled", False))
        self._camera_device = str(rospy.get_param("~camera_device", "/dev/video0"))
        self._restore_camera_on_shutdown = bool(
            rospy.get_param("~restore_camera_on_shutdown", True)
        )
        # Cache controls now: ROS parameter access is not reliable once the
        # ROS shutdown sequence has begun.
        self._led_camera_controls = (
            ("exposure_auto", rospy.get_param("~led_exposure_auto", 1)),
            ("exposure_auto_priority", rospy.get_param("~led_exposure_auto_priority", 0)),
            ("exposure_absolute", rospy.get_param("~led_exposure_absolute", 50)),
            ("white_balance_temperature_auto", rospy.get_param("~led_white_balance_auto", True)),
        )
        self._restore_camera_controls = (
            ("exposure_auto", rospy.get_param("~restore_exposure_auto", 3)),
            ("exposure_auto_priority", rospy.get_param("~restore_exposure_auto_priority", 0)),
            ("white_balance_temperature_auto", rospy.get_param("~restore_white_balance_auto", True)),
            ("brightness", rospy.get_param("~restore_brightness", 0)),
            ("contrast", rospy.get_param("~restore_contrast", 50)),
            ("saturation", rospy.get_param("~restore_saturation", 50)),
            ("gamma", rospy.get_param("~restore_gamma", 300)),
        )
        self._camera_profile_active = False
        self._enabled = bool(rospy.get_param("~enabled", True))
        if self._camera_profile_enabled and self._enabled:
            self._apply_led_camera_profile()
        rospy.on_shutdown(self._release_runtime)
        rospy.on_shutdown(self._restore_camera_profile)
        atexit.register(self._restore_camera_profile)

        self._direction_pub = rospy.Publisher("/traffic_sign/direction", String, queue_size=1)
        self._confidence_pub = rospy.Publisher("/traffic_sign/confidence", Float32, queue_size=1)
        self._debug_pub = rospy.Publisher("/traffic_sign/debug_image", Image, queue_size=1)
        self._enable_service = rospy.Service(
            "~set_enabled", SetBool, self._set_enabled
        )
        image_topic = rospy.get_param("~image_topic", "/usb_cam/image_raw")
        self._image_sub = rospy.Subscriber(image_topic, Image, self._image_callback, queue_size=1)
        rospy.loginfo(
            "traffic_sign_recognition ready: model=%s image_topic=%s calibration=%s classes=%s enabled=%s",
            self._model_path,
            image_topic,
            self._calibration_file,
            ",".join(self._classes),
            self._enabled,
        )

    def _set_enabled(self, request) -> SetBoolResponse:
        """Enable inference/LED exposure or restore the shared camera profile."""
        requested = bool(request.data)
        with self._state_lock:
            if requested == self._enabled:
                # A previous disable request may have stopped inference but
                # failed midway through restoring V4L2 controls. Allow a
                # repeated disable request (for example from FAULT handling)
                # to retry that safety-critical restoration.
                if not requested and self._camera_profile_active:
                    try:
                        self._restore_camera_profile(raise_on_error=True)
                    except RuntimeError as error:
                        return SetBoolResponse(success=False, message=str(error))
                return SetBoolResponse(
                    success=True,
                    message="traffic sign recognition already %s"
                    % ("enabled" if requested else "disabled"),
                )

            if requested:
                try:
                    if self._camera_profile_enabled:
                        self._apply_led_camera_profile()
                except RuntimeError as error:
                    return SetBoolResponse(success=False, message=str(error))
                self._history.clear()
                self._last_direction = "NONE"
                self._enabled = True
                rospy.set_param("~enabled", True)
                self._direction_pub.publish(String(data="NONE"))
                self._confidence_pub.publish(Float32(data=0.0))
                rospy.loginfo("Traffic sign recognition enabled.")
                return SetBoolResponse(success=True, message="recognition enabled")

            # Mark disabled before restoring exposure so no new callback can
            # start inference using a half-restored camera frame.
            self._enabled = False
            self._history.clear()
            self._last_direction = "NONE"
            self._direction_pub.publish(String(data="NONE"))
            self._confidence_pub.publish(Float32(data=0.0))
            try:
                self._restore_camera_profile(raise_on_error=True)
            except RuntimeError as error:
                return SetBoolResponse(success=False, message=str(error))
            rospy.set_param("~enabled", False)
            rospy.loginfo("Traffic sign recognition disabled; camera profile restored.")
            return SetBoolResponse(success=True, message="recognition disabled")

    def _release_runtime(self) -> None:
        if getattr(self, "_rknn", None) is not None:
            self._rknn.release()
            self._rknn = None

    def _set_camera_controls(self, controls: Sequence[Tuple[str, object]]) -> None:
        """Set the shared USB camera controls without opening a second camera node."""
        for name, value in controls:
            command = [
                "v4l2-ctl",
                "-d",
                self._camera_device,
                "--set-ctrl=%s=%s" % (name, int(value) if isinstance(value, bool) else value),
            ]
            try:
                subprocess.run(command, check=True, capture_output=True, text=True, timeout=3)
            except (OSError, subprocess.SubprocessError) as error:
                raise RuntimeError("Cannot set camera control %s: %s" % (name, error))

    def _apply_led_camera_profile(self) -> None:
        # Mark active first so a partially-applied V4L2 sequence can still be
        # restored by a subsequent disable/fault request.
        self._camera_profile_active = True
        self._set_camera_controls(self._led_camera_controls)
        rospy.loginfo("LED camera profile enabled on %s: manual exposure=%s", self._camera_device,
                      dict(self._led_camera_controls)["exposure_absolute"])

    def _restore_camera_profile(self, raise_on_error: bool = False) -> None:
        if not getattr(self, "_camera_profile_active", False) or not self._restore_camera_on_shutdown:
            return
        try:
            self._set_camera_controls(self._restore_camera_controls)
            self._camera_profile_active = False
            rospy.loginfo("Automatic USB camera profile restored on %s", self._camera_device)
        except RuntimeError as error:
            rospy.logerr("Unable to restore automatic camera profile: %s", error)
            if raise_on_error:
                raise

    @staticmethod
    def _load_camera_calibration(
        calibration_file: Path,
    ) -> Tuple[np.ndarray, np.ndarray, Tuple[int, int]]:
        """Load the USB camera calibration and reject incomplete files early."""
        try:
            with calibration_file.open("r", encoding="utf-8") as stream:
                calibration = yaml.safe_load(stream)
            width = int(calibration["image_width"])
            height = int(calibration["image_height"])
            matrix = np.asarray(calibration["camera_matrix"]["data"], dtype=np.float64).reshape(3, 3)
            coefficients = np.asarray(
                calibration["distortion_coefficients"]["data"], dtype=np.float64
            ).reshape(-1, 1)
        except (OSError, KeyError, TypeError, ValueError, yaml.YAMLError) as error:
            raise RuntimeError("Unable to load camera calibration %s: %s" % (calibration_file, error))
        return matrix, coefficients, (width, height)

    def _preprocess_camera_frame(self, frame: np.ndarray) -> np.ndarray:
        """Preserve LED colour while applying the shared, versioned LED preprocessing."""
        processed = cv2.flip(frame, 1) if self._flip_horizontal else frame.copy()
        if self._contrast != 1.0 or self._brightness != 0.0:
            processed = cv2.convertScaleAbs(
                processed, alpha=self._contrast, beta=self._brightness
            )
        if self._saturation != 1.0:
            hsv = cv2.cvtColor(processed, cv2.COLOR_BGR2HSV)
            hsv[:, :, 1] = np.clip(
                hsv[:, :, 1].astype(np.float32) * self._saturation, 0, 255
            ).astype(np.uint8)
            processed = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)
        if self._lab_clahe:
            lab = cv2.cvtColor(processed, cv2.COLOR_BGR2LAB)
            lab[:, :, 0] = self._clahe.apply(lab[:, :, 0])
            processed = cv2.cvtColor(lab, cv2.COLOR_LAB2BGR)
        if self._gamma != 1.0:
            lookup = np.array(
                [((value / 255.0) ** self._gamma) * 255.0 for value in range(256)], dtype=np.uint8
            )
            processed = cv2.LUT(processed, lookup)
        return processed

    def _decode_output(self, outputs: Sequence[np.ndarray]) -> Tuple[Optional[np.ndarray], Optional[np.ndarray], Optional[np.ndarray]]:
        # ``export.py --include onnx`` creates YOLOv5's decoded output
        # [batch, 25200, 5 + class_count].  Some RKNN-export pipelines expose
        # the three raw detection heads instead, so keep support for both.
        if len(outputs) == 1:
            # RKNN Lite may retain a singleton batch dimension and/or append
            # a singleton stride-alignment dimension: (1, 25200, 9, 1).
            # Remove only dimensions of length one, yielding (25200, 9).
            values = np.squeeze(np.asarray(outputs[0]))
            if values.ndim == 2 and values.shape[0] == 5 + len(self._classes):
                values = values.T
            if values.ndim != 2 or values.shape[1] != 5 + len(self._classes):
                raise RuntimeError(
                    "Unsupported decoded YOLOv5 RKNN output shape: %s" % (values.shape,)
                )

            objectness = values[:, 4:5]
            class_scores = values[:, 5:] * objectness
            classes = np.argmax(class_scores, axis=1)
            scores = np.max(class_scores, axis=1)
            selected = scores >= self._object_threshold
            if not np.any(selected):
                return None, None, None

            xywh = values[selected, :4]
            boxes = np.concatenate((xywh[:, :2] - xywh[:, 2:] / 2.0,
                                    xywh[:, :2] + xywh[:, 2:] / 2.0), axis=1)
            return self._classwise_nms(boxes, classes[selected], scores[selected])

        if len(outputs) != 3:
            raise RuntimeError(
                "Expected one decoded or three raw YOLOv5 RKNN outputs, got %d" % len(outputs)
            )

        all_boxes, all_classes, all_scores = [], [], []
        for output, mask in zip(outputs, DEFAULT_MASKS):
            values = np.asarray(output)
            if values.ndim == 4 and values.shape[0] == 1:
                values = values[0]
            if values.ndim != 3:
                raise RuntimeError("Unsupported RKNN output shape: %s" % (values.shape,))
            if values.shape[0] % 3 == 0:
                values = values.reshape(3, -1, values.shape[1], values.shape[2]).transpose(2, 3, 0, 1)
            elif values.shape[-1] % 3 == 0:
                values = values.reshape(values.shape[0], values.shape[1], 3, -1)
            else:
                raise RuntimeError("Cannot interpret RKNN output shape: %s" % (values.shape,))

            grid_height, grid_width = values.shape[:2]
            anchors = DEFAULT_ANCHORS[list(mask)]
            confidence = sigmoid(values[..., 4:5])
            class_probabilities = sigmoid(values[..., 5:])
            class_scores = class_probabilities * confidence
            classes = np.argmax(class_scores, axis=-1)
            scores = np.max(class_scores, axis=-1)
            selected = scores >= self._object_threshold
            if not np.any(selected):
                continue

            grid_x, grid_y = np.meshgrid(np.arange(grid_width), np.arange(grid_height))
            grid = np.stack((grid_x, grid_y), axis=-1)[..., None, :]
            xy = (sigmoid(values[..., :2]) * 2.0 - 0.5 + grid) * (self._input_size / grid_height)
            wh = (sigmoid(values[..., 2:4]) * 2.0) ** 2 * anchors[None, None, :, :]
            xyxy = np.concatenate((xy - wh / 2.0, xy + wh / 2.0), axis=-1)
            all_boxes.append(xyxy[selected])
            all_classes.append(classes[selected])
            all_scores.append(scores[selected])

        if not all_boxes:
            return None, None, None

        boxes = np.concatenate(all_boxes)
        classes = np.concatenate(all_classes)
        scores = np.concatenate(all_scores)
        return self._classwise_nms(boxes, classes, scores)

    def _classwise_nms(
        self,
        boxes: np.ndarray,
        classes: np.ndarray,
        scores: np.ndarray,
    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
        """Suppress overlapping candidates independently for each class."""
        kept_boxes, kept_classes, kept_scores = [], [], []
        for class_index in np.unique(classes):
            class_indices = np.where(classes == class_index)[0]
            keep = nms_boxes(boxes[class_indices], scores[class_indices], self._nms_threshold)
            kept_boxes.append(boxes[class_indices][keep])
            kept_classes.append(classes[class_indices][keep])
            kept_scores.append(scores[class_indices][keep])
        return np.concatenate(kept_boxes), np.concatenate(kept_classes), np.concatenate(kept_scores)

    def _stable_direction(self, candidate: Optional[str]) -> str:
        self._history.append(candidate)
        votes = Counter(value for value in self._history if value is not None)
        if not votes:
            return "NONE"
        direction, count = votes.most_common(1)[0]
        return direction.upper() if count >= self._stable_count else "NONE"

    @staticmethod
    def _restore_boxes(boxes: np.ndarray, scale: float, padding: Tuple[float, float], width: int, height: int) -> np.ndarray:
        restored = boxes.copy()
        restored[:, [0, 2]] = (restored[:, [0, 2]] - padding[0]) / scale
        restored[:, [1, 3]] = (restored[:, [1, 3]] - padding[1]) / scale
        restored[:, [0, 2]] = np.clip(restored[:, [0, 2]], 0, width - 1)
        restored[:, [1, 3]] = np.clip(restored[:, [1, 3]], 0, height - 1)
        return restored

    def _image_callback(self, message: Image) -> None:
        with self._state_lock:
            if not self._enabled:
                return
            self._process_enabled_image(message)

    def _process_enabled_image(self, message: Image) -> None:
        try:
            camera_frame = self._bridge.imgmsg_to_cv2(message, desired_encoding="bgr8")
        except CvBridgeError as error:
            rospy.logerr_throttle(5.0, "traffic sign image conversion failed: %s", error)
            return

        expected_width, expected_height = self._calibration_size
        if camera_frame.shape[:2] != (expected_height, expected_width):
            rospy.logwarn_throttle(
                5.0,
                "traffic sign image size %dx%d differs from calibration %dx%d; frame skipped",
                camera_frame.shape[1],
                camera_frame.shape[0],
                expected_width,
                expected_height,
            )
            return

        undistorted = cv2.undistort(camera_frame, self._camera_matrix, self._distortion_coefficients)
        processed = self._preprocess_camera_frame(undistorted)
        model_input, scale, padding = letterbox(processed, self._input_size)
        model_input = cv2.cvtColor(model_input, cv2.COLOR_BGR2RGB)
        try:
            # The OpenCV image is HWC RGB.  State this explicitly instead of
            # relying on RKNN Lite's default (which can be NCHW for ONNX
            # models and yields near-zero detections with an HWC buffer).
            outputs = self._rknn.inference(inputs=[model_input], data_format="nhwc")
            boxes, classes, scores = self._decode_output(outputs)
        except Exception as error:
            rospy.logerr_throttle(5.0, "traffic sign inference failed: %s", error)
            return

        candidate, candidate_score = None, 0.0
        debug = processed.copy()
        if boxes is not None:
            boxes = self._restore_boxes(boxes, scale, padding, debug.shape[1], debug.shape[0])
            best_index = int(np.argmax(scores))
            candidate = self._classes[int(classes[best_index])]
            candidate_score = float(scores[best_index])
            for box, class_index, score in zip(boxes, classes, scores):
                label = self._classes[int(class_index)].upper()
                x1, y1, x2, y2 = box.astype(int)
                cv2.rectangle(debug, (x1, y1), (x2, y2), (0, 255, 0), 2)
                cv2.putText(debug, "%s %.2f" % (label, score), (x1, max(20, y1 - 6)),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)

        stable = self._stable_direction(candidate)
        self._last_direction = stable
        self._direction_pub.publish(String(data=stable))
        self._confidence_pub.publish(Float32(data=candidate_score if stable != "NONE" else 0.0))
        cv2.putText(debug, "STABLE: %s" % stable, (12, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.8,
                    (0, 255, 255), 2)
        try:
            self._debug_pub.publish(self._bridge.cv2_to_imgmsg(debug, encoding="bgr8"))
        except CvBridgeError as error:
            rospy.logerr_throttle(5.0, "traffic sign debug image publish failed: %s", error)

    def run(self) -> None:
        rospy.spin()


if __name__ == "__main__":
    try:
        TrafficSignRecognizer().run()
    except (RuntimeError, FileNotFoundError, ValueError) as error:
        rospy.logfatal("traffic_sign_recognition did not start: %s", error)
        raise