#!/usr/bin/env python3
"""
Terrain-aware cmd_vel speed shaper.

Sits between raw velocity command producers (Nav2's controller_server, remapped in
fori_nav2.launch.py, and parking_controller_node's own patrol/docking/mock-drive
commands, remapped in fori_full.launch.py) and the real /cmd_vel that
serial_bridge_node drives the motors from. Looks up the 2.5D elevation map
(pcd_gridmap_converter output) at the robot's current position and a short
lookahead point in the direction of travel, computes the local grade, and scales
the commanded forward speed: faster uphill (more torque margin needed), slower
downhill (safer descent). serial_bridge_node's own target_linear_speed clamp is
still the final hard safety ceiling regardless of this node's output.
"""
import math

import numpy as np
import rclpy
import yaml
from geometry_msgs.msg import PoseWithCovarianceStamped, Twist
from nav_msgs.msg import Odometry
from rclpy.node import Node
from rclpy.qos import DurabilityPolicy, QoSProfile, QoSReliabilityPolicy, qos_profile_sensor_data


def _read_pgm_token(f):
    token = b''
    while True:
        c = f.read(1)
        if not c:
            raise EOFError('Unexpected EOF while reading PGM header')
        if c in b' \t\n\r':
            if token:
                return token
            continue
        if c == b'#':
            while c not in (b'\n', b''):
                c = f.read(1)
            continue
        token += c


class ElevationMap:
    """Loads a pcd_gridmap_converter elevation_map.pgm + elevation_map.yaml pair.

    Pixel 0 means unknown. Pixel 1..255 linearly encodes height between
    min_elevation and max_elevation (see pcd_to_gridmap_node.cpp's
    publishAndSaveMaps()). Image rows are stored top-to-bottom for the highest-y
    grid row first, mirroring how nav2_map_server decodes map.pgm.
    """

    def __init__(self, pgm_path, yaml_path, logger=None):
        self.logger = logger
        self.valid = False
        try:
            with open(yaml_path, 'r') as f:
                meta = yaml.safe_load(f)
            self.resolution = float(meta['resolution'])
            origin = meta['origin']
            self.origin_x = float(origin[0])
            self.origin_y = float(origin[1])
            self.width = int(meta['width'])
            self.height = int(meta['height'])
            self.min_elevation = float(meta['min_elevation'])
            self.max_elevation = float(meta['max_elevation'])
            self.unknown_value = int(meta.get('unknown_value', 0))

            with open(pgm_path, 'rb') as f:
                magic = _read_pgm_token(f)
                if magic != b'P5':
                    raise ValueError(f'Unsupported PGM magic: {magic!r}')
                w = int(_read_pgm_token(f))
                h = int(_read_pgm_token(f))
                _read_pgm_token(f)  # maxval, assumed 255
                data = f.read(w * h)

            if len(data) != w * h:
                raise ValueError('Truncated PGM pixel data')
            if (w, h) != (self.width, self.height):
                raise ValueError(f'PGM size {w}x{h} does not match yaml size {self.width}x{self.height}')

            self.pixels = np.frombuffer(data, dtype=np.uint8).reshape((h, w))
            self.valid = True
        except Exception as e:  # noqa: BLE001 - any load failure -> safe passthrough mode
            if self.logger:
                self.logger.error(f'[ELEVATION] 고도 지도 로드 실패 ({pgm_path}): {e}')

    def lookup(self, x, y):
        if not self.valid:
            return None
        col = int(math.floor((x - self.origin_x) / self.resolution))
        grid_y = int(math.floor((y - self.origin_y) / self.resolution))
        if col < 0 or col >= self.width or grid_y < 0 or grid_y >= self.height:
            return None
        image_row = (self.height - 1) - grid_y
        pixel = int(self.pixels[image_row, col])
        if pixel == self.unknown_value:
            return None
        return self.min_elevation + (pixel - 1) / 254.0 * (self.max_elevation - self.min_elevation)


class TerrainSpeedNode(Node):
    def __init__(self):
        super().__init__('terrain_speed_node')

        self.declare_parameter('elevation_pgm', '/home/yoo/pcd_gunsan_output/elevation_map.pgm')
        self.declare_parameter('elevation_yaml', '/home/yoo/pcd_gunsan_output/elevation_map.yaml')
        # 기본값은 /cmd_vel (목 모드 등 collision_monitor가 없는 구성에서 안전한 기본).
        # 실기체 모드에서는 fori_full.launch.py가 이 값을 /cmd_vel_terrain으로 넘겨서
        # collision_monitor(급출현 장애물 감속/정지)가 마지막 안전 게이트로 끼어들게 한다.
        self.declare_parameter('output_topic', '/cmd_vel')
        self.declare_parameter('lookahead_dist', 0.5)
        self.declare_parameter('uphill_gain', 3.0)
        self.declare_parameter('downhill_gain', 4.0)
        self.declare_parameter('max_boost', 1.3)
        self.declare_parameter('min_scale', 0.5)
        self.declare_parameter('flat_deadband', 0.02)
        self.declare_parameter('smoothing_alpha', 0.25)

        self.lookahead_dist = self.get_parameter('lookahead_dist').value
        self.uphill_gain = self.get_parameter('uphill_gain').value
        self.downhill_gain = self.get_parameter('downhill_gain').value
        self.max_boost = self.get_parameter('max_boost').value
        self.min_scale = self.get_parameter('min_scale').value
        self.flat_deadband = self.get_parameter('flat_deadband').value
        self.smoothing_alpha = self.get_parameter('smoothing_alpha').value

        pgm_path = self.get_parameter('elevation_pgm').value
        yaml_path = self.get_parameter('elevation_yaml').value
        self.elevation = ElevationMap(pgm_path, yaml_path, logger=self.get_logger())
        if self.elevation.valid:
            self.get_logger().info(
                f'[TERRAIN] 고도 지도 로드 완료: {pgm_path} '
                f'({self.elevation.width}x{self.elevation.height}px, '
                f'{self.elevation.min_elevation:.2f}m ~ {self.elevation.max_elevation:.2f}m)'
            )
        else:
            self.get_logger().warn(
                '[TERRAIN] 고도 지도를 사용할 수 없어 경사 보정 없이 /cmd_vel_raw를 그대로 통과시킵니다.'
            )

        self.pose_x = 0.0
        self.pose_y = 0.0
        self.pose_yaw = 0.0
        self.last_amcl_time = None
        self.scale = 1.0

        # AMCL(맵 기준 정밀 위치)을 우선 사용하고, 3초 이상 갱신이 없으면 휠 오도메트리로
        # 대체한다 - parking_controller_node.py의 amcl_pose_callback/odom_callback과 동일한
        # freshness 기준.
        amcl_qos = QoSProfile(depth=5,
                               reliability=QoSReliabilityPolicy.BEST_EFFORT,
                               durability=DurabilityPolicy.VOLATILE)
        self.create_subscription(PoseWithCovarianceStamped, '/amcl_pose', self.amcl_pose_callback, amcl_qos)
        self.create_subscription(Odometry, '/odom_wheels', self.odom_callback, qos_profile_sensor_data)
        self.create_subscription(Twist, '/cmd_vel_raw', self.cmd_vel_callback, 10)
        output_topic = self.get_parameter('output_topic').value
        self.cmd_vel_pub = self.create_publisher(Twist, output_topic, 10)
        self.get_logger().info(f'[TERRAIN] 출력 토픽: {output_topic}')

    def _update_pose(self, x, y, q):
        self.pose_x = x
        self.pose_y = y
        siny_cosp = 2 * (q.w * q.z + q.x * q.y)
        cosy_cosp = 1 - 2 * (q.y * q.y + q.z * q.z)
        self.pose_yaw = math.atan2(siny_cosp, cosy_cosp)

    def amcl_pose_callback(self, msg):
        self.last_amcl_time = self.get_clock().now()
        self._update_pose(msg.pose.pose.position.x, msg.pose.pose.position.y, msg.pose.pose.orientation)

    def odom_callback(self, msg):
        if self.last_amcl_time is not None:
            elapsed = (self.get_clock().now() - self.last_amcl_time).nanoseconds / 1e9
            if elapsed < 3.0:
                return
        self._update_pose(msg.pose.pose.position.x, msg.pose.pose.position.y, msg.pose.pose.orientation)

    def cmd_vel_callback(self, msg):
        out = Twist()
        out.angular.z = msg.angular.z

        # 제자리 회전이나 지도가 없는 경우엔 경사 보정 없이 그대로 통과 (안전 기본값)
        if not self.elevation.valid or abs(msg.linear.x) < 0.02:
            out.linear.x = msg.linear.x
            self.scale = 1.0
            self.cmd_vel_pub.publish(out)
            return

        direction = 1.0 if msg.linear.x >= 0.0 else -1.0
        to_x = self.pose_x + direction * math.cos(self.pose_yaw) * self.lookahead_dist
        to_y = self.pose_y + direction * math.sin(self.pose_yaw) * self.lookahead_dist

        z_from = self.elevation.lookup(self.pose_x, self.pose_y)
        z_to = self.elevation.lookup(to_x, to_y)

        if z_from is None or z_to is None:
            raw_scale = 1.0
        else:
            grade = (z_to - z_from) / self.lookahead_dist
            if abs(grade) < self.flat_deadband:
                raw_scale = 1.0
            elif grade > 0.0:
                raw_scale = 1.0 + self.uphill_gain * grade   # 오르막: 속도 부스트
            else:
                raw_scale = 1.0 + self.downhill_gain * grade  # 내리막: 속도 감쇠 (grade<0)
            raw_scale = max(self.min_scale, min(self.max_boost, raw_scale))

        # 셀 경계를 지날 때 속도가 튀지 않도록 EMA로 스무딩
        self.scale = self.scale * (1.0 - self.smoothing_alpha) + raw_scale * self.smoothing_alpha
        out.linear.x = msg.linear.x * self.scale
        self.cmd_vel_pub.publish(out)


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


if __name__ == '__main__':
    main()
