"""
Kabsch 검증: viewer_fk_TCP(103점) vs 세그 좌표(103점)
→ 변환 T 확인 (RMSE가 작으면 세그 = 로봇 좌표계 TCP, T=identity)

두 가지 FK 방식 비교:
  v1: identity 시작, sign없음, U rpy=[PI,0,0], R y=-1082, T y=-100 rpy=[PI/2,0,0]
      tcpOffset = [-133.246, -10.813, 504.281]
  v4: mTr(0,0,-505) 시작, sign:-1 (U,R,B,T), R y=+1082, T y=+100 rpy=[-PI/2,0,0]
      tcpOffset = [-136.826, +8.192, +506.982]  ← VERIFIED_LOCK §2 확정값

VERIFIED_LOCK.md §2: 후보B 4점 0.002mm 이내 (v4 방식)
VERIFIED_LOCK.md §10: v1 방식 메인, 4점 최대 0.0016mm
"""
import json, math, sys
import numpy as np

PI = math.pi

# ── 행렬 유틸 ─────────────────────────────────────────────────
def mI():
    return np.eye(4)

def mTr(x, y, z):
    m = np.eye(4)
    m[0,3]=x; m[1,3]=y; m[2,3]=z
    return m

def mRx(a):
    c,s = math.cos(a), math.sin(a)
    return np.array([[1,0,0,0],[0,c,-s,0],[0,s,c,0],[0,0,0,1]], dtype=float)

def mRy(a):
    c,s = math.cos(a), math.sin(a)
    return np.array([[c,0,s,0],[0,1,0,0],[-s,0,c,0],[0,0,0,1]], dtype=float)

def mRz(a):
    c,s = math.cos(a), math.sin(a)
    return np.array([[c,-s,0,0],[s,c,0,0],[0,0,1,0],[0,0,0,1]], dtype=float)

def mRxyz(rx, ry, rz):
    return mRx(rx) @ mRy(ry) @ mRz(rz)

# ── 뷰어 FK v1 (VERIFIED_LOCK §10) ───────────────────────────
URDF_V1 = [
    {'a':'S', 'xyz':[0, 0, 505],    'rpy':[0, 0, 0]},
    {'a':'L', 'xyz':[150, 0, 0],    'rpy':[-PI/2, 0, -PI/2]},
    {'a':'U', 'xyz':[760, 0, 0],    'rpy':[PI, 0, 0]},
    {'a':'R', 'xyz':[200, -1082, 0],'rpy':[-PI/2, 0, 0]},
    {'a':'B', 'xyz':[0, 0, 0],      'rpy':[PI/2, 0, 0]},
    {'a':'T', 'xyz':[0, -100, 0],   'rpy':[PI/2, 0, 0]},
]
TCP_V1 = [-133.246, -10.813, 504.281]

def fk_v1(joints_deg):
    j = dict(zip(['S','L','U','R','B','T'], joints_deg))
    m = mI()
    for u in URDF_V1:
        rx,ry,rz = u['rpy']
        ax,ay,az = u['xyz']
        angle_rad = math.radians(j[u['a']])
        jt = mTr(ax,ay,az) @ mRxyz(rx,ry,rz) @ mRz(angle_rad)
        m = m @ jt
    m = m @ mTr(*TCP_V1)
    return m[:3, 3].tolist()   # mm

# ── 뷰어 FK v4 (VERIFIED_LOCK §7) ────────────────────────────
URDF_V4 = [
    {'a':'S', 'xyz':[0, 0, 505],    'rpy':[0, 0, 0],         'sgn':+1},
    {'a':'L', 'xyz':[150, 0, 0],    'rpy':[-PI/2, 0, -PI/2], 'sgn':+1},
    {'a':'U', 'xyz':[760, 0, 0],    'rpy':[0, 0, 0],          'sgn':-1},
    {'a':'R', 'xyz':[200, +1082, 0],'rpy':[-PI/2, 0, 0],      'sgn':-1},
    {'a':'B', 'xyz':[0, 0, 0],      'rpy':[PI/2, 0, 0],       'sgn':-1},
    {'a':'T', 'xyz':[0, +100, 0],   'rpy':[-PI/2, 0, 0],      'sgn':-1},
]
TCP_V4 = [-136.826, +8.192, +506.982]

def fk_v4(joints_deg):
    j = dict(zip(['S','L','U','R','B','T'], joints_deg))
    m = mTr(0, 0, -505)   # ← v4: mTr(0,0,-505) 시작
    for u in URDF_V4:
        rx,ry,rz = u['rpy']
        ax,ay,az = u['xyz']
        angle_rad = math.radians(j[u['a']] * u['sgn'])
        jt = mTr(ax,ay,az) @ mRxyz(rx,ry,rz) @ mRz(angle_rad)
        m = m @ jt
    m = m @ mTr(*TCP_V4)
    return m[:3, 3].tolist()   # mm

# ── Kabsch 알고리즘 ───────────────────────────────────────────
def kabsch(src_pts, dst_pts):
    """src → dst 강체변환 R, t, RMSE"""
    S = np.array(src_pts); D = np.array(dst_pts)
    cS = S.mean(0); cD = D.mean(0)
    A = S - cS; B = D - cD
    H = A.T @ B
    U, _, Vt = np.linalg.svd(H)
    V = Vt.T
    d = np.linalg.det(V @ U.T)
    R = V @ np.diag([1, 1, d]) @ U.T
    t = cD - R @ cS
    res = np.linalg.norm((R @ S.T).T + t - D, axis=1)
    return R, t, float(np.sqrt(np.mean(res**2))), float(res.max()), res

# ── 데이터 로드 ───────────────────────────────────────────────
wp_path = r'E:\도진팩토리\3D스캔및티칭시스템\weld_path.json'
with open(wp_path, 'r', encoding='utf-8') as f:
    wp = json.load(f)

arc_poses = [p for p in wp['poses'] if p.get('arc') == 1]
segs_flat = []
for seg in wp['segs']:
    segs_flat.extend(seg)

print(f"arc=1 포즈: {len(arc_poses)}개")
print(f"세그 점수:  {len(segs_flat)}개")

if len(arc_poses) != len(segs_flat):
    print(f"⚠️  개수 불일치! arc1={len(arc_poses)}, seg={len(segs_flat)}")
    print("     arc=1 포즈와 세그 순서 매핑 확인 필요")

n = min(len(arc_poses), len(segs_flat))

# ── v1 FK 103점 계산 ──────────────────────────────────────────
print("\n[v1 FK 직접 비교] (처음 5점)")
print(f"{'#':3} {'FK_TCP':36} {'seg':36} {'오차mm':8}")
fk_v1_pts = []
seg_pts = []
for i in range(n):
    p = arc_poses[i]
    jd = [p['S'],p['L'],p['U'],p['R'],p['B'],p['T']]
    fk = fk_v1(jd)
    seg = segs_flat[i]
    fk_v1_pts.append(fk)
    seg_pts.append(seg)
    err = math.sqrt(sum((a-b)**2 for a,b in zip(fk,seg)))
    if i < 5:
        print(f"{i:3} {str([round(x,1) for x in fk]):36} {str([round(x,1) for x in seg]):36} {err:.2f}")

R1, t1, rmse1, maxe1, res1 = kabsch(fk_v1_pts, seg_pts)
print(f"\n★ v1 Kabsch RMSE={rmse1:.3f}mm  최대={maxe1:.3f}mm")
print(f"   t = [{t1[0]:.2f}, {t1[1]:.2f}, {t1[2]:.2f}]")
print(f"   R 대각 = [{R1[0,0]:.4f}, {R1[1,1]:.4f}, {R1[2,2]:.4f}]")

# ── v4 FK 103점 계산 ──────────────────────────────────────────
print("\n[v4 FK 직접 비교] (처음 5점)")
fk_v4_pts = []
for i in range(n):
    p = arc_poses[i]
    jd = [p['S'],p['L'],p['U'],p['R'],p['B'],p['T']]
    fk = fk_v4(jd)
    fk_v4_pts.append(fk)
    seg = segs_flat[i]
    err = math.sqrt(sum((a-b)**2 for a,b in zip(fk,seg)))
    if i < 5:
        print(f"{i:3} {str([round(x,1) for x in fk]):36} {str([round(x,1) for x in seg]):36} {err:.2f}")

R4, t4, rmse4, maxe4, res4 = kabsch(fk_v4_pts, seg_pts)
print(f"\n★ v4 Kabsch RMSE={rmse4:.3f}mm  최대={maxe4:.3f}mm")
print(f"   t = [{t4[0]:.2f}, {t4[1]:.2f}, {t4[2]:.2f}]")
print(f"   R 대각 = [{R4[0,0]:.4f}, {R4[1,1]:.4f}, {R4[2,2]:.4f}]")

# ── 판정 ──────────────────────────────────────────────────────
print("\n" + "="*60)
best_rmse = min(rmse1, rmse4)
best = "v1" if rmse1 <= rmse4 else "v4"
if best_rmse < 2.0:
    print(f"✅ {best} FK가 세그 좌표와 {best_rmse:.3f}mm 일치")
    print(f"   → 변환 T = identity (세그 좌표 = 로봇 좌표계 TCP)")
    print(f"   → IK 입력에 세그 좌표 직접 사용 가능")
elif best_rmse < 10.0:
    print(f"⚠️  {best} FK RMSE={best_rmse:.1f}mm — 오프셋 있음, Kabsch T 적용 필요")
    if best == "v4":
        print(f"   t = [{t4[0]:.2f}, {t4[1]:.2f}, {t4[2]:.2f}]")
    else:
        print(f"   t = [{t1[0]:.2f}, {t1[1]:.2f}, {t1[2]:.2f}]")
else:
    print(f"❌ FK 불일치 (RMSE={best_rmse:.1f}mm) — FK 파라미터 재확인 필요")

# ── 결과 저장 ─────────────────────────────────────────────────
result = {
    'n_points': n,
    'v1': {'rmse_mm': round(rmse1,4), 'max_err_mm': round(maxe1,4),
           't': [round(x,3) for x in t1.tolist()],
           'R_diag': [round(R1[i,i],6) for i in range(3)]},
    'v4': {'rmse_mm': round(rmse4,4), 'max_err_mm': round(maxe4,4),
           't': [round(x,3) for x in t4.tolist()],
           'R_diag': [round(R4[i,i],6) for i in range(3)]},
    'verdict': f"{best} 채택 (RMSE {best_rmse:.3f}mm)",
}
out_path = r'E:\도진팩토리\3D스캔및티칭시스템\kabsch_result.json'
with open(out_path, 'w', encoding='utf-8') as f:
    json.dump(result, f, ensure_ascii=False, indent=2)
print(f"\n결과 저장: {out_path}")