# -*- coding: utf-8 -*-
"""修正版 6 轴机械臂建模：几何方向全部按物理正确重建
J1 底座竖直 / J2 肩横轴 / 大臂上行 / J3 肘横轴 / 小臂前伸 / J4 前臂自转 / J5 腕摆 / J6 法兰自转
"""
import os, math, json
import FreeCAD as App
import Part, Mesh, Import

OUT = os.path.dirname(os.path.abspath(__file__))
PARTS = os.path.join(OUT, "parts")
os.makedirs(PARTS, exist_ok=True)

D = dict(
    base_r=140, plate_h=20, base_h=150,     # J1 底座
    waist_h=70, shoulder_z=250,             # 腰部 + J2 肩高度
    L1=520,                                 # 大臂 J2->J3
    L2=460,                                 # 小臂 J3->J5
    L3=130,                                 # 腕 J5->法兰
    L4=95,                                  # 法兰->TCP
    jr=68,                                  # 大关节半径
    jr_w=44,                                # 腕关节半径
)

doc = App.newDocument("Robot6Axis")

def add(name, shape, color):
    o = doc.addObject("Part::Feature", name)
    o.Shape = shape
    o.addProperty("App::PropertyString", "ColorHex", "design")
    o.ColorHex = color
    return o

def cyl(r, h, base=(0,0,0), axis=(0,0,1)):
    return Part.makeCylinder(r, h, App.Vector(*base), App.Vector(*axis))

# ===== J1 底座 =====
plate = cyl(D["base_r"], D["plate_h"])
for a in (45,135,225,315):
    x = (D["base_r"]-24)*math.cos(math.radians(a)); y = (D["base_r"]-24)*math.sin(math.radians(a))
    plate = plate.cut(cyl(7, D["plate_h"]+10, (x,y,-5)))
base = plate.fuse(cyl(D["base_r"]*0.70, D["base_h"]-D["plate_h"], (0,0,D["plate_h"])))
add("J1_底座", base, "#4a5560")

# ===== J1 回转腰 =====
waist = cyl(D["base_r"]*0.62, D["waist_h"], (0,0,D["base_h"]))
# 肩部立柱（腰部以上到 J2 轴心）
pylon = Part.makeBox(170, 150, D["shoulder_z"]-(D["base_h"]+D["waist_h"]),
                     App.Vector(-85,-75,D["base_h"]+D["waist_h"]))
waist = waist.fuse(pylon)
add("J2_腰部", waist, "#e8b923")   # 工业黄

# ===== J2 肩关节（横置，轴线沿 Y）=====
j2 = cyl(D["jr"], 190, (0,-95,D["shoulder_z"]), (0,1,0))
add("J2_肩关节", j2, "#2b2f33")

# ===== 大臂 J2->J3（向上前倾 62°）=====
a_t = math.radians(62)   # 与水平面夹角
j3 = App.Vector(D["L1"]*math.cos(a_t), 0, D["shoulder_z"]+D["L1"]*math.sin(a_t))
arm = Part.makeBox(D["L1"], 88, 104, App.Vector(0,-44,-52))
for i in range(3):   # 减重孔
    hole = cyl(27, 140, (110+i*115, -70, 0))
    arm = arm.cut(hole)
arm.rotate(App.Vector(0,0,0), App.Vector(0,1,0), -math.degrees(a_t))
arm.translate(App.Vector(0,0,D["shoulder_z"]))
add("J3_大臂", arm, "#d8dbe0")

# ===== J3 肘关节（横轴）=====
add("J3_肘关节", cyl(D["jr"]*0.82, 150, (j3.x,-75,j3.z), (0,1,0)), "#2b2f33")

# ===== 小臂 J3->J5（前伸，略下倾 10°）=====
a_f = math.radians(-10)
j5 = App.Vector(j3.x + D["L2"]*math.cos(a_f), 0, j3.z + D["L2"]*math.sin(a_f))
fore = Part.makeBox(D["L2"], 74, 86, App.Vector(0,-37,-43))
for i in range(2):
    fore = fore.cut(cyl(21, 120, (130+i*150, -60, 0)))
fore.rotate(App.Vector(0,0,0), App.Vector(0,1,0), -math.degrees(a_f))
fore.translate(j3)
add("J4_小臂", fore, "#d8dbe0")

# ===== J4 前臂自转（沿小臂轴向）=====
ax4 = App.Vector(math.cos(a_f), 0, math.sin(a_f))
add("J4_旋转关节", cyl(D["jr_w"], 110, (j5.x - ax4.x*55, 0, j5.z - ax4.z*55), ax4), "#2b2f33")

# ===== J5 腕摆（横轴）=====
add("J5_腕关节", cyl(D["jr_w"]*0.92, 120, (j5.x, -60, j5.z), (0,1,0)), "#2b2f33")

# ===== 腕体 + 法兰（沿延伸方向）=====
fl = App.Vector(j5.x + ax4.x*D["L3"], 0, j5.z + ax4.z*D["L3"])
wrist = cyl(D["jr_w"]*0.86, D["L3"], (j5.x + ax4.x*55, 0, j5.z + ax4.z*55), ax4)
add("腕体", wrist, "#4a5560")

flange = cyl(48, 18, (fl.x - ax4.x*9, 0, fl.z - ax4.z*9), ax4)
for a in (45,135,225,315):   # 4×M6 螺栓孔
    off = App.Vector(0, 30*math.cos(math.radians(a)), 30*math.sin(math.radians(a)))
    flange = flange.cut(cyl(3.2, 26, (fl.x - ax4.x*11 + off.x, off.y, fl.z - ax4.z*11 + off.z), ax4))
add("法兰ISO9409-50-4-M6", flange, "#8d949c")

# ===== J6 末端夹爪 =====
tcp0 = App.Vector(fl.x + ax4.x*10, 0, fl.z + ax4.z*10)
tool = cyl(33, D["L4"]-30, (tcp0.x, 0, tcp0.z), ax4)
for s in (-1,1):   # 两指
    fing = Part.makeBox(64, 15, 44, App.Vector(0, s*13-7.5, -22))
    fing.rotate(App.Vector(0,0,0), App.Vector(0,1,0), -math.degrees(a_f))
    fing.translate(App.Vector(tcp0.x + ax4.x*(D["L4"]-30), 0, tcp0.z + ax4.z*(D["L4"]-30)))
    tool = tool.fuse(fing)
add("J6_夹爪", tool, "#c0392b")

tcp = App.Vector(fl.x + ax4.x*D["L4"], 0, fl.z + ax4.z*D["L4"])

doc.recompute()

# ---- 导出 ----
meta = []
merged = Mesh.Mesh()
for o in doc.Objects:
    m = Mesh.Mesh(o.Shape.tessellate(0.35))
    m.write(os.path.join(PARTS, o.Name + ".obj"))
    o.Shape.exportStl(os.path.join(PARTS, o.Name + ".stl"))
    merged.addMesh(m)
    bb = o.Shape.BoundBox
    meta.append(dict(name=o.Name, color=o.ColorHex, verts=len(m.Points), faces=len(m.Facets),
                     vol=round(o.Shape.Volume/1e6, 4),
                     bb=[round(v,1) for v in (bb.XMin,bb.YMin,bb.ZMin,bb.XMax,bb.YMax,bb.ZMax)]))
merged.write(os.path.join(OUT, "robot_arm.obj"))
Import.export(doc.Objects, os.path.join(OUT, "robot_arm.step"))
doc.saveAs(os.path.join(OUT, "robot_arm.FCStd"))

kin = dict(params=D, joints=dict(
    J2=[0,0,round(D["shoulder_z"],1)], J3=[round(j3.x,1),0,round(j3.z,1)],
    J5=[round(j5.x,1),0,round(j5.z,1)], flange=[round(fl.x,1),0,round(fl.z,1)],
    tcp=[round(tcp.x,1),0,round(tcp.z,1)],
    reach_lateral=round(tcp.x,1), reach_height=round(tcp.z,1)))
json.dump(dict(meta=meta, kin=kin), open(os.path.join(OUT,"model_data.json"),"w"), indent=1, ensure_ascii=False)

print("=== BUILD OK ===")
print("零件:", len(doc.Objects))
print("J2 肩轴:", [0,0,D["shoulder_z"]], " J3 肘:", [round(j3.x,1),round(j3.z,1)])
print("J5 腕:", [round(j5.x,1),round(j5.z,1)], " TCP:", [round(tcp.x,1),round(tcp.z,1)])
print("最大水平伸展:", round(tcp.x,1), "mm   TCP 高度:", round(tcp.z,1), "mm")
print("总体积:", round(sum(o.Shape.Volume for o in doc.Objects)/1e6,3), "L")