#!/usr/bin/env python3
"""
terrain_classification_node.py — Terrain classification (Issue #469).

Fuses RealSense D435i depth + RPLIDAR A1M8 to classify terrain and recommend speed.

Pipeline (10 Hz)
────────────────
  1. Extract depth-based roughness features
  2. Extract lidar-based reflectance features
  3. Classify terrain using rule-based matcher
  4. Publish /saltybot/terrain_type (JSON + string)
  5. Publish speed recommendations

Publishes
─────────
  /saltybot/terrain_type         std_msgs/String  — JSON: type, confidence, speed_scale
  /saltybot/terrain_type_string  std_msgs/String  — Human-readable type name
  /saltybot/terrain_speed_scale  std_msgs/Float32 — Speed multiplier [0.0, 1.0]
"""

import json
import time
import numpy as np

import rclpy
from rclpy.node import Node
from rclpy.qos import QoSProfile, ReliabilityPolicy, HistoryPolicy

from sensor_msgs.msg import Image, LaserScan
from std_msgs.msg import String, Float32
import cv2
from cv_bridge import CvBridge

from saltybot_terrain_classification.terrain_classifier import TerrainClassifier
from saltybot_terrain_classification.depth_extractor import DepthExtractor
from saltybot_terrain_classification.lidar_extractor import LidarExtractor


class TerrainClassificationNode(Node):

    def __init__(self):
        super().__init__("terrain_classification")

        self._declare_params()
        p = self._load_params()

        self._classifier = TerrainClassifier(
            depth_std_threshold=p["depth_std_threshold"],
            hysteresis_count=p["hysteresis_samples"],
        )
        self._depth_extractor = DepthExtractor(
            roi_width_px=320,
            roi_height_px=240,
            min_depth_m=p["min_depth_m"],
            max_depth_m=p["max_depth_m"],
        )
        self._lidar_extractor = LidarExtractor()
        self._cv_bridge = CvBridge()

        self._latest_depth_image = None
        self._latest_scan = None
        self._last_depth_t = 0.0
        self._last_scan_t = 0.0

        depth_qos = QoSProfile(
            reliability=ReliabilityPolicy.BEST_EFFORT,
            history=HistoryPolicy.KEEP_LAST,
            depth=1,
        )
        lidar_qos = QoSProfile(
            reliability=ReliabilityPolicy.BEST_EFFORT,
            history=HistoryPolicy.KEEP_LAST,
            depth=1,
        )
        pub_qos = QoSProfile(
            reliability=ReliabilityPolicy.RELIABLE,
            history=HistoryPolicy.KEEP_LAST,
            depth=10,
        )

        self.create_subscription(
            Image, p["depth_topic"], self._depth_cb, depth_qos
        )
        self.create_subscription(
            LaserScan, p["lidar_topic"], self._scan_cb, lidar_qos
        )

        self._terrain_type_pub = self.create_publisher(
            String, "/saltybot/terrain_type", pub_qos
        )
        self._terrain_type_string_pub = self.create_publisher(
            String, "/saltybot/terrain_type_string", pub_qos
        )
        self._speed_scale_pub = self.create_publisher(
            Float32, "/saltybot/terrain_speed_scale", pub_qos
        )

        rate = p["control_rate"]
        self._timer = self.create_timer(1.0 / rate, self._control_cb)

        self.get_logger().info(
            f"TerrainClassificationNode ready  rate={rate}Hz  "
            f"depth={p['depth_topic']}  lidar={p['lidar_topic']}"
        )

    def _declare_params(self) -> None:
        self.declare_parameter("control_rate", 10.0)
        self.declare_parameter("depth_topic", "/camera/depth/image_rect_raw")
        self.declare_parameter("lidar_topic", "/scan")
        self.declare_parameter("depth_std_threshold", 0.02)
        self.declare_parameter("min_depth_m", 0.1)
        self.declare_parameter("max_depth_m", 3.0)
        self.declare_parameter("confidence_threshold", 0.5)
        self.declare_parameter("hysteresis_samples", 5)

    def _load_params(self) -> dict:
        g = self.get_parameter
        return {k: g(k).value for k in [
            "control_rate", "depth_topic", "lidar_topic",
            "depth_std_threshold", "min_depth_m", "max_depth_m",
            "confidence_threshold", "hysteresis_samples",
        ]}

    def _depth_cb(self, msg: Image) -> None:
        self._latest_depth_image = msg
        self._last_depth_t = time.monotonic()

    def _scan_cb(self, msg: LaserScan) -> None:
        self._latest_scan = msg
        self._last_scan_t = time.monotonic()

    def _control_cb(self) -> None:
        p = self._load_params()
        now = time.monotonic()

        depth_age = (now - self._last_depth_t) if self._last_depth_t > 0.0 else 1e9
        scan_age = (now - self._last_scan_t) if self._last_scan_t > 0.0 else 1e9

        if depth_age > 1.0 or scan_age > 1.0:
            return

        depth_roughness = None
        lidar_reflectance = None

        if self._latest_depth_image is not None:
            try:
                depth_cv = self._cv_bridge.imgmsg_to_cv2(
                    self._latest_depth_image, desired_encoding="passthrough"
                )
                depth_roughness = self._depth_extractor.extract_roughness(depth_cv)
            except Exception:
                pass

        if self._latest_scan is not None:
            try:
                ranges = np.array(self._latest_scan.ranges, dtype=np.float32)
                intensities = np.array(self._latest_scan.intensities, dtype=np.float32)
                lidar_reflectance = self._lidar_extractor.extract_reflectance(
                    ranges, intensities
                )
            except Exception:
                pass

        if depth_roughness is None:
            depth_roughness = 0.3
        if lidar_reflectance is None:
            lidar_reflectance = 0.5

        classification = self._classifier.update(depth_roughness, lidar_reflectance)

        if classification.confidence < p["confidence_threshold"]:
            return

        self._publish_terrain_type(classification)
        self._publish_speed_scale(classification.speed_scale)

        self.get_logger().info(
            f"Terrain: {classification.surface_type}  "
            f"confidence={classification.confidence:.2f}  "
            f"speed_scale={classification.speed_scale:.2f}"
        )

    def _publish_terrain_type(self, classification) -> None:
        """Publish terrain classification JSON."""
        msg = String()
        msg.data = json.dumps({
            "surface_type": classification.surface_type,
            "confidence": round(classification.confidence, 3),
            "roughness": round(classification.roughness, 3),
            "reflectance": round(classification.reflectance, 3),
            "speed_scale": round(classification.speed_scale, 3),
        })
        self._terrain_type_pub.publish(msg)

        msg_str = String()
        msg_str.data = classification.surface_type
        self._terrain_type_string_pub.publish(msg_str)

    def _publish_speed_scale(self, scale: float) -> None:
        """Publish speed scale multiplier."""
        msg = Float32()
        msg.data = float(scale)
        self._speed_scale_pub.publish(msg)


def main(args=None):
    rclpy.init(args=args)
    node = TerrainClassificationNode()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.try_shutdown()


if __name__ == "__main__":
    main()
