# -*- coding: utf-8 -*-
"""关节电机力矩与负载能力校核。

与旧版的区别（这些是旧版实际出错的地方）：

  1. 拓扑。旧版把关节表手抄成 "J3 肘俯仰 / J4 前臂自转 / J5 腕俯仰 / J6 夹爪"，
     并据此虚构了一个夹爪机构和 135mm 大臂——与 human_layout.py 的实际拓扑
     （J3 大臂自转 / J4 肘俯仰 / J6 腕自转）完全不符，整份报告的结论都是错的。
     本版不再复制拓扑：零件表、关节轴、运动学链全部复用 strength_check，
     而 strength_check 的轴又直接来自 human_layout，单一事实源。
  2. 力臂。旧版把每个关节的质量与力臂写成硬编码常数（J2 = 1.75kg @ 464mm），
     几何一改就失效且不会报错。本版按每个零件的真实质心、在可达姿态网格上
     搜索最大重力力矩。
  3. 惯量。旧版用一组手填的 kg·m²。本版按 Σm·d²（对关节轴的平行轴距离）
     在同一姿态网格上取最大值。零件自身转动惯量未计入，对大臂量级的关节
     是小修正，已在报告中标注。

电机参数：仓库内没有原始数据手册，以下为工程假定，可用环境变量覆盖。
    ARMT_CONT  连续额定输出力矩 (N·m)
    ARMT_PEAK  峰值堵转输出力矩 (N·m)
    ARMT_ALLOW 封闭塑料外壳热衰减后的许用连续力矩 (N·m)
    ARMT_ALPHA 角加速度 (rad/s²)
    ARMPAY     末端负载 (kg)

输出：docs/torque_balance.json + docs/torque_report.md
"""
import io
import json
import math
import os
import sys
try:
    sys.stdout.reconfigure(encoding="utf-8")
except Exception:
    pass
from arm import kinematics as HL
from arm.paths import DOCS
from arm.analysis import strength_check as SC
G = 9.81
T_CONT = float(os.environ.get("ARMT_CONT", "2.5"))
T_PEAK = float(os.environ.get("ARMT_PEAK", "8.0"))
T_ALLOW = float(os.environ.get("ARMT_ALLOW", "2.0"))
ALPHA = float(os.environ.get("ARMT_ALPHA", "5.0"))
PAYLOAD = float(os.environ.get("ARMPAY", "0.5"))

# 关节名与轴向，全部取自 human_layout（不手抄）
JOINT_CN = {t[0]: t[1] for t in HL.JOINTS}
JOINT_D = {t[0]: t[3] for t in HL.JOINTS}


def _axis_kind(tag):
    d = JOINT_D[tag]
    if abs(d[2]) == 1:
        return "roll(绕Z)"
    if abs(d[0]) == 1:
        return "roll(绕X)"
    return "pitch(绕Y)"


def analyze(payload_kg=PAYLOAD):
    parts = SC.part_table(payload_kg)
    jl = SC.joint_loads(parts)
    jl0 = SC.joint_loads(SC.part_table(0.0))   # 去掉末端负载，用于反推负载能力
    masses = {"J%d" % j: sum(m for _, m, _, g in parts if g >= j) for j in range(1, 7)}
    arm_mass = {k: v - payload_kg * 1000.0 for k, v in masses.items()}

    res = {"assumptions": {
        "motor": "60mm 方形标准关节电机模组，2 级行星减速",
        "T_cont_Nm": T_CONT, "T_peak_Nm": T_PEAK, "T_allow_Nm": T_ALLOW,
        "alpha_rad_s2": ALPHA, "payload_kg": payload_kg,
        "note": "电机参数为工程假定，仓库内无数据手册；结论对此敏感，见报告第 4 节",
    }, "joints": {}}

    for j in range(1, 7):
        tag = "J%d" % j
        r = jl[tag]
        tau_g = r["tau"] / 1000.0                  # N*m，重力
        tau_dyn = r["I"] * ALPHA                   # N*m，惯性
        tau_g_arm = jl0[tag]["tau"] / 1000.0       # 去掉负载后的重力力矩
        req_static = tau_g
        req_peak = tau_g + tau_dyn
        res["joints"][tag] = {
            "cn": JOINT_CN[tag], "axis": _axis_kind(tag),
            "distal_mass_kg": round(masses[tag] / 1000.0, 3),
            "arm_mass_kg": round(arm_mass[tag] / 1000.0, 3),
            "weight_N": round(r["F"], 1),
            "reach_tip_mm": round(r["Lp"], 0),
            "tau_gravity_Nm": round(tau_g, 3),
            "tau_gravity_arm_only_Nm": round(tau_g_arm, 3),
            "I_distal_kgm2": round(r["I"], 5),
            "tau_dynamic_Nm": round(tau_dyn, 3),
            "tau_required_static_Nm": round(req_static, 3),
            "tau_required_peak_Nm": round(req_peak, 3),
            "duty_vs_Tcont_pct": round(100.0 * req_static / T_CONT, 1),
            "duty_vs_Tallow_pct": round(100.0 * req_static / T_ALLOW, 1),
            "peak_margin_x": round(T_PEAK / req_peak, 2) if req_peak > 1e-3 else None,
            "status": "OK" if (req_static <= T_ALLOW and req_peak <= T_PEAK) else "OVERLOAD",
        }
    return res, jl0, jl


def max_payload(res):
    """在许用连续力矩下，末端还能带多少 kg（按各关节取最小）。

    线性化：tau_g(p) ≈ tau_g_arm + p·g·L。用的是该关节实测的最大 TCP 力臂，
    即最不利姿态下的力臂。
    """
    out = {}
    for tag, d in res["joints"].items():
        L = d["reach_tip_mm"] / 1000.0
        if L <= 0:
            out[tag] = None
            continue
        head = T_ALLOW - d["tau_gravity_arm_only_Nm"]
        out[tag] = round(max(head, 0.0) / (G * L), 3) if head > 0 else 0.0
    return out


def generate_report(res, mp):
    L = []
    L.append("# 6 自由度机械臂 · 关节电机力矩与负载能力校核\n")
    a = res["assumptions"]
    L.append("## 1. 前提\n")
    L.append("- 驱动：%s" % a["motor"])
    L.append("- 连续额定输出力矩 `T_cont` = %.2f N·m；峰值堵转 `T_peak` = %.2f N·m"
             % (a["T_cont_Nm"], a["T_peak_Nm"]))
    L.append("- 封闭打印外壳热衰减后许用连续力矩 `T_allow` = %.2f N·m；末端负载 %.2f kg；"
             "角加速度 %.1f rad/s²" % (a["T_allow_Nm"], a["payload_kg"], a["alpha_rad_s2"]))
    L.append("- **%s**\n" % a["note"])
    L.append("载荷模型：每个零件按网格体积加权质心定位，随姿态做刚体变换，在可达姿态"
             "网格（俯仰轴取行程端点/中值，自转轴取 0/±60/±120/180°）上取重力力矩最大值。"
             "力矩为绕**关节自身轴**的分量——这才是电机要出的力。\n")

    L.append("## 2. 各关节校核结果\n")
    L.append("| 关节 | 名称 | 轴 | 以下质量(kg) | TCP力臂(mm) | 重力力矩(N·m) | 惯性力矩(N·m) | 峰值需求(N·m) | 负荷率(对T_allow) | 峰值裕度 | 判定 |")
    L.append("|:---:|:---|:---|---:|---:|---:|---:|---:|---:|---:|:---:|")
    for tag, d in res["joints"].items():
        L.append("| %s | %s | %s | %.2f | %.0f | **%.2f** | %.2f | %.2f | **%.0f%%** | %s | %s |"
                 % (tag, d["cn"], d["axis"], d["distal_mass_kg"], d["reach_tip_mm"],
                    d["tau_gravity_Nm"], d["tau_dynamic_Nm"], d["tau_required_peak_Nm"],
                    d["duty_vs_Tallow_pct"],
                    ("%.2f×" % d["peak_margin_x"]) if d["peak_margin_x"] else "—", d["status"]))
    L.append("")
    L.append("## 3. 许用连续工况下的末端负载能力\n")
    L.append("去掉末端负载后各关节仍需 `tau_gravity_arm_only` 的重力力矩，余额除以该关节"
             "实测最大 TCP 力臂，即为还能挂的负载：\n")
    L.append("| 关节 | 仅臂身重力力矩(N·m) | 余额(N·m) | 可挂负载(kg) |")
    L.append("|:---:|---:|---:|---:|")
    for tag, d in res["joints"].items():
        head = T_ALLOW - d["tau_gravity_arm_only_Nm"]
        L.append("| %s | %.2f | %+.2f | %s |"
                 % (tag, d["tau_gravity_arm_only_Nm"], head,
                    ("%.2f" % mp[tag]) if mp[tag] is not None else "—"))
    ok = [t for t, v in mp.items() if v is not None]
    worst = min(ok, key=lambda t: mp[t]) if ok else None
    if worst:
        L.append("")
        L.append("**整臂瓶颈：%s，可挂末端负载 %.2f kg。**\n" % (worst, mp[worst]))
    L.append("## 4. 结论要点\n")
    over = [t for t, d in res["joints"].items() if d["status"] == "OVERLOAD"]
    if over:
        L.append("- **超载关节：%s**。" % "、".join(over))
        for t in over:
            d = res["joints"][t]
            L.append("  - %s（%s）重力 %.2f N·m，超 `T_allow` %.0f%%；即使按峰值堵转算，"
                     "裕度也只有 %s。" % (t, d["cn"], d["tau_gravity_Nm"],
                                            d["duty_vs_Tallow_pct"] - 100.0,
                                            ("%.2f×" % d["peak_margin_x"]) if d["peak_margin_x"] else "—"))
        L.append("- 结论对电机参数高度敏感：若实际连续额定高于假定值，"
                 "请用 `ARMT_CONT` / `ARMT_ALLOW` 覆盖后重跑本脚本。")
    else:
        L.append("- 全部关节在许用连续力矩内。")
    L.append("- 惯量按 Σm·d² 估，**未计入零件自身转动惯量**，对肩/肘是大臂量级的关节"
             "约低估一成，惯性力矩一栏应视为下限。")
    L.append("- 本版不含被动弹簧平衡。旧版报告提出的 J2 拉簧预补偿（2.48 N·m）在几何上"
             "**尚未实现**；若要保留该方案，需要在 P04/P05 之间增加拉簧连杆机构。")
    return "\n".join(L)


def main():
    res, jl0, jl = analyze(PAYLOAD)
    mp = max_payload(res)
    io.open(DOCS + "/torque_balance.json", "w", encoding="utf-8").write(
        json.dumps(res, ensure_ascii=False, indent=2))
    io.open(DOCS + "/torque_report.md", "w", encoding="utf-8").write(generate_report(res, mp))

    print("末端负载 %.2f kg   T_cont=%.2f  T_allow=%.2f  T_peak=%.2f  alpha=%.1f\n"
          % (PAYLOAD, T_CONT, T_ALLOW, T_PEAK, ALPHA))
    print("%-4s %-8s %-12s %8s %8s %8s %9s %8s  %s"
          % ("关节", "名称", "轴", "重力N·m", "惯性N·m", "峰值N·m", "负荷率", "峰值裕度", "判定"))
    for tag, d in res["joints"].items():
        print("%-4s %-8s %-12s %8.2f %8.2f %8.2f %8.0f%% %8s  %s"
              % (tag, d["cn"], d["axis"], d["tau_gravity_Nm"], d["tau_dynamic_Nm"],
                 d["tau_required_peak_Nm"], d["duty_vs_Tallow_pct"],
                 ("%.2fx" % d["peak_margin_x"]) if d["peak_margin_x"] else "-",
                 d["status"]))
    print("\n许用连续力矩下可挂末端负载 (kg):")
    for tag in res["joints"]:
        print("  %-4s %s" % (tag, ("%.2f" % mp[tag]) if mp[tag] is not None else "—"))
    ok = [t for t, v in mp.items() if v is not None]
    if ok:
        w = min(ok, key=lambda t: mp[t])
        print("  整臂瓶颈 %s -> %.2f kg" % (w, mp[w]))
    print("\n已写出 docs/torque_balance.json + docs/torque_report.md")
if __name__ == "__main__":
    main()

    # OCC/cadquery 的静态析构在解释器正常退出时必定段错误（实测：光是
    # `python -c "import arm.cad.helpers"` 就返回 139，与本项目代码无关）。所以入口
    # 脚本跑完后显式刷缓冲再 os._exit —— 否则全流程已经成功、exit code 却是 139,
    # 会把上游脚本的 set -e 判成失败。只在 __main__ 出口生效, 被 import 时不触发。
    sys.stdout.flush()
    os._exit(0)