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

# 转换为弧度
mcp_angles_rad = mcp_angles_deg * DEG_TO_RAD
pip_angles_rad = {}
for name in finger_names:
    pip_angles_rad[name] = pip_angles_deg[name] * DEG_TO_RAD

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")

# 获取MCP=0时的实际PIP值（弧度）
pip_at_zero_rad = {}
for name in finger_names:
    pip_at_zero_rad[name] = pip_angles_rad[name][0]

print("MCP=0时的实际PIP值（弧度）:")
for name in finger_names:
    print(f"  {name}: PIP = {pip_at_zero_rad[name]:.6f} rad = {pip_at_zero_rad[name]/DEG_TO_RAD:.2f}°")
print()

# 定义带固定截距的二次拟合函数
def quadratic_func_fixed_intercept(x, a, b, c_fixed):
    """二次函数: y = a*x^2 + b*x + c_fixed (截距固定)"""
    return a * x**2 + b * x + c_fixed

# 存储拟合结果
fitting_results = {}

# 对每个手指进行二次拟合（固定截距为MCP=0时的实际值）
for finger_name in finger_names:
    print(f"\n{'='*60}")
    print(f"分析: {finger_name}")
    print(f"{'='*60}")
    
    y_data_rad = pip_angles_rad[finger_name]
    x_data_rad = mcp_angles_rad
    c_fixed = pip_at_zero_rad[finger_name]  # 固定截距为MCP=0时的实际值
    
    # 二次拟合: PIP = a*MCP² + b*MCP + c_fixed
    # 转换为: PIP - c_fixed = a*MCP² + b*MCP
    y_adjusted = y_data_rad - c_fixed
    
    # 构建矩阵进行最小二乘拟合
    # y_adjusted = a*x² + b*x
    # 可以写成: [x², x] * [a, b]^T = y_adjusted
    A = np.vstack([x_data_rad**2, x_data_rad]).T
    coeffs, residuals, rank, s = np.linalg.lstsq(A, y_adjusted, rcond=None)
    a = coeffs[0]
    b = coeffs[1]
    
    # 计算拟合结果
    y_pred_rad = a * x_data_rad**2 + b * x_data_rad + c_fixed
    r2 = r2_score(y_data_rad, y_pred_rad)
    mae = np.mean(np.abs(y_data_rad - y_pred_rad))
    
    fitting_results[finger_name] = {
        'a': a,  # 二次项系数
        'b': b,  # 一次项系数
        'c': c_fixed,  # 常数项（固定）
        'r2': r2,
        'mae': mae,
        'formula_rad': f"PIP(rad) = {a:.6f} * MCP(rad)² + {b:.6f} * MCP(rad) + {c_fixed:.6f}"
    }
    
    print(f"二次拟合（固定截距）: {fitting_results[finger_name]['formula_rad']}")
    print(f"  R² = {r2:.6f}, MAE = {mae:.6f} rad ({mae/DEG_TO_RAD:.4f}°)")
    print(f"  验证: 当MCP=0时, PIP = {c_fixed:.6f} rad = {c_fixed/DEG_TO_RAD:.2f}° (实际值)")

# 详细验证：计算多个关键点的拟合值并与实际值对比
print("\n" + "="*60)
print("详细验证：计算关键点的拟合值")
print("="*60)

# 选择几个关键点进行验证：0°, 10°, 30°, 50°, 70°
test_points_deg = [0, 10, 30, 50, 70]
test_points_rad = [deg * DEG_TO_RAD for deg in test_points_deg]

for finger_name in finger_names:
    if finger_name in fitting_results:
        result = fitting_results[finger_name]
        a = result['a']
        b = result['b']
        c = result['c']
        
        print(f"\n{finger_name}:")
        print(f"  公式: PIP(rad) = {a:.6f} * MCP(rad)² + {b:.6f} * MCP(rad) + {c:.6f}")
        print(f"\n  验证点对比:")
        print(f"  {'MCP(°)':<10} {'实际PIP(°)':<15} {'拟合PIP(°)':<15} {'误差(°)':<12} {'误差(%)':<10}")
        print(f"  {'-'*70}")
        
        max_error = 0
        for mcp_deg in test_points_deg:
            # 找到对应的实际值
            if mcp_deg in mcp_angles_deg:
                idx = np.where(mcp_angles_deg == mcp_deg)[0][0]
                actual_pip_deg = pip_angles_deg[finger_name][idx]
            else:
                # 如果不在数据点中，使用插值
                actual_pip_deg = np.interp(mcp_deg, mcp_angles_deg, pip_angles_deg[finger_name])
            
            # 计算拟合值
            mcp_rad = mcp_deg * DEG_TO_RAD
            pip_fit_rad = a * mcp_rad**2 + b * mcp_rad + c
            pip_fit_deg = pip_fit_rad / DEG_TO_RAD
            
            # 计算误差
            error_deg = abs(pip_fit_deg - actual_pip_deg)
            error_percent = (error_deg / actual_pip_deg * 100) if actual_pip_deg > 0 else 0
            max_error = max(max_error, error_deg)
            
            print(f"  {mcp_deg:<10.0f} {actual_pip_deg:<15.2f} {pip_fit_deg:<15.2f} {error_deg:<12.4f} {error_percent:<10.2f}%")
        
        print(f"  最大误差: {max_error:.4f}° ({max_error*DEG_TO_RAD:.6f} rad)")

# 可视化结果
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]
        a = result['a']
        b = result['b']
        c = result['c']
        
        # 生成平滑曲线
        x_smooth_rad = np.linspace(mcp_angles_rad.min(), mcp_angles_rad.max(), 200)
        y_smooth_rad = a * x_smooth_rad**2 + b * x_smooth_rad + c
        x_smooth_deg = x_smooth_rad / DEG_TO_RAD
        y_smooth_deg = y_smooth_rad / DEG_TO_RAD
        ax.plot(x_smooth_deg, y_smooth_deg, 'r-', linewidth=2, 
               label=f"二次拟合 (R²={result['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格式（4个系数）
print("\n" + "="*60)
print("MuJoCo XML格式的polycoef系数（4个值）")
print("="*60)
print("MuJoCo公式: PIP = a₀ + a₁*MCP + a₂*MCP² + a₃*MCP³")
print("polycoef格式: [a₀, a₁, a₂, a₃] (4个系数，弧度)")
print("二次拟合: PIP = c + b*MCP + a*MCP²")
print("因此: polycoef = [c, b, a, 0]\n")

# 输出XML格式的代码片段
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]
        a = result['a']  # 二次项系数
        b = result['b']  # 一次项系数
        c = result['c']  # 常数项（固定）
        
        # 二次: PIP(rad) = a*MCP(rad)² + b*MCP(rad) + c
        # polycoef = [c, b, a, 0] (4个系数)
        polycoef = [c, b, a, 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(f"    <!-- {finger_name}: PIP = {c:.6f} + {b:.6f}*MCP + {a:.6f}*MCP², R²={result['r2']:.4f}, MAE={result['mae']/DEG_TO_RAD:.2f}° -->")

print("</equality>")

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