import numpy as np
import matplotlib.pyplot as plt
from matplotlib.widgets import Slider

# --- 全局常数定义 (单位: mm 或 N) ---
L_AB = 215.0
L_AC = 42.0
L_BD = 50.0
L_BE = 254.0  # 50 + 204
Fs = 280.0    # 弹簧力 220N

def calculate_force_ae(L_mm):
    # 1. 几何计算
    cos_theta = (L_AB**2 + L_BE**2 - L_mm**2) / (2 * L_AB * L_BE)
    
    if abs(cos_theta) > 1:
        return np.nan
        
    theta = np.arccos(cos_theta) # 修正：定义 theta
    sin_theta = np.sin(theta)

    # 2. 计算弹簧长度 L_CD (米)
    dx = L_AB - L_BD * cos_theta
    dy = L_AC - L_BD * sin_theta
    L_CD_m = np.sqrt(dx**2 + dy**2) / 1000.0

    # 3. 计算 B 点扭矩 Mb (N·m)
    # 基于之前推导的力臂公式: |2.1*cos(theta) - 10.75*sin(theta)|
    # 这里的系数 2.1 是 L_BD*L_AC/1000, 10.75 是 L_AB*L_BD/1000
    arm = abs(L_BD * cos_theta * L_AC - L_AB * L_BD * sin_theta) / 1000000.0
    Mb = (Fs / L_CD_m) * arm

    # 4. 计算 AE 方向的力分量 F_AE (N)
    L_BE_m = L_BE / 1000.0
    sin_psi = (L_AB * sin_theta) / L_mm if L_mm != 0 else np.nan
    
    F_AE = (Mb / L_BE_m) * sin_psi
    return F_AE

# --- 可视化部分 ---
def update_plot(val):
    L_ae = slider.val
    ax_mech.clear()
    
    A = np.array([0, 0])
    B = np.array([L_AB, 0])
    C = np.array([0, L_AC])
    
    cos_theta = (L_AB**2 + L_BE**2 - L_ae**2) / (2 * L_AB * L_BE)
    
    if abs(cos_theta) <= 1:
        theta = np.arccos(cos_theta)
        # E 和 D 坐标计算 (对应手绘图趋势，y轴向上)
        E = np.array([L_AB - L_BE * np.cos(theta), L_BE * np.sin(theta)])
        D = np.array([L_AB - L_BD * np.cos(theta), L_BD * np.sin(theta)])
        
        ax_mech.plot([A[0], B[0]], [A[1], B[1]], 'ro-', lw=3, label='Base AB')
        ax_mech.plot([A[0], C[0]], [A[1], C[1]], 'go-', lw=3, label='Link AC')
        ax_mech.plot([B[0], E[0]], [B[1], E[1]], 'bo-', lw=3, label='Link BE')
        ax_mech.plot([C[0], D[0]], [C[1], D[1]], 'k--', lw=2, label='Spring CD')
        ax_mech.plot([A[0], E[0]], [A[1], E[1]], 'y:', label='Distance AE')

        nodes = {'A': A, 'B': B, 'C': C, 'D': D, 'E': E}
        for name, pos in nodes.items():
            ax_mech.text(pos[0], pos[1] + 5, name, fontsize=10, fontweight='bold')

    ax_mech.set_aspect('equal')
    ax_mech.set_xlim(-60, L_AB + 100)
    ax_mech.set_ylim(-50, L_BE + 50)
    ax_mech.grid(True, linestyle=':')
    ax_mech.set_title(f'Mechanism Pose (AE = {L_ae:.1f}mm)')
    
    # 更新力曲线上的红点
    current_f = calculate_force_ae(L_ae)
    line_force.set_data([L_ae], [current_f])
    fig.canvas.draw_idle()

# --- 初始化界面 ---
fig, (ax_mech, ax_force) = plt.subplots(1, 2, figsize=(15, 6))
plt.subplots_adjust(bottom=0.2)

L_min, L_max = 39.1, 468.9
L_range = np.linspace(L_min, L_max, 500)
F_range = [calculate_force_ae(l) for l in L_range]

ax_force.plot(L_range, F_range, 'b-', label='Force $F_{AE}$ (N)')
ax_force.axhline(0, color='k', alpha=0.2)
line_force, = ax_force.plot([], [], 'ro')
ax_force.set_xlabel('AE Length (mm)')
ax_force.set_ylabel('Force (N)')
ax_force.grid(True, alpha=0.3)
ax_force.set_title('Static Force Characteristic')

ax_slider = plt.axes([0.2, 0.05, 0.6, 0.03])
slider = Slider(ax_slider, 'AE Length', L_min, L_max, valinit=250.0)
slider.on_changed(update_plot)

update_plot(250.0)
plt.show()