#!/usr/bin/env python3
"""Interactive raw-frame collector for the LED direction-sign dataset.

The collector never opens /dev/video0.  It consumes the common ROS image
topic, previews the same colour-preserving preprocessing used by the future
recogniser, and saves the processed training frame for every capture.
"""

from __future__ import print_function

import csv
import threading
import time
from pathlib import Path

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


CLASS_KEYS = {
    ord("0"): "negative",
    ord("1"): "left",
    ord("2"): "right",
    ord("3"): "straight",
    ord("4"): "stop",
}


class LedDatasetCollector(object):
    def __init__(self):
        rospy.init_node("led_dataset_collector")
        self._bridge = CvBridge()
        self._package_dir = Path(__file__).resolve().parent.parent
        calibration_default = self._package_dir / "config" / "head_camera.yaml"
        self._calibration_file = Path(
            rospy.get_param("~calibration_file", str(calibration_default))
        ).expanduser()
        self._camera_matrix, self._distortion, self._calibration_size = self._load_calibration()

        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._use_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._camera_profile = rospy.get_param("~camera_profile", "usb-cam-auto-exposure-auto-white-balance")
        self._data_root = Path(
            rospy.get_param("~data_root", str(Path.home() / "traffic_sign_data"))
        ).expanduser()
        self._processed_root = self._data_root / "processed"
        # Keep new processed-only records separate from the earlier raw-pair log.
        self._metadata_path = self._data_root / "processed_metadata.csv"
        for label in CLASS_KEYS.values():
            (self._processed_root / label).mkdir(parents=True, exist_ok=True)

        self._selected_label = "left"
        self._lock = threading.Lock()
        self._latest_raw = None
        self._latest_processed = None
        self._latest_stamp_ns = 0
        self._write_metadata_header()

        self._preview_pub = rospy.Publisher(
            "/traffic_sign/collector/processed_image", Image, queue_size=1
        )
        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(
            "LED collector ready: image_topic=%s processed_root=%s keys: 0=negative 1=left 2=right 3=straight 4=stop s=save q=quit",
            image_topic,
            self._processed_root,
        )

    def _load_calibration(self):
        try:
            with self._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)
            distortion = 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" % (self._calibration_file, error))
        return matrix, distortion, (width, height)

    def _write_metadata_header(self):
        self._data_root.mkdir(parents=True, exist_ok=True)
        if self._metadata_path.exists():
            return
        with self._metadata_path.open("w", newline="", encoding="utf-8") as stream:
            csv.writer(stream).writerow(
                [
                    "processed_filename",
                    "label",
                    "saved_at_utc",
                    "camera_profile",
                    "calibration_file",
                    "flip_horizontal",
                    "lab_clahe",
                    "clahe_clip_limit",
                    "gamma",
                ]
            )

    def _process(self, raw):
        undistorted = cv2.undistort(raw, self._camera_matrix, self._distortion)
        processed = cv2.flip(undistorted, 1) if self._flip_horizontal else undistorted
        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._use_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 _image_callback(self, message):
        try:
            raw = self._bridge.imgmsg_to_cv2(message, desired_encoding="bgr8")
        except CvBridgeError as error:
            rospy.logerr_throttle(5.0, "collector image conversion failed: %s", error)
            return
        width, height = self._calibration_size
        if raw.shape[:2] != (height, width):
            rospy.logwarn_throttle(
                5.0,
                "collector image size %dx%d differs from calibration %dx%d; frame skipped",
                raw.shape[1], raw.shape[0], width, height,
            )
            return
        processed = self._process(raw)
        with self._lock:
            self._latest_raw = raw.copy()
            self._latest_processed = processed
            self._latest_stamp_ns = message.header.stamp.to_nsec()
        try:
            self._preview_pub.publish(self._bridge.cv2_to_imgmsg(processed, encoding="bgr8"))
        except CvBridgeError as error:
            rospy.logerr_throttle(5.0, "collector preview publish failed: %s", error)

    def _save_current_frame(self):
        with self._lock:
            if self._latest_processed is None:
                rospy.logwarn("No camera frame received; nothing saved")
                return
            processed = self._latest_processed.copy()
            stamp_ns = self._latest_stamp_ns or time.time_ns()
        filename = "%s_%019d.png" % (self._selected_label, stamp_ns)
        processed_destination = self._processed_root / self._selected_label / filename
        if not cv2.imwrite(str(processed_destination), processed):
            rospy.logerr("Failed to save processed frame: %s", processed_destination)
            return
        with self._metadata_path.open("a", newline="", encoding="utf-8") as stream:
            csv.writer(stream).writerow(
                [
                    str(processed_destination.relative_to(self._data_root)),
                    self._selected_label,
                    time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
                    self._camera_profile,
                    str(self._calibration_file),
                    self._flip_horizontal,
                    self._use_lab_clahe,
                    self._clahe.getClipLimit(),
                    self._gamma,
                ]
            )
        rospy.loginfo(
            "Saved processed %s frame: %s",
            self._selected_label,
            processed_destination,
        )

    def _display(self, raw, processed):
        preview = np.hstack((raw, processed))
        cv2.putText(preview, "RAW", (12, 28), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 255), 2)
        offset = raw.shape[1]
        cv2.putText(preview, "LED PREPROCESS", (offset + 12, 28), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 255), 2)
        cv2.putText(
            preview,
            "label=%s | 0:negative 1:left 2:right 3:straight 4:stop | s:save q:quit" % self._selected_label,
            (12, preview.shape[0] - 14),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.55,
            (0, 255, 255),
            2,
        )
        cv2.imshow("LED dataset collector", preview)

    def run(self):
        cv2.namedWindow("LED dataset collector", cv2.WINDOW_NORMAL)
        while not rospy.is_shutdown():
            with self._lock:
                raw = None if self._latest_raw is None else self._latest_raw.copy()
                processed = None if self._latest_processed is None else self._latest_processed.copy()
            if raw is not None:
                self._display(raw, processed)
            key = cv2.waitKey(20) & 0xFF
            if key in CLASS_KEYS:
                self._selected_label = CLASS_KEYS[key]
                rospy.loginfo("Collector label selected: %s", self._selected_label)
            elif key in (ord("s"), ord("S")):
                self._save_current_frame()
            elif key in (ord("q"), ord("Q"), 27):
                break
        cv2.destroyAllWindows()


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