"""
작업2B+3: Klampt IK → FK + RoboDK 정답지 비교 → DB 축적
★ 전제: test_klampt_fk_gate.py FK 게이트 통과 후 실행
★ 비교: 위치오차(mm) + 방향오차(deg) — 관절각 직접비교 금지
★ 방향오차: so3.distance() (geodesic, 오일러각 차이 X)

실행: python test_klampt_ik_compare.py
"""
import sys, json, os, math

print("=" * 65)
print("작업2B+3: Klampt IK 비교 + DB 축적")
print("=" * 65)

# ── 임포트 ───────────────────────────────────────────────────
try:
    import klampt
    from klampt import WorldModel
    from klampt.model import ik
    from klampt.math import so3, se3
    print(f"[init] klampt {klampt.__version__} ✅")
except ImportError as e:
    print(f"❌ {e}  → pip install klampt"); sys.exit(1)

# ── URDF 로드 ─────────────────────────────────────────────────
URDF_CANDIDATES = [
    # ★ Klampt는 한국어 경로 처리 불가 → 영문 경로 사용
    r'C:\Users\user\AppData\Local\klampt_robots\ar2010_kinematic.urdf',
]
urdf_path = next((p for p in URDF_CANDIDATES if os.path.exists(p)), None)
if not urdf_path:
    print("❌ URDF 없음 → test_klampt_fk_gate.py 먼저 실행"); sys.exit(1)

world = WorldModel()
world.loadRobot(urdf_path)
robot_k = world.robot(0)
ee_link = robot_k.numLinks() - 1
print(f"[init] URDF 로드: {os.path.basename(urdf_path)}, DOF={robot_k.numDrivers()}")

# ── RoboDK 정답지 로드 ───────────────────────────────────────
ref_path = r'E:\도진팩토리\3D스캔및티칭시스템\robodk_ref_joints.json'
if not os.path.exists(ref_path):
    print("❌ robodk_ref_joints.json 없음 → test_robodk_simple_joints.py 먼저"); sys.exit(1)
with open(ref_path, 'r', encoding='utf-8') as f:
    ref = json.load(f)
ok_poses = [r for r in ref['results'] if r.get('ok')]
print(f"[init] RoboDK 정답지: {len(ok_poses)}개 자세\n")

def mat3x3_from_robodk_pose(pose_mat_flat):
    """RoboDK 4x4 Mat에서 3x3 회전행렬 추출 (so3 형식: column-major 9원소)"""
    # RoboDK Pose: 4x4, 첫 3열이 회전
    # se3 형식 R은 column-major [r00,r10,r20, r01,r11,r21, r02,r12,r22]
    # 여기선 fk_pose를 직접 못 가져오므로 joints → Klampt FK의 R 사용
    pass  # 사용 안 함 — 아래에서 Klampt FK R만 사용

print(f"{'자세':10} {'목표xyz':26} {'RDK관절(deg)':35} {'Klp관절(deg)':35} {'위치오차':8} {'방향오차':8}")
print("-" * 130)

db = []
for r in ok_poses:
    name = r['name']
    tgt_xyz = r['fk_flange_xyz']  # mm, 플랜지 기준 (TCP 오프셋 제거)
    rdk_j = r['joints_deg']       # 6축 관절각 (도)

    try:
        # ── Klampt IK ────────────────────────────────────────
        tgt_m = [v / 1000.0 for v in tgt_xyz]  # mm → m

        goal = ik.objective(robot_k.link(ee_link),
                            local=[0, 0, 0], world=tgt_m)
        solved = ik.solve(goal, iters=200, tol=1e-5)

        if not solved:
            print(f"{name:10} IK 해 없음")
            db.append({'name': name, 'ok': False, 'reason': 'klampt_no_ik'})
            continue

        # Klampt 관절각 (라디안 → 도)
        q_rad = robot_k.getConfig()
        klp_j = [math.degrees(q_rad[robot_k.driver(i).getAffectedLink()]) for i in range(min(6, robot_k.numDrivers()))]

        # ── Klampt FK ────────────────────────────────────────
        T_k = robot_k.link(ee_link).getTransform()
        R_k, t_k = T_k
        klp_fk_mm = [v * 1000.0 for v in t_k]  # m → mm

        # ── RoboDK FK (정답지에 저장된 값) ───────────────────
        rdk_fk_mm = r['fk_flange_xyz']  # 플랜지 기준

        # ── 위치오차 (mm) ─────────────────────────────────────
        pos_err = sum((a-b)**2 for a,b in zip(rdk_fk_mm, klp_fk_mm)) ** 0.5

        # ── 방향오차 (deg) — geodesic so3.distance ────────────
        # RoboDK 방향: 정답지에 R 행렬이 없으므로 Klampt FK R만 있음
        # → 현재는 위치오차만 비교, 방향오차는 URDF 로드 후 R도 추출 시 추가
        # ★ 추후: robodk_ref_joints.json에 R_mat 추가하면 so3.distance 사용
        orient_err_deg = None  # 정답지 R 없음 (향후 확장)

        rdk_j_str = str([round(j,1) for j in rdk_j])
        klp_j_str = str([round(j,1) for j in klp_j])
        tgt_str = f"[{tgt_xyz[0]:.0f},{tgt_xyz[1]:.0f},{tgt_xyz[2]:.0f}]"
        status = "✅" if pos_err < 5.0 else "⚠️"
        print(f"{name:10} {tgt_str:26} {rdk_j_str:35} {klp_j_str:35} {pos_err:6.2f}mm {status}")

        db.append({
            'name': name,
            'target_xyz_mm': tgt_xyz,
            'robodk': {
                'joints_deg': rdk_j,
                'fk_xyz_mm': rdk_fk_mm,
                'branch': r.get('branch', ''),
            },
            'klampt': {
                'joints_deg': [round(j, 4) for j in klp_j],
                'fk_xyz_mm': [round(v, 3) for v in klp_fk_mm],
            },
            'comparison': {
                'pos_err_mm': round(pos_err, 4),
                'orient_err_deg': orient_err_deg,
                'note': 'orient_err 미구현 — robodk_ref_joints에 R_mat 추가 후 확장'
            },
            'ok': pos_err < 5.0
        })

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

# ── DB 저장 ───────────────────────────────────────────────────
ok_db = [d for d in db if d.get('ok')]
out = {
    'robot': 'AR2010', 'tcp_mm': 525.2,
    'urdf': os.path.basename(urdf_path),
    'total': len(db), 'ok': len(ok_db),
    'comparison_db': db
}
db_path = r'E:\도진팩토리\3D스캔및티칭시스템\comparison_db.json'
with open(db_path, 'w', encoding='utf-8') as f:
    json.dump(out, f, ensure_ascii=False, indent=2)

print(f"\n★ 비교 결과: {len(ok_db)}/{len(db)} 자세 일치 (기준: 5mm)")
print(f"★ DB 저장: {db_path}")
if ok_db:
    s = ok_db[0]
    print(f"\n  샘플 [{s['name']}]:")
    print(f"    RoboDK 관절각: {s['robodk']['joints_deg']}")
    print(f"    Klampt  관절각: {s['klampt']['joints_deg']}")
    print(f"    위치오차: {s['comparison']['pos_err_mm']}mm")
    print(f"    → 관절각이 달라도 위치오차 작으면 다른 IK브랜치 (정상)")
print("=" * 65)