#!/usr/bin/env python3
"""实时监听 LowState (只读, 不控制电机).

用法:
    python examples/monitor_state.py --duration 30 --verbose
"""
import argparse
import math
import signal
import sys
import time

from go1_pro_sdk import MCUClient, JOINT_NAMES, decode_sn

running = True


def sigint(s, f):
    global running
    running = False


def main():
    p = argparse.ArgumentParser()
    p.add_argument('--state', default=None, help='Blowfish state 文件 (默认用包内置)')
    p.add_argument('--duration', type=float, default=30)
    p.add_argument('--rate', type=float, default=5, help='打印频率 Hz')
    p.add_argument('--verbose', '-v', action='store_true')
    p.add_argument('--show-feet', action='store_true',
                   help='显示 4 脚足端力 (raw ADC) + 估计值')
    p.add_argument('--feet-only', action='store_true',
                   help='只显示足端力, 不显示关节角 (聚焦脚力)')
    p.add_argument('--imu-only', action='store_true',
                   help='只显示 IMU + 姿态 (四元数/角速度/加速度/RPY)')
    args = p.parse_args()
    signal.signal(signal.SIGINT, sigint)

    print('📡 监听 LowState')
    print(f'  时长 {args.duration}s, 打印 {args.rate}Hz')

    with MCUClient(state_path=args.state) as client:
        recv_count = client.wake_mcu(50)
        print(f'唤醒: 收到 {recv_count}/50 帧')
        if recv_count == 0:
            print('❌ 无回包. 检查 Pi 抢占源或 state 是否正确')
            return 1

        start = time.time()
        last_print = 0
        n_decoded = 0
        damping = __import__('go1_pro_sdk').LowCmd().all_damping()
        # 足端力统计
        ff_hist = {'FR': [], 'FL': [], 'RR': [], 'RL': []}
        ffest_hist = {'FR': [], 'FL': [], 'RR': [], 'RL': []}

        if args.imu_only:
            print(f"\n{'时间':>6} {'四元数(w,x,y,z)':>44} {'gyro(°/s)':>28} {'acc(m/s²)':>28} {'RPY(°)':>28}")
            print('-' * 140)
        elif args.feet_only:
            print(f"\n{'时间':>6} {'FR 右前':>10} {'FL 左前':>10} {'RR 右后':>10} {'RL 左后':>10}")
            print('-' * 65)
        else:
            print(f"\n{'时间':>6} {'电量':>5} {'FR_0':>8} {'FR_1':>8} {'FR_2':>8} {'rpy[°]':>20}")
            print('-' * 70)

        while running and time.time() - start < args.duration:
            client.send(damping)
            time.sleep(0.005)
            state = client.recv_latest()
            if state is None:
                continue
            n_decoded += 1
            ff = state.footForce
            ffest = state.footForceEst
            for i, leg in enumerate(['FR', 'FL', 'RR', 'RL']):
                ff_hist[leg].append(ff[i])
                ffest_hist[leg].append(ffest[i])

            now = time.time()
            if now - last_print >= 1.0 / args.rate:
                t = now - start
                imu = state.imu

                if args.imu_only:
                    qw,qx,qy,qz = imu.quaternion
                    gx,gy,gz = imu.gyroscope
                    ax,ay,az = imu.accelerometer
                    r,p,y = imu.rpy
                    print(f'{t:6.1f} '
                          f'({qw:+7.4f},{qx:+7.4f},{qy:+7.4f},{qz:+7.4f}) '
                          f'({math.degrees(gx):+7.2f},{math.degrees(gy):+7.2f},{math.degrees(gz):+7.2f}) '
                          f'({ax:+7.2f},{ay:+7.2f},{az:+7.2f}) '
                          f'({math.degrees(r):+7.2f},{math.degrees(p):+7.2f},{math.degrees(y):+7.2f})')
                elif args.feet_only:
                    print(f'{t:6.1f}       {ff[0]:>5d}     {ff[1]:>5d}     '
                          f'{ff[2]:>5d}     {ff[3]:>5d}')
                else:
                    rpy_str = (f"{math.degrees(state.imu.rpy[0]):+5.1f},"
                               f"{math.degrees(state.imu.rpy[1]):+5.1f},"
                               f"{math.degrees(state.imu.rpy[2]):+5.1f}")
                    print(f'{t:6.1f} {state.bms.SOC:4d}% '
                          f'{state.motorState[0].q:+8.3f} '
                          f'{state.motorState[1].q:+8.3f} '
                          f'{state.motorState[2].q:+8.3f} '
                          f'{rpy_str:>20}')

                if args.show_feet or args.feet_only:
                    print(f"     🦶 足力: [FR={ff[0]:5d} FL={ff[1]:5d} "
                          f"RR={ff[2]:5d} RL={ff[3]:5d}]  "
                          f"估: [FR={ffest[0]:5d} FL={ffest[1]:5d} "
                          f"RR={ffest[2]:5d} RL={ffest[3]:5d}]")

                if args.verbose:
                    print(f'    SN={decode_sn(state.SN)} '
                          f'acc_z={state.imu.accelerometer[2]:.2f} m/s²  '
                          f'BMS current={state.bms.current_a:.2f}A')
                last_print = now

        print(f'\n解码 {n_decoded} 帧, 退出.')
        if ff_hist['FR']:
            print(f'\n🦶 足端力统计 (全程 {n_decoded} 帧):')
            for leg in ['FR', 'FL', 'RR', 'RL']:
                vals = ff_hist[leg]
                vest = ffest_hist[leg]
                print(f'  {leg}: 力 min={min(vals):5d} max={max(vals):5d} range={max(vals)-min(vals)}  '
                      f'估 min={min(vest):5d} max={max(vest):5d} range={max(vest)-min(vest)}')


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