#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""底盘标定工具：直线 / 横移 / 原地旋转。

用法（真机）：

    ros2 run stardust_clean_robot_base calibrate_chassis
    ros2 run stardust_clean_robot_base calibrate_chassis --mode linear \\
        --speed 0.2 --distance 3.0 --yes
    ros2 run stardust_clean_robot_base calibrate_chassis --mode yaw \\
        --speed 0.5 --turns 0.5 --arm 1.0

为什么只需要标两个量（见 src/chassis/kinematics/mecanum.cpp）：

    k_lin = 实际位移 / odom 位移   -> 修正 2*pi*r/G（wheel_radius 或 gear_ratio）
    k_ang = 实际转角 / odom 转角   -> 修正 (wheel_separation+axle_separation)/2

odom 只是把编码器计数乘上这两个尺度，所以标定必须用**地面实测值**当基准，
不能拿指令速度 x 时间当基准（那会把"驱动器没跟到指令"和"结构参数错"混在一起）。

注意：几何参数只在 chassis 节点启动时读取，改完参数要重启节点。
"""

import argparse
import json
import math
import sys
import time

import rclpy
from geometry_msgs.msg import Twist
from nav_msgs.msg import Odometry
from rcl_interfaces.msg import ParameterType
from rcl_interfaces.srv import GetParameters
from std_msgs.msg import String

DEFAULT_CMD_TOPIC = '/cmd_vel'
DEFAULT_ODOM_TOPIC = '/robot_base/wheel/odom'
DEFAULT_STATUS_TOPIC = '/robot_base/chassis/cmd_status'
DEFAULT_CHASSIS_NODE = '/robot_base/chassis'

# 读不到节点参数时用的兜底值（= params/chassis/mecanum_evt3.yaml）
FALLBACK_GEOMETRY = {
    'wheel_radius': 0.076,
    'gear_ratio': 25.0,
    'wheel_separation': 0.335,
    'axle_separation': 0.330,
}

MODE_LINEAR = 'linear'
MODE_LATERAL = 'lateral'
MODE_YAW = 'yaw'


# ── 小工具 ──────────────────────────────────────────────────────────


def wrap_angle(a):
    """把角度折到 [-pi, pi)。"""
    return math.atan2(math.sin(a), math.cos(a))


def yaw_from_quaternion(q):
    """四元数 -> yaw [rad]。"""
    siny = 2.0 * (q.w * q.z + q.x * q.y)
    cosy = 1.0 - 2.0 * (q.y * q.y + q.z * q.z)
    return math.atan2(siny, cosy)


def ask(prompt, default=None, cast=float):
    """交互问答；直接回车用默认值。"""
    suffix = '' if default is None else ' ({}): '.format(default)
    while True:
        try:
            raw = input(prompt + suffix).strip()
        except EOFError:
            raise SystemExit('\n输入结束，退出。')
        if not raw:
            if default is not None:
                return default
            print('  请输入一个值。')
            continue
        try:
            return cast(raw)
        except ValueError:
            print('  无法解析为数字，请重试。')


def ask_text(prompt, default, choices=('y', 'n')):
    """有限选项的交互问答。"""
    while True:
        try:
            raw = input('{} ({}): '.format(prompt, default)).strip().lower()
        except EOFError:
            raise SystemExit('\n输入结束，退出。')
        if not raw:
            return default
        if raw in choices:
            return raw
        print('  请从 {} 里选。'.format('/'.join(choices)))


# ── 运动 + 记录 ─────────────────────────────────────────────────────


class ChassisCalibrator(object):
    def __init__(self, node, args):
        self.node = node
        self.args = args
        self.rate = max(args.rate, 1.0)
        self.period = 1.0 / self.rate

        self.pub = node.create_publisher(Twist, args.cmd_topic, 10)
        self.sub = node.create_subscription(
            Odometry, args.odom_topic, self._on_odom, 10)
        self.status_sub = node.create_subscription(
            String, args.status_topic, self._on_status, 10)

        self.odom = None          # 最近一帧 odom
        self.odom_count = 0
        self.status = None        # 最近一次 chassis cmd_status(JSON)
        self.status_count = 0

        # 单次运行的累积量
        self.yaw_acc = 0.0
        self.yaw_prev = None
        self.vx_sum = 0.0
        self.vy_sum = 0.0
        self.vx_n = 0
        self.wz_sum = 0.0
        self.wz_n = 0

    # ── 回调 ──
    def _on_odom(self, msg):
        self.odom = msg
        self.odom_count += 1
        yaw = yaw_from_quaternion(msg.pose.pose.orientation)
        if self.yaw_prev is not None:
            self.yaw_acc += wrap_angle(yaw - self.yaw_prev)
        self.yaw_prev = yaw
        if self.recording:
            self.vx_sum += msg.twist.twist.linear.x
            self.vy_sum += msg.twist.twist.linear.y
            self.vx_n += 1
            self.wz_sum += msg.twist.twist.angular.z
            self.wz_n += 1

    def _on_status(self, msg):
        self.status_count += 1
        try:
            self.status = json.loads(msg.data)
        except ValueError:
            self.status = None

    # ── 等待 ──
    def spin_for(self, seconds):
        end = time.monotonic() + seconds
        while time.monotonic() < end:
            rclpy.spin_once(self.node, timeout_sec=0.01)

    def wait_odom(self, timeout=5.0):
        end = time.monotonic() + timeout
        while time.monotonic() < end:
            rclpy.spin_once(self.node, timeout_sec=0.05)
            if self.odom is not None:
                return True
        return False

    def wait_status(self, timeout=2.0):
        end = time.monotonic() + timeout
        while time.monotonic() < end:
            rclpy.spin_once(self.node, timeout_sec=0.05)
            if self.status is not None:
                return True
        return False

    # ── 发布 ──
    def publish_zero(self, seconds=0.5):
        zero = Twist()
        end = time.monotonic() + seconds
        while time.monotonic() < end:
            self.pub.publish(zero)
            rclpy.spin_once(self.node, timeout_sec=self.period)

    # ── 参数读取 ──
    def fetch_geometry(self):
        geometry = dict(FALLBACK_GEOMETRY)
        cli = self.node.create_client(
            GetParameters, self.args.chassis_node + '/get_parameters')
        if not cli.wait_for_service(timeout_sec=2.0):
            print('  [提示] 未连上 {} 的参数服务，用兜底几何值。'.format(
                self.args.chassis_node))
            return geometry, False
        req = GetParameters.Request()
        names = ['chassis.wheel_radius', 'chassis.gear_ratio',
                 'chassis.wheel_separation', 'chassis.axle_separation']
        req.names = names
        fut = cli.call_async(req)
        rclpy.spin_until_future_complete(self.node, fut, timeout_sec=2.0)
        res = fut.result()
        if res is None or len(res.values) != len(names):
            print('  [提示] 读参数失败，用兜底几何值。')
            return geometry, False
        for name, value in zip(names, res.values):
            key = name.split('.', 1)[1]
            if value.type == ParameterType.PARAMETER_DOUBLE:
                geometry[key] = value.double_value
            elif value.type == ParameterType.PARAMETER_INTEGER:
                geometry[key] = float(value.integer_value)
            else:
                print('  [提示] {} 未设置，用兜底值 {}'.format(
                    name, geometry[key]))
        return geometry, True

    # ── 单次运行 ──
    def run_once(self, mode, speed, duration):
        """跑一次并返回 odom 侧的原始观测。"""
        self.publish_zero(0.5)              # 先确保处于"零指令"状态
        self.yaw_acc = 0.0
        self.yaw_prev = yaw_from_quaternion(self.odom.pose.pose.orientation)
        start = self.odom
        self.vx_sum = 0.0
        self.vy_sum = 0.0
        self.vx_n = 0
        self.wz_sum = 0.0
        self.wz_n = 0

        cmd = Twist()
        if mode == MODE_LINEAR:
            cmd.linear.x = speed
        elif mode == MODE_LATERAL:
            cmd.linear.y = speed
        else:
            cmd.angular.z = speed

        self.recording = True
        t_end = time.monotonic() + duration
        next_pub = time.monotonic()
        while rclpy.ok() and time.monotonic() < t_end:
            now = time.monotonic()
            if now >= next_pub:
                self.pub.publish(cmd)
                next_pub = now + self.period
            rclpy.spin_once(self.node, timeout_sec=0.005)
        self.recording = False

        self.publish_zero(1.2)              # 停稳（含减速段）
        end = self.odom

        sy = yaw_from_quaternion(start.pose.pose.orientation)
        ey = yaw_from_quaternion(end.pose.pose.orientation)
        return {
            'dx': end.pose.pose.position.x - start.pose.pose.position.x,
            'dy': end.pose.pose.position.y - start.pose.pose.position.y,
            'yaw_delta_pose': wrap_angle(ey - sy),   # 只用于量级核对
            'yaw_delta_int': self.yaw_acc,           # 积分值（可跨多圈）
            'mean_vx': (self.vx_sum / self.vx_n) if self.vx_n else float('nan'),
            'mean_vy': (self.vy_sum / self.vx_n) if self.vx_n else float('nan'),
            'mean_wz': (self.wz_sum / self.wz_n) if self.wz_n else float('nan'),
            'odom_samples': self.vx_n,
            'duration': duration,
        }


# ── 结果解释 ────────────────────────────────────────────────────────


def print_observations(res, mode):
    """跑完先报 odom 侧观测值（用户量完尺子后再输入实测值）。"""
    print()
    if mode == MODE_YAW:
        print('  odom 累积转角 : {:.4f} rad (={:.2f} deg, 积分值)'.format(
            res['yaw_delta_int'], math.degrees(res['yaw_delta_int'])))
        print('  odom pose 差角: {:+.4f} rad （仅核对量级，多圈会折算）'.format(
            res['yaw_delta_pose']))
    else:
        print('  odom dx       : {:+.4f} m   (采样 {} 帧)'.format(
            res['dx'], res['odom_samples']))
        print('  odom dy       : {:+.4f} m   （直行偏离，>0.01 说明逐轮/打滑）'
              .format(res['dy']))
        print('  odom dyaw     : {:+.4f} rad （直行转角，>0.017 说明轮子不一致）'
              .format(res['yaw_delta_pose']))
    print('  odom 稳态速度 : vx={:+.4f} vy={:+.4f} wz={:+.4f}'.format(
        res['mean_vx'], res['mean_vy'], res['mean_wz']))


def report_linear(res, measured, geometry, mode):
    if mode == MODE_LINEAR:
        odom_d = res['dx']
        sign_name = 'dx'
    else:
        odom_d = abs(res['dy'])
        sign_name = '|dy|'
    if abs(odom_d) < 1e-3:
        print('  [失败] odom {} 几乎为 0，车没动（检查 speed_enable / 许可）。'
              .format(sign_name))
        return None
    k = measured / odom_d
    r = geometry['wheel_radius']
    g = geometry['gear_ratio']
    print()
    print('  实测位移      : {:.4f} m   (odom {} {:.4f} m)'.format(
        measured, sign_name, odom_d))
    print('  ==> k_lin     = {:.5f}'.format(k))
    print()
    print('  当前 wheel_radius={:.6f}  gear_ratio={:.4f}'.format(r, g))
    print('    每电机转 {:.4f} mm, counts/m {:.0f}'.format(
        2.0 * math.pi * r / g * 1e3,
        g * 10000.0 / (2.0 * math.pi * r)))
    print('  方案A（只改轮径，减速比保持 {:.4f}）'.format(g))
    print('     wheel_radius: {:.6f}  ->  {:.6f}   (对应外径 {:.1f} mm)'
          .format(r, r * k, r * k * 2e3))
    print('  方案B（只改减速比，轮径保持 {:.6f}）'.format(r))
    print('     gear_ratio:   {:.4f}   ->  {:.4f}'.format(g, g * k))
    print('  两者只差一个自由度（代码里只用 r/G 的比值），任选一个；'
          '用卡尺量轮子外径决定改哪个。')
    return k


def report_yaw(res, theta_real, geometry):
    odom_yaw = abs(res['yaw_delta_int'])
    if odom_yaw < 1e-3:
        print('  [失败] odom 累积转角几乎为 0，车没转（检查许可）。')
        return None
    k = theta_real / odom_yaw
    w = geometry['wheel_separation']
    ax = geometry['axle_separation']
    print()
    print('  实测转角      : {:.4f} rad (={:.2f} deg)   (odom {:.4f} rad)'
          .format(theta_real, math.degrees(theta_real), odom_yaw))
    print('  ==> k_ang     = {:.5f}'.format(k))
    print()
    print('  当前 wheel_separation={:.4f}  axle_separation={:.4f}  '
          'half_diagonal={:.4f}'.format(w, ax, (w + ax) / 2.0))
    print('  按比例同时修改（保持第 1 步量出的比例关系）：')
    print('     wheel_separation: {:.4f} -> {:.4f}'.format(w, w * k))
    print('     axle_separation:  {:.4f} -> {:.4f}'.format(ax, ax * k))
    return k


def tracking_warning(res, cmd_speed, mode):
    """指令跟随性检查：odom 速度 vs 指令速度。"""
    if mode == MODE_YAW:
        mean = res['mean_wz']
    elif mode == MODE_LINEAR:
        mean = res['mean_vx']
    else:
        mean = res['mean_vy']
    if math.isnan(mean) or abs(cmd_speed) < 1e-6:
        return
    ratio = mean / cmd_speed
    if abs(ratio - 1.0) > 0.02:
        print()
        print('  [警告] odom 稳态速度 / 指令速度 = {:.4f}（偏离 >2%）。'.format(
            ratio))
        print('         若这个偏差和目标比例接近，说明是**驱动器没跟到指令转速**'
              '（速度单位/整定/限幅），改结构参数只会把它掩盖掉。')
        print('         先把这一项查清再改 wheel_radius / gear_ratio。')


# ── 主流程 ──────────────────────────────────────────────────────────


def plan_mode(args):
    """交互补全 模式/速度/时长/臂长，返回 (mode, speed, duration, arm)。"""
    if args.mode is None:
        print('模式：')
        print('  1) 直线前进 (linear.x)  -> 标 linear 尺度（wheel_radius/gear_ratio）')
        print('  2) 横移     (linear.y)  -> 校核辊子方向 + linear 尺度')
        print('  3) 原地旋转 (angular.z) -> 标 rotation 尺度（轮距/轴距）')
        choice = ask('请选择', 1, cast=int)
        args.mode = {1: MODE_LINEAR, 2: MODE_LATERAL, 3: MODE_YAW}.get(
            choice, MODE_LINEAR)

    arm = args.arm if args.arm is not None else 1.0
    if args.mode == MODE_YAW:
        if args.speed is None:
            args.speed = ask('角速度 [rad/s]', 0.5)
        if args.turns is None:
            args.turns = ask('转几圈（0.5 = 半圈，负数表示反向）', 0.5)
        duration = abs(args.turns) * 2.0 * math.pi / abs(args.speed)
        if args.arm is None:
            arm = ask('地面标定杆的臂长 [m]（杆端划弧，越长约准）', arm)
        speed = abs(args.speed) * (1.0 if args.turns >= 0 else -1.0)
    else:
        if args.speed is None:
            args.speed = ask('速度 [m/s]', 0.2)
        if args.distance is None:
            args.distance = ask('计划行程 [m]（起点标记 -> 终点标记）', 3.0)
        duration = abs(args.distance) / abs(args.speed)
        speed = args.speed
    return args.mode, speed, duration, arm


def ask_angle(prompt, arm):
    """输入杆端弧长 / deg:XX / turn:X，返回实际转角 [rad]。"""
    while True:
        try:
            raw = input(prompt).strip()
        except EOFError:
            raise SystemExit('\n输入结束，退出。')
        if not raw:
            print('  请输入一个值。')
            continue
        try:
            if raw.startswith('deg:'):
                return math.radians(float(raw[4:]))
            if raw.startswith('turn:'):
                return float(raw[5:]) * 2.0 * math.pi
            return float(raw) / arm
        except ValueError:
            print('  无法解析；支持 米数 / deg:180 / turn:0.5。')


def check_ready(cal, args):
    """动之前的安全检查。"""
    ok = True
    if not cal.wait_odom(5.0):
        print('[失败] 5s 内没收到 {} 的 odom。'.format(args.odom_topic))
        return False
    if cal.wait_status(2.0):
        st = cal.status or {}
        if isinstance(st, dict):
            print('  chassis 状态: speed_enable={} allowed={} '
                  'dispatched={} cmd_timeout={}'.format(
                      st.get('speed_enable'), st.get('allowed'),
                      st.get('dispatched'), st.get('cmd_timeout')))
            if not st.get('speed_enable') and not args.force:
                print('[拒绝] speed_enable=false：速度开关没开，车不会动。')
                print('       先执行：ros2 service call {} '
                      'std_srvs/srv/SetBool "{{data: true}}"'
                      .format('/robot_base/speed_enable'))
                print('       （或确认用 --force 跳过本检查）')
                ok = False
    else:
        print('  [提示] 没收到 {}，跳过 speed_enable 检查。'
              .format(args.status_topic))
    return ok


def main():
    parser = argparse.ArgumentParser(
        description='底盘标定：直线 / 横移 / 原地旋转（用地面实测值反算结构参数）')
    parser.add_argument('--mode', choices=[MODE_LINEAR, MODE_LATERAL, MODE_YAW],
                        help='标定模式；不给则交互选择')
    parser.add_argument('--speed', type=float,
                        help='线速度 [m/s] 或角速度 [rad/s]（按模式）')
    parser.add_argument('--distance', type=float, help='直线/横移的计划行程 [m]')
    parser.add_argument('--turns', type=float, help='旋转圈数（0.5 = 半圈）')
    parser.add_argument('--arm', type=float, default=None,
                        help='旋转标定杆臂长 [m]；不给则交互询问（默认 1.0）')
    parser.add_argument('--cmd-topic', default=DEFAULT_CMD_TOPIC)
    parser.add_argument('--odom-topic', default=DEFAULT_ODOM_TOPIC)
    parser.add_argument('--status-topic', default=DEFAULT_STATUS_TOPIC)
    parser.add_argument('--chassis-node', default=DEFAULT_CHASSIS_NODE)
    parser.add_argument('--rate', type=float, default=20.0,
                        help='cmd_vel 发布频率 [Hz]，默认 20')
    parser.add_argument('--yes', action='store_true', help='跳过启动确认')
    parser.add_argument('--force', action='store_true',
                        help='跳过 speed_enable=false 的拦截')
    parser.add_argument('--dry-run', action='store_true',
                        help='只打印计划，不动车')
    args = parser.parse_args()

    rclpy.init()
    node = rclpy.create_node('chassis_calibrator')
    cal = ChassisCalibrator(node, args)
    cal.recording = False

    try:
        print('=== 底盘标定工具 ===')
        print('  cmd_vel : {}'.format(args.cmd_topic))
        print('  odom    : {}'.format(args.odom_topic))

        geometry = dict(FALLBACK_GEOMETRY)
        got_params = False
        if cal.wait_odom(5.0):
            geometry, got_params = cal.fetch_geometry()
            print('  几何参数（{}）: wheel_radius={:.6f} gear_ratio={:.4f} '
                  'wheel_separation={:.4f} axle_separation={:.4f}'.format(
                      '运行时读取' if got_params else '兜底值',
                      geometry['wheel_radius'], geometry['gear_ratio'],
                      geometry['wheel_separation'],
                      geometry['axle_separation']))

        mode, speed, duration, arm = plan_mode(args)
        unit = 'rad/s' if mode == MODE_YAW else 'm/s'
        print()
        print('计划：mode={} speed={:.3f} {} duration={:.2f} s{}'.format(
            mode, speed, unit, duration,
            '' if mode != MODE_YAW else ' arm={:.3f} m'.format(arm)))

        if args.dry_run:
            print('[dry-run] 不动车，退出。')
            return 0

        if not check_ready(cal, args):
            return 1

        if not args.yes:
            ans = ask_text('\n车要动了，确认开始?', 'n')
            if ans != 'y':
                print('已取消。')
                return 0

        ks = []
        while True:
            print('\n--- 开始运行（{:.2f} s）---'.format(duration))
            res = cal.run_once(mode, speed, duration)
            print_observations(res, mode)
            tracking_warning(res, speed, mode)

            if mode == MODE_YAW:
                theta = ask_angle('杆端划过的弧长 [m]'
                                  '（臂长 {:.3f} m；也可输入 deg:180 或 turn:0.5）: '
                                  .format(arm), arm)
                k = report_yaw(res, abs(theta), geometry)
            else:
                measured = ask('卷尺实测位移 [m]', None)
                k = report_linear(res, measured, geometry, mode)

            if k is not None:
                ks.append(k)
                if len(ks) > 1:
                    ordered = sorted(ks)
                    med = ordered[len(ordered) // 2]
                    spread = (max(ks) - min(ks)) / med if med else 0.0
                    print('\n  已测 {} 次: k = {} -> 中位数 {:.5f}，'
                          '极差 {:.2%}'.format(
                              len(ks),
                              ', '.join('{:.5f}'.format(v) for v in ks),
                              med, spread))
                    if spread > 0.02:
                        print('  [提示] 离散度 >2%，建议再测一次；'
                              '打滑/中途掉速会污染结果。')

            if ask_text('\n再测一次?', 'n') != 'y':
                break

        print('\n=== 完成 ===')
        print('改动只写 params/chassis/<机型>.yaml；几何参数在 chassis 节点'
              '启动时读取，改完必须重启节点。')
        print('（ros2 param set 对 odom 不生效）')
        return 0
    finally:
        try:
            cal.recording = False
            cal.publish_zero(0.3)
        except Exception:                       # noqa: BLE001 - 退出路径尽力而为
            pass
        node.destroy_node()
        rclpy.shutdown()


if __name__ == '__main__':
    sys.exit(main())
