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
warnings.filterwarnings('ignore')

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

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

# 提取数据
mcp_angles = df.iloc[:, 0].values  # MCP角度（主动关节）
finger_names = ['食指', '中指', '无名指', '尾指']
pip_angles = {}
for i, name in enumerate(finger_names):
    pip_angles[name] = df.iloc[:, i+1].values  # PIP角度（被动关节）

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

# 定义拟合函数
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

def cubic_func(x, a, b, c, d):
    """三次函数: y = a*x^3 + b*x^2 + c*x + d"""
    return a * x**3 + b * x**2 + c * x + d

# 存储拟合结果
fitting_results = {}

# 对每个手指进行拟合
for finger_name in finger_names:
    print(f"\n{'='*60}")
    print(f"分析: {finger_name}")
    print(f"{'='*60}")
    
    y_data = pip_angles[finger_name]
    
    # 尝试不同的拟合方法
    fits = {}
    
    # 1. 线性拟合
    try:
        popt_linear, _ = curve_fit(linear_func, mcp_angles, y_data)
        y_pred_linear = linear_func(mcp_angles, *popt_linear)
        r2_linear = r2_score(y_data, y_pred_linear)
        fits['线性'] = {
            'params': popt_linear,
            'r2': r2_linear,
            'func': linear_func,
            'formula': f"PIP = {popt_linear[0]:.6f} * MCP + {popt_linear[1]:.6f}"
        }
        print(f"线性拟合: {fits['线性']['formula']}")
        print(f"  R² = {r2_linear:.6f}")
    except:
        pass
    
    # 2. 二次拟合
    try:
        popt_quad, _ = curve_fit(quadratic_func, mcp_angles, y_data)
        y_pred_quad = quadratic_func(mcp_angles, *popt_quad)
        r2_quad = r2_score(y_data, y_pred_quad)
        fits['二次'] = {
            'params': popt_quad,
            'r2': r2_quad,
            'func': quadratic_func,
            'formula': f"PIP = {popt_quad[0]:.6f} * MCP² + {popt_quad[1]:.6f} * MCP + {popt_quad[2]:.6f}"
        }
        print(f"二次拟合: {fits['二次']['formula']}")
        print(f"  R² = {r2_quad:.6f}")
    except:
        pass
    
    # 3. 三次拟合
    try:
        popt_cubic, _ = curve_fit(cubic_func, mcp_angles, y_data)
        y_pred_cubic = cubic_func(mcp_angles, *popt_cubic)
        r2_cubic = r2_score(y_data, y_pred_cubic)
        fits['三次'] = {
            'params': popt_cubic,
            'r2': r2_cubic,
            'func': cubic_func,
            'formula': f"PIP = {popt_cubic[0]:.6f} * MCP³ + {popt_cubic[1]:.6f} * MCP² + {popt_cubic[2]:.6f} * MCP + {popt_cubic[3]:.6f}"
        }
        print(f"三次拟合: {fits['三次']['formula']}")
        print(f"  R² = {r2_cubic:.6f}")
    except:
        pass
    
    # 选择最佳拟合（R²最高）
    if fits:
        best_fit = max(fits.items(), key=lambda x: x[1]['r2'])
        best_name, best_data = best_fit
        fitting_results[finger_name] = {
            'best': best_name,
            'all_fits': fits,
            'best_fit': best_data
        }
        print(f"\n最佳拟合: {best_name} (R² = {best_data['r2']:.6f})")
        print(f"函数关系: {best_data['formula']}")

# 可视化结果
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, pip_angles[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.min(), mcp_angles.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})")
        
        # 显示所有拟合结果
        for fit_name, fit_data in result['all_fits'].items():
            if fit_name != result['best']:
                y_fit = fit_data['func'](mcp_angles, *fit_data['params'])
                ax.plot(mcp_angles, y_fit, '--', alpha=0.5, 
                       label=f"{fit_name} (R²={fit_data['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")

# 输出总结
print("\n" + "="*60)
print("拟合结果总结")
print("="*60)

for finger_name in finger_names:
    if finger_name in fitting_results:
        result = fitting_results[finger_name]
        print(f"\n{finger_name}:")
        print(f"  最佳拟合方法: {result['best']}")
        print(f"  函数关系: {result['best_fit']['formula']}")
        print(f"  拟合优度 (R²): {result['best_fit']['r2']:.6f}")

# 生成用于MuJoCo XML的polycoef格式
print("\n" + "="*60)
print("MuJoCo XML格式的polycoef系数")
print("="*60)
print("格式: <joint joint1=\"MCP\" joint2=\"PIP\" polycoef=\"...\"/>")
print("polycoef格式: [c0, c1, c2, c3, ...] 表示 PIP = c0 + c1*MCP + c2*MCP² + c3*MCP³ + ...\n")

for finger_name in finger_names:
    if finger_name in fitting_results:
        result = fitting_results[finger_name]
        best_fit = result['best_fit']
        params = best_fit['params']
        
        # 根据拟合类型生成polycoef
        if result['best'] == '线性':
            # 线性: y = a*x + b -> polycoef = [b, a, 0, 0]
            polycoef = [params[1], params[0], 0, 0]
        elif result['best'] == '二次':
            # 二次: y = a*x² + b*x + c -> polycoef = [c, b, a, 0]
            polycoef = [params[2], params[1], params[0], 0]
        elif result['best'] == '三次':
            # 三次: y = a*x³ + b*x² + c*x + d -> polycoef = [d, c, b, a]
            polycoef = [params[3], params[2], params[1], params[0]]
        else:
            polycoef = [0, 1, 0, 0]  # 默认
        
        # 格式化输出
        polycoef_str = ' '.join([f"{coef:.6f}" for coef in polycoef])
        joint_name = finger_name.replace('指', '')
        print(f"{finger_name}: polycoef=\"{polycoef_str}\"")

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