#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
快速提取PWM-速度关系
"""
import sys
import numpy as np

def main():
    if len(sys.argv) < 2:
        print("用法: python extract_pwm_speed.py calibration.log")
        return
    
    log_file = sys.argv[1]
    
    # 存储各档位数据
    pwm_data = {}
    
    with open(log_file, 'r', encoding='utf-8', errors='ignore') as f:
        lines = f.readlines()
    
    # 解析数据
    for i, line in enumerate(lines):
        if not line.startswith('$CAL'):
            continue
        
        try:
            parts = line.strip().split(',')
            if len(parts) < 18:
                continue
            
            state = parts[3]
            throttle_pwm = int(parts[4])
            
            # 只处理巡航状态（速度稳定）
            if '_C' not in state or 'FWD' not in state:
                continue
            
            # 计算速度（从ENU坐标变化）
            if i + 10 < len(lines):  # 确保有后续数据
                enu_x1 = float(parts[6])
                enu_y1 = float(parts[7])
                time1 = int(parts[2])
                
                # 找10个样本后的数据（约0.5秒后）
                for j in range(i+5, min(i+15, len(lines))):
                    if lines[j].startswith('$CAL'):
                        parts2 = lines[j].strip().split(',')
                        if len(parts2) >= 18 and parts2[3] == state:
                            enu_x2 = float(parts2[6])
                            enu_y2 = float(parts2[7])
                            time2 = int(parts2[2])
                            
                            dt = (time2 - time1) / 1000.0  # 毫秒→秒
                            if dt > 0:
                                dx = enu_x2 - enu_x1
                                dy = enu_y2 - enu_y1
                                speed = np.sqrt(dx*dx + dy*dy) / dt
                                
                                if throttle_pwm not in pwm_data:
                                    pwm_data[throttle_pwm] = []
                                pwm_data[throttle_pwm].append(speed)
                            break
        except:
            continue
    
    # 统计结果
    print("\n========== PWM-速度关系 ==========")
    pwm_speed_pairs = []
    for pwm in sorted(pwm_data.keys()):
        speeds = pwm_data[pwm]
        avg_speed = np.mean(speeds)
        std_speed = np.std(speeds)
        pwm_speed_pairs.append((pwm, avg_speed))
        print(f"PWM {pwm}: 平均速度 {avg_speed:.3f} m/s (±{std_speed:.3f}, {len(speeds)}个样本)")
    
    # 线性拟合：v = k × (1500 - pwm) + bias
    if len(pwm_speed_pairs) >= 2:
        pwms = np.array([p[0] for p in pwm_speed_pairs])
        speeds = np.array([p[1] for p in pwm_speed_pairs])
        
        # 转换为 (1500-pwm) vs speed
        x = 1500 - pwms
        y = speeds
        
        # 线性拟合
        coef = np.polyfit(x, y, 1)
        k = coef[0]
        bias = coef[1]
        
        print(f"\n========== 拟合模型 ==========")
        print(f"v = {k:.6f} × (1500 - pwm) + {bias:.6f}")
        print(f"\n反解PWM公式：")
        print(f"pwm = 1500 - (v - {bias:.6f}) / {k:.6f}")
        
        print(f"\n========== 推荐参数 ==========")
        print(f"#define MC_CFG_FORWARD_K        ({k:.6f}f)")
        print(f"#define MC_CFG_FORWARD_BIAS     ({bias:.6f}f)")
        
        # 预测极限速度
        max_speed_1000 = k * (1500 - 1000) + bias
        print(f"\n预测 1000 PWM 最大速度: {max_speed_1000:.3f} m/s")
        print(f"#define MC_CFG_MAX_FORWARD_MPS  ({max_speed_1000:.2f}f)")

if __name__ == '__main__':
    main()