#!/usr/bin/env python3
"""示教程序 (teach_points)：记录机器人在 map 坐标系下的位姿，存为 YAML 供导航复用。

依赖定位栈在运行(localization.launch.py 提供 AMCL)，位姿来源为 TF: map -> base_footprint
(= AMCL 实时定位结果)。用法是另开终端用 keyboard_control 把车开到目标点，回到本程序按回车记录。

预设要记录的点(可用参数 points 覆盖)：B_1..B_6, C_1, D_1, E_1, F_1。
也可随时用 `name <名字>` 记录任意自定义点。

运行：
  # 终端1：定位栈
  ros2 launch craic_localization localization.launch.py map:=.../craic.yaml
  # 终端2：遥控(开车到点)
  ros2 run udp_teleop keyboard_control --ros-args --params-file src/udp_teleop/config/params.yaml
  # 终端3：示教(记录)，建议把结果存到包的 config 下供导航读取
  ros2 run craic_localization teach_points --ros-args \
      -p output_file:=$PWD/src/craic_localization/config/taught_points.yaml

交互命令：
  回车 / r      记录"当前待记录点"的位姿，并前进到下一个
  p             只打印当前位姿，不记录
  name <X>      把当前位姿记录为自定义名字 X(已存在则覆盖)
  del <X>       删除已记录的点 X
  list          列出已记录的点
  skip / back   跳过 / 回退当前待记录点
  save          立即保存到文件
  h             帮助
  q             保存并退出
"""

import math
import os
import threading

import rclpy
import yaml
from rclpy.duration import Duration
from rclpy.node import Node
from rclpy.time import Time
from tf2_ros import (
    Buffer,
    ConnectivityException,
    ExtrapolationException,
    LookupException,
    TransformListener,
)

DEFAULT_POINTS = ['B_1', 'B_2', 'B_3', 'B_4', 'B_5', 'B_6', 'C_1', 'D_1', 'E_1', 'F_1']


def quat_to_yaw(qx, qy, qz, qw):
    """四元数 -> 绕 Z 偏航角(弧度)。"""
    return math.atan2(2.0 * (qw * qz + qx * qy), 1.0 - 2.0 * (qy * qy + qz * qz))


class TeachPoints(Node):
    def __init__(self):
        super().__init__('teach_points')
        self.map_frame = self.declare_parameter('map_frame', 'map').value
        self.base_frame = self.declare_parameter('base_frame', 'base_footprint').value
        self.output_file = os.path.abspath(
            self.declare_parameter('output_file', 'taught_points.yaml').value)
        self.points = list(self.declare_parameter('points', DEFAULT_POINTS).value)

        self.tf_buffer = Buffer()
        self.tf_listener = TransformListener(self.tf_buffer, self)

        self.recorded = {}  # name -> {x, y, yaw}
        self._load_existing()

    # -------- 位姿读取 -------- #
    def lookup_pose(self):
        """读取 map->base_frame 的最新变换，返回 ({x,y,yaw}, None) 或 (None, 错误信息)。"""
        try:
            t = self.tf_buffer.lookup_transform(
                self.map_frame, self.base_frame, Time(),
                timeout=Duration(seconds=1.0))
        except (LookupException, ConnectivityException, ExtrapolationException) as exc:
            return None, str(exc)
        tr = t.transform.translation
        q = t.transform.rotation
        yaw = quat_to_yaw(q.x, q.y, q.z, q.w)
        return {'x': round(tr.x, 4), 'y': round(tr.y, 4), 'yaw': round(yaw, 4)}, None

    # -------- 文件读写 -------- #
    def _load_existing(self):
        if os.path.isfile(self.output_file):
            try:
                with open(self.output_file, 'r', encoding='utf-8') as f:
                    data = yaml.safe_load(f) or {}
                self.recorded = dict(data.get('points', {}))
                if self.recorded:
                    print(f'已载入已有记录 {len(self.recorded)} 个点：{self.output_file}')
            except Exception as exc:  # noqa: BLE001
                print(f'载入已有文件失败(忽略)：{exc}')

    def save(self):
        os.makedirs(os.path.dirname(self.output_file) or '.', exist_ok=True)
        header = (
            '# 示教记录的点位（map 坐标系；x,y 单位米，yaw 单位弧度）\n'
            '# 由 craic_localization/teach_points 生成\n'
        )
        body = yaml.safe_dump({'points': self.recorded},
                              default_flow_style=False, sort_keys=False,
                              allow_unicode=True)
        with open(self.output_file, 'w', encoding='utf-8') as f:
            f.write(header + body)


HELP = __doc__.split('交互命令：', 1)[1]


def interactive(node: TeachPoints):
    idx = 0
    print('=' * 60)
    print('  示教程序 teach_points —— 记录 map 系下位姿')
    print(f'  输出文件: {node.output_file}')
    print(f'  预设点: {", ".join(node.points)}')
    print('  (输入 h 查看命令)')
    print('=' * 60)

    while rclpy.ok():
        target = node.points[idx] if idx < len(node.points) else None
        tag = target if target else '(预设已完，用 name <X> 记录自定义点)'
        try:
            line = input(f'[已记录 {len(node.recorded)}] 下一个: {tag} > ').strip()
        except EOFError:
            break
        cmd = line.lower()

        if cmd in ('q', 'quit', 'exit'):
            node.save()
            print(f'已保存到 {node.output_file}，退出。')
            break

        elif cmd in ('', 'r'):
            if target is None:
                print('  预设点已记录完。用 "name <X>" 记录自定义点，或 q 退出。')
                continue
            pose, err = node.lookup_pose()
            if err:
                print(f'  ✗ 读取位姿失败(定位栈在运行吗?): {err}')
                continue
            node.recorded[target] = pose
            node.save()
            print(f'  ✓ {target} = x:{pose["x"]} y:{pose["y"]} '
                  f'yaw:{pose["yaw"]}rad ({math.degrees(pose["yaw"]):.1f}°)  [已存盘]')
            idx += 1

        elif cmd == 'p':
            pose, err = node.lookup_pose()
            if err:
                print(f'  ✗ 读取位姿失败: {err}')
            else:
                print(f'  当前位姿 x:{pose["x"]} y:{pose["y"]} '
                      f'yaw:{pose["yaw"]}rad ({math.degrees(pose["yaw"]):.1f}°)')

        elif cmd.startswith('name '):
            name = line[5:].strip()
            if not name:
                print('  用法: name <名字>')
                continue
            pose, err = node.lookup_pose()
            if err:
                print(f'  ✗ 读取位姿失败: {err}')
                continue
            node.recorded[name] = pose
            node.save()
            print(f'  ✓ {name} = x:{pose["x"]} y:{pose["y"]} '
                  f'yaw:{pose["yaw"]}rad  [已存盘]')

        elif cmd.startswith('del '):
            name = line[4:].strip()
            if node.recorded.pop(name, None) is not None:
                node.save()
                print(f'  已删除 {name}  [已存盘]')
            else:
                print(f'  没有名为 {name} 的记录')

        elif cmd == 'list':
            if not node.recorded:
                print('  (空)')
            for k, v in node.recorded.items():
                print(f'    {k}: x:{v["x"]} y:{v["y"]} yaw:{v["yaw"]}')

        elif cmd == 'skip':
            idx += 1

        elif cmd == 'back':
            idx = max(0, idx - 1)

        elif cmd == 'save':
            node.save()
            print(f'  已保存到 {node.output_file}')

        elif cmd in ('h', 'help', '?'):
            print(HELP)

        else:
            print('  未知命令，输入 h 查看帮助。')


def main(args=None):
    rclpy.init(args=args)
    node = TeachPoints()
    # 后台 spin 以填充 TF 缓冲
    spin_thread = threading.Thread(target=rclpy.spin, args=(node,), daemon=True)
    spin_thread.start()
    try:
        interactive(node)
    except KeyboardInterrupt:
        node.save()
        print(f'\n中断，已保存到 {node.output_file}')
    finally:
        node.destroy_node()
        if rclpy.ok():
            rclpy.shutdown()


if __name__ == '__main__':
    main()
