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

# 定义正确的拟合函数
def polynomial_func(x, a0, a1, a2, a3):
    """多项式函数: y = a0 + a1*x + a2*x² + a3*x³"""
    return a0 + a1*x + a2*x**2 + a3*x**3

def linear_func(x, a, b):
    """线性函数: y = a*x + b"""
    return a * x + b

def quadratic_func(x, a, b, c):
    """二次函数: y = a*x² + 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_rad = pip_angles_rad[finger_name]
    x_data_rad = mcp_angles_rad
    
    # 尝试不同的拟合方法
    models = {
        '线性': linear_func,
        '二次': quadratic_func,
        '三次': polynomial_func
    }
    
    best_r2 = -np.inf
    best_result = None
    best_model_name = None
    
    for model_name, model_func in models.items():
        try:
            if model_name == '线性':
                # 线性拟合
                popt, pcov = curve_fit(model_func, x_data_rad, y_data_rad)
                y_pred_rad = model_func(x_data_rad, *popt)
                params = popt
            elif model_name == '二次':
                # 二次拟合
                popt, pcov = curve_fit(model_func, x_data_rad, y_data_rad)
                y_pred_rad = model_func(x_data_rad, *popt)
                params = [popt[1], popt[0], popt[2]]  # 重新排列为 [a1, a2, a0]
            else:  # 三次
                # 三次多项式拟合
                popt, pcov = curve_fit(model_func, x_data_rad, y_data_rad)
                y_pred_rad = model_func(x_data_rad, *popt)
                params = popt
            
            r2 = r2_score(y_data_rad, y_pred_rad)
            mae = np.mean(np.abs(y_data_rad - y_pred_rad))
            
            if r2 > best_r2:
                best_r2 = r2
                best_result = {
                    'model': model_name,
                    'params': params,
                    'r2': r2,
                    'mae': mae,
                    'y_pred': y_pred_rad
                }
                best_model_name = model_name
                
            print(f"  {model_name}拟合: R² = {r2:.6f}, MAE = {mae:.6f} rad")
            
        except Exception as e:
            print(f"  {model_name}拟合失败: {e}")
    
    # 保存最佳拟合结果
    if best_result:
        fitting_results[finger_name] = best_result
        print(f"\n最佳模型: {best_model_name}")
        
        if best_model_name == '线性':
            a, b = best_result['params']
            print(f"  拟合公式: PIP = {a:.6f} * MCP + {b:.6f}")
        elif best_model_name == '二次':
            a, b, c = best_result['params']
            print(f"  拟合公式: PIP = {a:.6f} * MCP + {b:.6f} * MCP² + {c:.6f}")
        else:  # 三次
            a0, a1, a2, a3 = best_result['params']
            print(f"  拟合公式: PIP = {a0:.6f} + {a1:.6f} * MCP + {a2:.6f} * MCP² + {a3:.6f} * MCP³")
        
        print(f"  R² = {best_result['r2']:.6f}, MAE = {best_result['mae']:.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]
        
        # 生成平滑曲线
        x_smooth_rad = np.linspace(mcp_angles_rad.min(), mcp_angles_rad.max(), 200)
        
        if result['model'] == '线性':
            a, b = result['params']
            y_smooth_rad = linear_func(x_smooth_rad, a, b)
        elif result['model'] == '二次':
            a, b, c = result['params']
            y_smooth_rad = quadratic_func(x_smooth_rad, a, b, c)
        else:  # 三次
            a0, a1, a2, a3 = result['params']
            y_smooth_rad = polynomial_func(x_smooth_rad, a0, a1, a2, a3)
        
        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"{result['model']}拟合 (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_corrected.png', dpi=300, bbox_inches='tight')
print(f"\n\n可视化结果已保存到: scripts/fitting_results_corrected.png")

# 生成用于MuJoCo XML的polycoef格式（4个系数）
print("\n" + "="*60)
print("MuJoCo XML格式的polycoef系数")
print("="*60)
print("MuJoCo公式: PIP = a₀ + a₁*MCP + a₂*MCP² + a₃*MCP³")
print("polycoef格式: [a₀, a₁, a₂, a₃] (4个系数，弧度)\n")

# 输出XML格式的代码片段
print("<equality>")
print("  <!--  mimic distal -> proximal  -->")
print("  <!-- 基于实际数据拟合的多项式关系（弧度） -->")

for finger_name in finger_names:
    if finger_name in fitting_results:
        result = fitting_results[finger_name]
        
        # 映射手指名称到关节名称
        joint_map = {
            '食指': 'index',
            '中指': 'middle',
            '无名指': 'ring',
            '尾指': 'little'
        }
        joint_name = joint_map.get(finger_name, finger_name)
        
        if result['model'] == '线性':
            a, b = result['params']
            # 线性: PIP(rad) = a*MCP(rad) + b
            # 转换为三次多项式格式: a₀ = b, a₁ = a, a₂ = 0, a₃ = 0
            polycoef = [b, a, 0.0, 0.0]
        elif result['model'] == '二次':
            a, b, c = result['params']
            # 二次: PIP(rad) = a*MCP + b*MCP² + c
            # 转换为三次多项式格式: a₀ = c, a₁ = a, a₂ = b, a₃ = 0
            polycoef = [c, a, b, 0.0]
        else:  # 三次
            a0, a1, a2, a3 = result['params']
            # 三次: PIP(rad) = a₀ + a₁*MCP + a₂*MCP² + a₃*MCP³
            polycoef = [a0, a1, a2, a3]
        
        polycoef_str = ' '.join([f"{coef:.10f}" for coef in polycoef])
        
        print(f"  <joint joint1=\"{joint_name}MCP\" joint2=\"{joint_name}PIP\" polycoef=\"{polycoef_str}\"/>")
        print(f"    <!-- {finger_name}: {result['model']}拟合, R²={result['r2']:.4f} -->")

print("</equality>")

# 验证拟合质量
print("\n" + "="*60)
print("拟合质量验证")
print("="*60)
for finger_name in finger_names:
    if finger_name in fitting_results:
        result = fitting_results[finger_name]
        actual_pip_at_zero_deg = pip_angles_deg[finger_name][0]
        
        # 计算MCP=0时的拟合值
        if result['model'] == '线性':
            a, b = result['params']
            fitted_pip_at_zero_rad = b
        elif result['model'] == '二次':
            a, b, c = result['params']
            fitted_pip_at_zero_rad = c
        else:  # 三次
            a0, a1, a2, a3 = result['params']
            fitted_pip_at_zero_rad = a0
        
        fitted_pip_at_zero_deg = fitted_pip_at_zero_rad / DEG_TO_RAD
        error = abs(fitted_pip_at_zero_deg - actual_pip_at_zero_deg)
        
        print(f"{finger_name}:")
        print(f"  实际MCP=0时PIP值: {actual_pip_at_zero_deg:.2f}°")
        print(f"  拟合MCP=0时PIP值: {fitted_pip_at_zero_deg:.2f}°")
        print(f"  误差: {error:.4f}°")
        print(f"  R²: {result['r2']:.6f}")
        print()

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