import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit
from sklearn.metrics import r2_score
import warnings
import math
warnings.filterwarnings('ignore')

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False

# 度数到弧度的转换常数
DEG_TO_RAD = math.pi / 180.0

# 读取CSV文件
csv_path = 'csv/各关节质量与运动角度相关数据表格.csv'
df = pd.read_csv(csv_path, encoding='utf-8')

# 提取数据（度数）
mcp_angles_deg = df.iloc[:, 0].values
finger_names = ['食指', '中指', '无名指', '尾指']
pip_angles_deg = {}
for i, name in enumerate(finger_names):
    pip_angles_deg[name] = df.iloc[:, i+1].values

print("="*60)
print("主动关节(MCP)与被动关节(PIP)函数关系拟合分析")
print("="*60)
print(f"\n数据点数量: {len(mcp_angles_deg)}")
print(f"MCP角度范围: {mcp_angles_deg.min():.2f}° - {mcp_angles_deg.max():.2f}°\n")

# 检查初始值
print("初始值检查（MCP=0°时的PIP值）:")
for name in finger_names:
    pip_at_zero = pip_angles_deg[name][0]
    print(f"  {name}: PIP = {pip_at_zero:.2f}° (当MCP=0°时)")
print()

# 定义拟合函数（使用度数，最后转换为弧度）
def linear_func(x, a, b):
    """线性函数: y = a*x + b"""
    return a * x + b

def quadratic_func(x, a, b, c):
    """二次函数: y = a*x^2 + b*x + c"""
    return a * x**2 + b * x + c

# 存储拟合结果
fitting_results = {}

# 对每个手指进行拟合
for finger_name in finger_names:
    print(f"\n{'='*60}")
    print(f"分析: {finger_name}")
    print(f"{'='*60}")
    
    y_data_deg = pip_angles_deg[finger_name]
    
    # 尝试不同的拟合方法
    fits = {}
    
    # 1. 线性拟合
    try:
        popt_linear, _ = curve_fit(linear_func, mcp_angles_deg, y_data_deg)
        y_pred_linear = linear_func(mcp_angles_deg, *popt_linear)
        r2_linear = r2_score(y_data_deg, y_pred_linear)
        
        # 计算平均绝对误差
        mae_linear = np.mean(np.abs(y_data_deg - y_pred_linear))
        
        fits['线性'] = {
            'params': popt_linear,
            'r2': r2_linear,
            'mae': mae_linear,
            'func': linear_func,
            'formula_deg': f"PIP(°) = {popt_linear[0]:.6f} * MCP(°) + {popt_linear[1]:.6f}"
        }
        print(f"线性拟合: {fits['线性']['formula_deg']}")
        print(f"  R² = {r2_linear:.6f}, MAE = {mae_linear:.4f}°")
    except:
        pass
    
    # 2. 二次拟合
    try:
        popt_quad, _ = curve_fit(quadratic_func, mcp_angles_deg, y_data_deg)
        y_pred_quad = quadratic_func(mcp_angles_deg, *popt_quad)
        r2_quad = r2_score(y_data_deg, y_pred_quad)
        
        # 计算平均绝对误差
        mae_quad = np.mean(np.abs(y_data_deg - y_pred_quad))
        
        fits['二次'] = {
            'params': popt_quad,
            'r2': r2_quad,
            'mae': mae_quad,
            'func': quadratic_func,
            'formula_deg': f"PIP(°) = {popt_quad[0]:.6f} * MCP(°)² + {popt_quad[1]:.6f} * MCP(°) + {popt_quad[2]:.6f}"
        }
        print(f"二次拟合: {fits['二次']['formula_deg']}")
        print(f"  R² = {r2_quad:.6f}, MAE = {mae_quad:.4f}°")
    except:
        pass
    
    # 选择最佳拟合（优先线性，除非二次明显更好）
    if fits:
        best_fit = None
        best_name = None
        
        if '线性' in fits:
            best_fit = fits['线性']
            best_name = '线性'
        
        if '二次' in fits:
            # 如果二次拟合的R²提升超过0.0001，且MAE更小，使用二次
            if best_fit is None or (fits['二次']['r2'] > best_fit['r2'] + 0.0001 and 
                                   fits['二次']['mae'] < best_fit['mae']):
                best_fit = fits['二次']
                best_name = '二次'
        
        if best_fit:
            fitting_results[finger_name] = {
                'best': best_name,
                'best_fit': best_fit
            }
            print(f"\n最佳拟合: {best_name} (R² = {best_fit['r2']:.6f}, MAE = {best_fit['mae']:.4f}°)")

# 可视化结果
fig, axes = plt.subplots(2, 2, figsize=(15, 12))
axes = axes.flatten()

for idx, finger_name in enumerate(finger_names):
    ax = axes[idx]
    
    # 绘制原始数据点
    ax.scatter(mcp_angles_deg, pip_angles_deg[finger_name], 
              alpha=0.6, s=30, label='原始数据', color='blue')
    
    # 绘制拟合曲线
    if finger_name in fitting_results:
        result = fitting_results[finger_name]
        best_fit = result['best_fit']
        
        # 生成平滑曲线
        x_smooth = np.linspace(mcp_angles_deg.min(), mcp_angles_deg.max(), 200)
        y_smooth = best_fit['func'](x_smooth, *best_fit['params'])
        ax.plot(x_smooth, y_smooth, 'r-', linewidth=2, 
               label=f"{result['best']}拟合 (R²={best_fit['r2']:.4f})")
    
    ax.set_xlabel('MCP角度 (°)', fontsize=12)
    ax.set_ylabel('PIP角度 (°)', fontsize=12)
    ax.set_title(f'{finger_name} - MCP vs PIP', fontsize=14, fontweight='bold')
    ax.legend(fontsize=9)
    ax.grid(True, alpha=0.3)

plt.tight_layout()
plt.savefig('scripts/fitting_results.png', dpi=300, bbox_inches='tight')
print(f"\n\n可视化结果已保存到: scripts/fitting_results.png")

# 生成用于MuJoCo XML的polycoef格式
print("\n" + "="*60)
print("MuJoCo XML格式的polycoef系数")
print("="*60)
print("重要说明:")
print("1. MuJoCo公式: PIP - PIP₀ = a₀ + a₁(MCP - MCP₀) + a₂(MCP - MCP₀)² + ...")
print("2. 如果初始配置中 MCP₀=0, PIP₀=0，则: PIP = a₀ + a₁·MCP + a₂·MCP² + ...")
print("3. 数据中MCP=0°时PIP≠0°，这意味着a₀应该包含初始偏移")
print("4. 系数需要从度数转换为弧度\n")

for finger_name in finger_names:
    if finger_name in fitting_results:
        result = fitting_results[finger_name]
        best_fit = result['best_fit']
        params_deg = best_fit['params']  # 度数的参数
        
        # 将度数参数转换为弧度参数
        if result['best'] == '线性':
            # 线性: PIP(°) = a*MCP(°) + b
            # 转换为弧度: PIP(rad) = a*MCP(rad) + b*π/180
            # 因为: PIP(rad) = PIP(°)*π/180 = (a*MCP(°)+b)*π/180 = a*MCP(rad) + b*π/180
            a_rad = params_deg[0]  # 线性系数不变（因为两边都乘以π/180）
            b_rad = params_deg[1] * DEG_TO_RAD  # 常数项需要转换
            polycoef = [b_rad, a_rad, 0.0, 0.0, 0.0]
            
        elif result['best'] == '二次':
            # 二次: PIP(°) = a*MCP(°)² + b*MCP(°) + c
            # 转换为弧度: PIP(rad) = a*(MCP(rad)*180/π)² + b*(MCP(rad)*180/π) + c
            #            = a*(180/π)²*MCP(rad)² + b*(180/π)*MCP(rad) + c*π/180
            a_deg = params_deg[0]
            b_deg = params_deg[1]
            c_deg = params_deg[2]
            a_rad = a_deg * (DEG_TO_RAD ** 2)
            b_rad = b_deg * DEG_TO_RAD
            c_rad = c_deg * DEG_TO_RAD
            polycoef = [c_rad, b_rad, a_rad, 0.0, 0.0]
        else:
            polycoef = [0.0, 1.0, 0.0, 0.0, 0.0]
        
        # 格式化输出
        polycoef_str = ' '.join([f"{coef:.10f}" for coef in polycoef])
        print(f"\n{finger_name}:")
        print(f"  拟合方法: {result['best']}")
        print(f"  度数公式: {best_fit['formula_deg']}")
        print(f"  polycoef=\"{polycoef_str}\"")
        print(f"  弧度公式: PIP(rad) = {polycoef[0]:.6f} + {polycoef[1]:.6f}*MCP(rad) + {polycoef[2]:.6f}*MCP(rad)²")

# 输出XML格式的代码片段
print("\n" + "="*60)
print("XML代码片段（可直接复制）")
print("="*60)
print("<equality>")
print("  <!--  mimic distal -> proximal  -->")
print("  <!-- 基于实际数据拟合的函数关系（弧度）: PIP = a₀ + a₁*MCP + a₂*MCP² + ... -->")
print("  <!-- 注意: a₀包含了MCP=0时PIP的初始偏移 -->")

for finger_name in finger_names:
    if finger_name in fitting_results:
        result = fitting_results[finger_name]
        best_fit = result['best_fit']
        params_deg = best_fit['params']
        
        # 转换为弧度polycoef
        if result['best'] == '线性':
            a_rad = params_deg[0]
            b_rad = params_deg[1] * DEG_TO_RAD
            polycoef = [b_rad, a_rad, 0.0, 0.0, 0.0]
        elif result['best'] == '二次':
            a_rad = params_deg[0] * (DEG_TO_RAD ** 2)
            b_rad = params_deg[1] * DEG_TO_RAD
            c_rad = params_deg[2] * DEG_TO_RAD
            polycoef = [c_rad, b_rad, a_rad, 0.0, 0.0]
        else:
            polycoef = [0.0, 1.0, 0.0, 0.0, 0.0]
        
        polycoef_str = ' '.join([f"{coef:.10f}" for coef in polycoef])
        
        # 映射手指名称到关节名称
        joint_map = {
            '食指': 'index',
            '中指': 'middle',
            '无名指': 'ring',
            '尾指': 'little'
        }
        joint_name = joint_map.get(finger_name, finger_name)
        print(f"  <joint joint1=\"{joint_name}MCP\" joint2=\"{joint_name}PIP\" polycoef=\"{polycoef_str}\"/>")
print("</equality>")

print("\n分析完成！")