"""
작업1: RoboDK 정답지 추출 — 플랜지 FK + 회전행렬(R) 포함
★ robodk_ref_joints.json 에 저장:
   - fk_tcp_xyz: TCP 기준 위치 (tool=525.2mm 포함)
   - fk_flange_xyz: 플랜지 기준 위치 (tool=eye = 공구 없음)
   - fk_R: 플랜지 회전행렬 9원소 (so3 column-major)
★ Klampt 비교는 플랜지 기준으로 (양쪽 공구 없는 상태)

실행 전: RoboDK GUI 열린 상태 (headless 금지)
"""
import sys, json

print("=" * 65)
print("작업1: RoboDK 정답지 추출 (플랜지FK + R행렬 포함)")
print("=" * 65)

# ── 1. 연결 ──────────────────────────────────────────────────
print("\n[1] RoboDK 연결...")
try:
    from robodk.robolink import Robolink, ITEM_TYPE_ROBOT
    from robodk.robomath import transl, roty, rotx, rotz, eye, pi, Mat
    RDK = Robolink()
    print(f"    ✅ 버전: {RDK.Version()}")
except Exception as e:
    print(f"    ❌ {e}\n    → RoboDK GUI를 먼저 여세요.")
    sys.exit(1)

# ── 2. AR2010 + TCP ──────────────────────────────────────────
print("\n[2] AR2010 로드 + TCP 525.2mm...")
try:
    robot = RDK.Item('', ITEM_TYPE_ROBOT)
    if not robot.Valid():
        robot = RDK.AddFile(
            r'E:\도진팩토리\VAULT\robots\YASKAWA\Yaskawa-Motoman-AR2010.robot'
        )
    print(f"    로봇: {robot.Name()}")

    # TCP 도구 설정
    tcp_pose = transl(0, 0, 525.2)
    tool = robot.AddTool(tcp_pose, 'TCP_525')
    robot.setTool(tool)

    ref = RDK.AddFrame('WeldRef', robot)
    ref.setPose(eye(4))
    robot.setPoseFrame(ref)
    print(f"    ✅ TCP 525.2mm + 기준프레임 설정")
except Exception as e:
    print(f"    ❌ {e}")
    sys.exit(1)

# ── 3. 영점 자세 먼저 (규약 확인용) ──────────────────────────
print("\n[3] ★영점 자세(관절각=0) FK 출력 (Klampt 규약 비교 기준)...")
try:
    zero_joints = [0.0] * 6
    robot.setJoints(zero_joints)
    fk_zero_tcp = robot.SolveFK(zero_joints)
    robot.setTool(eye(4))
    fk_zero_flange = robot.SolveFK(zero_joints)
    robot.setTool(tool)

    pos_tcp = fk_zero_tcp.Pos()
    pos_flg = fk_zero_flange.Pos()
    print(f"    TCP 위치:     [{pos_tcp[0]:.2f}, {pos_tcp[1]:.2f}, {pos_tcp[2]:.2f}] mm")
    print(f"    플랜지 위치:  [{pos_flg[0]:.2f}, {pos_flg[1]:.2f}, {pos_flg[2]:.2f}] mm")
    print(f"    TCP-플랜지 차이: [{pos_tcp[0]-pos_flg[0]:.2f}, {pos_tcp[1]-pos_flg[1]:.2f}, {pos_tcp[2]-pos_flg[2]:.2f}] mm")
    print(f"    → Klampt FK는 플랜지 기준이므로 이 값으로 비교해야 함")
except Exception as e:
    print(f"    ❌ {e}")

# ── 4. 테스트 자세 8개 ─────────────────────────────────────
test_poses = [
    ("A_중앙",   transl(800,    0, 500) * roty(pi)),
    ("B_우측",   transl(800,  300, 500) * roty(pi)),
    ("C_좌측",   transl(800, -300, 500) * roty(pi)),
    ("D_위쪽",   transl(700,    0, 700) * roty(pi)),
    ("E_아래",   transl(900,    0, 300) * roty(pi)),
    ("F_전방",   transl(1000,   0, 400) * roty(pi)),
    ("G_대각1",  transl(800,  200, 600) * roty(pi) * rotx(0.2)),
    ("H_대각2",  transl(750, -200, 650) * roty(pi) * rotz(0.15)),
]

SEED_JOINTS = [0, -45, 0, 0, -90, 0]

print(f"\n[4] {len(test_poses)}개 자세 SolveIK → 플랜지FK + R 저장...")
print(f"    {'자세':10} {'플랜지xyz(mm)':32} {'관절각(도)':40} {'오차':6}")
print("    " + "-"*95)

results = []
for name, pose in test_poses:
    try:
        pos_tgt = pose.Pos()
        robot.setJoints(SEED_JOINTS)

        # IK (TCP 기준 자세로)
        jmat = robot.SolveIK(pose)
        jlist = jmat.list()
        if len(jlist) == 0:
            print(f"    {name:10} IK 해 없음")
            results.append({'name': name, 'ok': False, 'reason': 'no_ik'})
            continue

        j6 = jlist[:6]

        # FK — TCP 기준 (원래 tool)
        fk_tcp = robot.SolveFK(jlist)
        pos_tcp_fk = fk_tcp.Pos()
        err_tcp = sum((a-b)**2 for a, b in zip(pos_tgt, pos_tcp_fk)) ** 0.5

        # FK — 플랜지 기준 (tool=eye)
        robot.setTool(eye(4))
        fk_flange = robot.SolveFK(jlist)
        robot.setTool(tool)

        pos_flg = fk_flange.Pos()

        # 회전행렬 R (플랜지 기준, 9원소 column-major for so3)
        # fk_flange[row, col]: 4x4 Mat, 첫 3x3이 회전
        # so3 column-major: [r00,r10,r20, r01,r11,r21, r02,r12,r22]
        R9 = []
        for col in range(3):
            for row in range(3):
                R9.append(round(float(fk_flange[row, col]), 8))

        # IK 브랜치
        try:
            cfg = robot.JointsConfig(jlist)
            branch = str(cfg.list()) if hasattr(cfg, 'list') else str(cfg)
        except Exception:
            branch = "unknown"

        j6r = [round(float(j), 4) for j in j6]
        flg_str = f"[{pos_flg[0]:.1f},{pos_flg[1]:.1f},{pos_flg[2]:.1f}]"
        j_str = str([round(j, 1) for j in j6r])
        status = "✅" if err_tcp < 1.0 else "⚠️"
        print(f"    {name:10} {flg_str:32} {j_str:40} {err_tcp:.4f}mm {status}")

        results.append({
            'name': name,
            'target_tcp_xyz': [round(p, 3) for p in pos_tgt],
            'joints_deg': j6r,
            'fk_tcp_xyz': [round(p, 3) for p in pos_tcp_fk],
            'fk_flange_xyz': [round(p, 3) for p in pos_flg],
            'fk_R_so3': R9,
            'tcp_err_mm': round(err_tcp, 4),
            'branch': branch,
            'ok': err_tcp < 1.0
        })

    except Exception as e:
        import traceback
        print(f"    {name:10} ❌ {type(e).__name__}: {e}")
        traceback.print_exc()
        results.append({'name': name, 'ok': False, 'reason': str(e)})

# ── 5. 저장 ──────────────────────────────────────────────────
ok_list = [r for r in results if r.get('ok')]
out = {
    'robot': 'AR2010', 'tcp_mm': 525.2,
    'seed_joints': SEED_JOINTS,
    'note': 'fk_flange_xyz + fk_R_so3 = Klampt 비교 기준 (플랜지, 공구 없음)',
    'total': len(results), 'ok': len(ok_list),
    'results': results
}
out_path = r'E:\도진팩토리\3D스캔및티칭시스템\robodk_ref_joints.json'
with open(out_path, 'w', encoding='utf-8') as f:
    json.dump(out, f, ensure_ascii=False, indent=2)

print(f"\n[5] ✅ 저장: {out_path}")
print(f"★ 결과: {len(ok_list)}/{len(results)} 자세 TCP 오차 1mm 이내")
print(f"★ fk_flange_xyz + fk_R_so3 포함 → Klampt FK게이트 비교 가능")
print("=" * 65)