# Code block 0
import numpy as np
from scipy.spatial.transform import Rotation as R

def rot(q):
    return R.from_quat(np.asarray(q)[[1, 2, 3, 0]]).as_matrix()

def quat(m):
    return R.from_matrix(m).as_quat()[[3, 0, 1, 2]]

leg = "square_table_leg4"
top = "square_table_top"
lp, lq, le = get_object_pose(leg, return_bbox_extent=True)
tp, tq, te = get_object_pose(top, return_bbox_extent=True)
holes = get_holes(top)
gp, gq = sample_grasp_pose(leg)
print("leg", lp, lq, le)
print("top", tp, tq, te)
print("holes", holes)
print("grasp", gp, gq, "reachable", is_reachable(gp, gq))
print("ee", get_ee_pose())
print("joints", get_joint_positions())

# Code block 1
open_gripper()
goto_pose(gp, gq, z_approach=0.10)
close_gripper()
goto_pose(np.asarray(gp) + [0, 0, 0.17], gq)
lp, lq = get_object_pose(leg)
ep, eq = get_ee_pose()
print("lifted leg", lp, lq, "ee", ep, eq)
rel_R = rot(eq).T @ rot(lq)
rel_p = rot(eq).T @ (np.asarray(lp) - ep)
hole, axis = holes[2]
target_R = rot(tq)
target_c = hole + axis * (le[1]/2 + 0.10)
ee_R = target_R @ rel_R.T
ee_p = target_c - ee_R @ rel_p
print("upright target", ee_p, quat(ee_R), is_reachable(ee_p, quat(ee_R)))
if is_reachable(ee_p, quat(ee_R)):
    goto_pose(ee_p, quat(ee_R))
print("leg after orient", get_object_pose(leg))
print("ee after orient", get_ee_pose())

# Code block 2
lp, lq, le = get_object_pose(leg, return_bbox_extent=True)
tp, tq = get_object_pose(top)
holes = get_holes(top)
hole, axis = [np.asarray(v) for v in holes[2]]
axis = axis / np.linalg.norm(axis)
print("top", tp, tq, "hole", hole, "axis", axis)
print("leg", lp, lq, le, "long axis", rot(lq)[:, 1])
ep, eq = get_ee_pose()
rel_R = rot(eq).T @ rot(lq)
rel_p = rot(eq).T @ (np.asarray(lp) - np.asarray(ep))
target_R = rot(tq)
target_c = hole + axis * (le[1] / 2)
ee_R = target_R @ rel_R.T
ee_p = target_c - ee_R @ rel_p
print("contact target", ee_p, quat(ee_R), is_reachable(ee_p, quat(ee_R)))
if is_reachable(ee_p, quat(ee_R)):
    goto_pose(ee_p, quat(ee_R))
print("contact leg", get_object_pose(leg))
print("contact ee", get_ee_pose())
print("joints", get_joint_positions())

# Code block 3
import numpy as np
from scipy.spatial.transform import Rotation as R

lp, lq, le = get_object_pose(leg, return_bbox_extent=True)
ep, eq = get_ee_pose()
print("Before regrasp", lp, lq)

# Switch to a downward grip around the upper part of the upright leg.
open_gripper()
goto_pose(np.asarray(ep) + [0, 0, 0.14], eq)
lp, lq = get_object_pose(leg)
print("Released leg", lp, lq)

long_axis = rot(lq)[:, 1]
if abs(long_axis[2]) > 0.9:
    grasp_p = np.asarray(lp) + [0, 0, 0.022]
    candidates = []
    for angle in np.linspace(-np.pi, np.pi, 9)[:-1]:
        gr = R.from_rotvec([0, 0, angle]).as_matrix() @ rot([0, 1, 0, 0])
        gquat = quat(gr)
        if is_reachable(grasp_p, gquat):
            candidates.append(gquat)
    if candidates:
        gquat = min(candidates, key=lambda q: abs(np.dot(q, eq)))
        goto_pose(grasp_p, gquat, z_approach=0.09)
        close_gripper()
        lp, lq = get_object_pose(leg)
        ep, eq = get_ee_pose()
        rel_R = rot(eq).T @ rot(lq)
        rel_p = rot(eq).T @ (np.asarray(lp) - np.asarray(ep))
        hole, axis = [np.asarray(v) for v in get_holes(top)[2]]
        _, tq = get_object_pose(top)
        target_R = rot(tq)
        target_c = hole + axis * (le[1] / 2)
        ee_R = target_R @ rel_R.T
        ee_p = target_c - ee_R @ rel_p
        if is_reachable(ee_p, quat(ee_R)):
            goto_pose(ee_p, quat(ee_R))
        print("Downward contact leg", get_object_pose(leg))
        print("Downward contact ee", get_ee_pose())
        print("Joints", get_joint_positions())
else:
    print("Leg is no longer upright; needs another pickup.")

# Code block 4
import numpy as np
from scipy.spatial.transform import Rotation as R

def rot(q):
    return R.from_quat(np.asarray(q)[[1, 2, 3, 0]]).as_matrix()

def quat(m):
    return R.from_matrix(m).as_quat()[[3, 0, 1, 2]]

leg, top = "square_table_leg4", "square_table_top"
lp, lq, le = get_object_pose(leg, return_bbox_extent=True)
tp, tq, te = get_object_pose(top, return_bbox_extent=True)
print("Leg:", lp, lq, le)
print("Table:", tp, tq, te)
print("Holes:", get_holes(top))

gp, gq = sample_grasp_pose(leg)
if is_reachable(gp, gq):
    open_gripper()
    goto_pose(gp, gq, z_approach=0.10)
    close_gripper()
    goto_pose(np.asarray(gp) + [0, 0, 0.16], gq)

    lp, lq = get_object_pose(leg)
    ep, eq = get_ee_pose()
    offset = rot(eq).T @ (np.asarray(lp) - np.asarray(ep))
    relative = rot(eq).T @ rot(lq)
    target_R = rot(tq)
    options = []
    for i, (h, a) in enumerate(get_holes(top)):
        h, a = np.asarray(h), np.asarray(a)
        for angle in [0.0, np.pi / 2, -np.pi / 2, np.pi]:
            lr = R.from_rotvec(a * angle).as_matrix() @ target_R
            er = lr @ relative.T
            center = h + a * (le[1] / 2)
            contact = center - er @ offset
            above = contact + a * 0.09
            if is_reachable(above, quat(er)) and is_reachable(contact, quat(er)):
                options.append((np.linalg.norm(above - ep), i, above, contact, er))
    if options:
        _, hole_index, above, contact, er = min(options, key=lambda x: x[0])
        goto_pose(above, quat(er))
        lp, lq = get_object_pose(leg)
        ep, eq = get_ee_pose()
        h, a = [np.asarray(v) for v in get_holes(top)[hole_index]]
        print("Selected hole:", hole_index, h, a)
        print("Upright leg:", lp, lq, "long axis:", rot(lq)[:, 1])
        if abs(np.dot(rot(lq)[:, 1], a)) > 0.96:
            correction = h + a * (le[1] / 2) - np.asarray(lp)
            contact = np.asarray(ep) + correction
            if is_reachable(contact, eq):
                goto_pose(contact, eq)
                print("At contact:", get_object_pose(leg))
                for k in range(3):
                    lp, lq = get_object_pose(leg)
                    ep, eq = get_ee_pose()
                    h, a = [np.asarray(v) for v in get_holes(top)[hole_index]]
                    if abs(np.dot(rot(lq)[:, 1], a)) < 0.94:
                        print("Stopped: leg tilted.")
                        break
                    turn = R.from_rotvec(-a * np.pi / 6).as_matrix()
                    desired_center = h + a * (le[1] / 2 - 0.0015)
                    new_ep = desired_center + turn @ (np.asarray(ep) - np.asarray(lp))
                    new_eq = quat(turn @ rot(eq))
                    if not is_reachable(new_ep, new_eq):
                        print("Stopped: turning pose unreachable.")
                        break
                    goto_pose(new_ep, new_eq)
                    print("Turn", k + 1, "leg:", get_object_pose(leg))
        else:
            print("Stopped: orientation slipped during pickup.")
    else:
        print("No reachable upright contact with this grasp.")
else:
    print("Pickup pose unreachable.")

print("Final EE:", get_ee_pose())
print("Final joints:", get_joint_positions())
print("Final leg:", get_object_pose(leg))

# Code block 5
import numpy as np
from scipy.spatial.transform import Rotation as R

def rot(q):
    return R.from_quat(np.asarray(q)[[1, 2, 3, 0]]).as_matrix()

def quat(m):
    return R.from_matrix(m).as_quat()[[3, 0, 1, 2]]

leg, top = "square_table_leg4", "square_table_top"
open_gripper()
home_pose()
lp, lq, le = get_object_pose(leg, return_bbox_extent=True)
gp, gq = sample_grasp_pose(leg)
print("Before pickup:", lp, lq, "grasp:", gp, gq)
if is_reachable(gp, gq):
    goto_pose(gp, gq, z_approach=0.10)
    close_gripper()
    print("Closed:", get_object_pose(leg), "EE:", get_ee_pose())
    ep, eq = get_ee_pose()
    goto_pose(np.asarray(ep) + [0, 0, 0.12], eq)
    lp, lq = get_object_pose(leg)
    ep, eq = get_ee_pose()
    print("Lift:", lp, lq, "EE:", ep, eq)
    if lp[2] > 0.08 and np.linalg.norm(np.asarray(lp) - ep) < 0.09:
        relative = rot(eq).T @ rot(lq)
        offset = rot(eq).T @ (np.asarray(lp) - ep)
        _, tq = get_object_pose(top)
        candidates = []
        for i, (h, a) in enumerate(get_holes(top)):
            h, a = np.asarray(h), np.asarray(a)
            for angle in [0, np.pi / 2, -np.pi / 2, np.pi]:
                lr = R.from_rotvec(a * angle).as_matrix() @ rot(tq)
                er = lr @ relative.T
                center = h + a * (le[1] / 2 + 0.15)
                target = center - er @ offset
                if is_reachable(target, quat(er)):
                    delta = R.from_matrix(er @ rot(eq).T).magnitude()
                    candidates.append((delta, i, target, er))
        if candidates:
            _, hole_index, target, er = min(candidates, key=lambda x: x[0])
            start_p, start_r = np.asarray(ep), rot(eq)
            rv = R.from_matrix(er @ start_r.T).as_rotvec()
            for f in [0.25, 0.5, 0.75, 1.0]:
                step_r = R.from_rotvec(f * rv).as_matrix() @ start_r
                goto_pose(start_p + f * (target - start_p), quat(step_r))
                p, q = get_object_pose(leg)
                e, _ = get_ee_pose()
                print("Orient", f, "leg:", p, q)
                if np.linalg.norm(np.asarray(p) - e) > 0.09:
                    print("Lost grasp; stopped.")
                    break
            print("Selected hole:", hole_index)
    else:
        print("Pickup failed; no orientation attempted.")
print("Current holes:", get_holes(top))
print("Current EE:", get_ee_pose())
print("Current leg:", get_object_pose(leg))

# Code block 6
import numpy as np
from scipy.spatial.transform import Rotation as R

def rot(q):
    return R.from_quat(np.asarray(q)[[1, 2, 3, 0]]).as_matrix()

def quat(m):
    return R.from_matrix(m).as_quat()[[3, 0, 1, 2]]

leg, top = "square_table_leg4", "square_table_top"
hole_index = 0
_, tq = get_object_pose(top)
target_R = rot(tq)
lp, lq, le = get_object_pose(leg, return_bbox_extent=True)
h, a = [np.asarray(v) for v in get_holes(top)[hole_index]]
a /= np.linalg.norm(a)

# Correct the measured grasp transform, rather than the original grasp.
for k in range(4):
    lp, lq = get_object_pose(leg)
    ep, eq = get_ee_pose()
    if np.linalg.norm(np.asarray(lp) - ep) > 0.07:
        print("Grasp lost.")
        break
    candidates = [
        R.from_rotvec(a * theta).as_matrix() @ target_R
        for theta in [0, np.pi / 2, np.pi, -np.pi / 2]
    ]
    desired = min(candidates, key=lambda m: R.from_matrix(m @ rot(lq).T).magnitude())
    correction = desired @ rot(lq).T
    center = h + a * (le[1] / 2 + 0.11)
    new_p = center + correction @ (np.asarray(ep) - lp)
    new_q = quat(correction @ rot(eq))
    if not is_reachable(new_p, new_q):
        print("Correction unreachable.")
        break
    goto_pose(new_p, new_q)
    lp, lq = get_object_pose(leg)
    error = R.from_matrix(desired @ rot(lq).T).magnitude()
    print("Orientation correction", k, "error", error, "center", lp)
    if error < 0.045:
        break

lp, lq = get_object_pose(leg)
if np.dot(rot(lq)[:, 1], a) > 0.985:
    for clearance in [0.045, 0.012, 0.0]:
        lp, lq = get_object_pose(leg)
        ep, eq = get_ee_pose()
        h, a = [np.asarray(v) for v in get_holes(top)[hole_index]]
        desired_center = h + a * (le[1] / 2 + clearance)
        new_p = np.asarray(ep) + desired_center - lp
        if not is_reachable(new_p, eq):
            print("Descent unreachable.")
            break
        goto_pose(new_p, eq)
        print("Descent", clearance, get_object_pose(leg))

    for k in range(4):
        lp, lq = get_object_pose(leg)
        ep, eq = get_ee_pose()
        h, a = [np.asarray(v) for v in get_holes(top)[hole_index]]
        delta = np.asarray(lp) - h
        lateral = np.linalg.norm(delta - a * np.dot(delta, a))
        if np.dot(rot(lq)[:, 1], a) < 0.97 or lateral > 0.012:
            print("Stopped: alignment requires correction.", lateral)
            break
        turn = R.from_rotvec(-a * np.pi / 6).as_matrix()
        center = h + a * min(np.dot(delta, a) - 0.001, le[1] / 2 - 0.001)
        new_p = center + turn @ (np.asarray(ep) - lp)
        new_q = quat(turn @ rot(eq))
        if not is_reachable(new_p, new_q):
            print("Wrist reset needed.")
            break
        goto_pose(new_p, new_q)
        print("Screw turn", k, get_object_pose(leg))
else:
    print("Not upright enough for contact.")

print("Leg:", get_object_pose(leg))
print("EE:", get_ee_pose())
print("Joints:", get_joint_positions())
print("Table:", get_object_pose(top))

# Code block 7
import numpy as np
from scipy.spatial.transform import Rotation as R

leg, top = "square_table_leg4", "square_table_top"

def rot(q):
    return R.from_quat(np.asarray(q)[[1, 2, 3, 0]]).as_matrix()

def quat(m):
    return R.from_matrix(m).as_quat()[[3, 0, 1, 2]]

_, _, extent = get_object_pose(leg, return_bbox_extent=True)
h, a = [np.asarray(v) for v in get_holes(top)[0]]
a /= np.linalg.norm(a)

# Replace the near-limit sideways grip with a downward grip.
open_gripper()
ep, eq = get_ee_pose()
goto_pose(np.asarray(ep) + [0, 0, 0.12], eq)
lp, lq = get_object_pose(leg)
print("After release:", lp, lq)

if np.dot(rot(lq)[:, 1], a) > 0.96:
    for cycle in range(5):
        lp, lq = get_object_pose(leg)
        h, a = [np.asarray(v) for v in get_holes(top)[0]]
        a /= np.linalg.norm(a)
        delta = np.asarray(lp) - h
        if (np.dot(rot(lq)[:, 1], a) < 0.96 or
                np.linalg.norm(delta - a * np.dot(delta, a)) > 0.012):
            print("Stopped: leg is no longer aligned.", lp, lq)
            break

        grip_p = np.asarray(lp) + a * 0.022
        yaw = np.arctan2(rot(lq)[1, 0], rot(lq)[0, 0])
        options = []
        for offset in [0, np.pi / 2, np.pi, -np.pi / 2]:
            gr = R.from_rotvec(a * (yaw + offset)).as_matrix() @ rot([0, 1, 0, 0])
            if all(is_reachable(grip_p, quat(
                    R.from_rotvec(-a * angle).as_matrix() @ gr))
                   for angle in [0, np.pi / 6, np.pi / 3, np.pi / 2]):
                options.append(gr)
        if not options:
            print("No reachable downward screw stroke.")
            break

        _, eq = get_ee_pose()
        gr = max(options, key=lambda m: abs(np.dot(quat(m), eq)))
        goto_pose(grip_p, quat(gr), z_approach=0.075)
        close_gripper()

        for step in range(3):
            lp, lq = get_object_pose(leg)
            ep, eq = get_ee_pose()
            h, a = [np.asarray(v) for v in get_holes(top)[0]]
            a /= np.linalg.norm(a)
            delta = np.asarray(lp) - h
            if np.dot(rot(lq)[:, 1], a) < 0.96:
                print("Stopped stroke: tilted.")
                break
            turn = R.from_rotvec(-a * np.pi / 6).as_matrix()
            # Apply modest pressure while rotating, without assuming insertion depth.
            center = h + a * (min(np.dot(delta, a), extent[1] / 2) - 0.001)
            target_p = center + turn @ (np.asarray(ep) - lp)
            target_q = quat(turn @ rot(eq))
            if not is_reachable(target_p, target_q):
                print("Reset required.")
                break
            goto_pose(target_p, target_q)
            print("Stroke", cycle, step, get_object_pose(leg))

        open_gripper()
        ep, eq = get_ee_pose()
        goto_pose(np.asarray(ep) + a * 0.09, eq)
        print("Released verification:", get_object_pose(leg))
else:
    print("Leg requires recovery before further turning.")

print("Final leg:", get_object_pose(leg))
print("Final tabletop:", get_object_pose(top))
print("Final holes:", get_holes(top))

# Code block 8
import numpy as np
from scipy.spatial.transform import Rotation as R

def rot(q):
    return R.from_quat(np.asarray(q)[[1, 2, 3, 0]]).as_matrix()

def quat(m):
    return R.from_matrix(m).as_quat()[[3, 0, 1, 2]]

leg, top = "square_table_leg4", "square_table_top"
lp, lq, extent = get_object_pose(leg, return_bbox_extent=True)
h, a = [np.asarray(v) for v in get_holes(top)[0]]
a /= np.linalg.norm(a)
_, tq = get_object_pose(top)
base = rot(tq)

# Grip the upper portion directly, retaining a downward wrist for screwing.
grip = np.asarray(lp) + rot(lq)[:, 1] * 0.018
yaw = np.arctan2(rot(lq)[1, 0], rot(lq)[0, 0])
choices = []
for offset in [0, np.pi / 2, -np.pi / 2, np.pi]:
    gr = R.from_rotvec(a * (yaw + offset)).as_matrix() @ rot([0, 1, 0, 0])
    if is_reachable(grip, quat(gr)):
        choices.append(gr)

if choices:
    _, eq = get_ee_pose()
    gr = min(choices, key=lambda m: R.from_matrix(m @ rot(eq).T).magnitude())
    open_gripper()
    goto_pose(grip, quat(gr), z_approach=0.08)
    close_gripper()
    ep, eq = get_ee_pose()
    goto_pose(np.asarray(ep) + a * 0.08, eq)

    aligned = False
    for k in range(4):
        lp, lq = get_object_pose(leg)
        ep, eq = get_ee_pose()
        if np.linalg.norm(np.asarray(lp) - ep) > 0.065:
            print("Pickup slipped.", lp)
            break
        targets = [R.from_rotvec(a * t).as_matrix() @ base
                   for t in [0, np.pi / 2, np.pi, -np.pi / 2]]
        desired = min(targets, key=lambda m: R.from_matrix(m @ rot(lq).T).magnitude())
        correction = desired @ rot(lq).T
        center = h + a * (extent[1] / 2 + 0.065)
        p = center + correction @ (np.asarray(ep) - lp)
        q = quat(correction @ rot(eq))
        if not is_reachable(p, q):
            print("Alignment unreachable.")
            break
        goto_pose(p, q)
        lp, lq = get_object_pose(leg)
        print("Alignment", k, lp, lq)
        if np.dot(rot(lq)[:, 1], a) > 0.997:
            aligned = True
            break

    if aligned:
        for clearance in [0.015, 0.0]:
            lp, lq = get_object_pose(leg)
            ep, eq = get_ee_pose()
            h, a = [np.asarray(v) for v in get_holes(top)[0]]
            p = np.asarray(ep) + h + a * (extent[1] / 2 + clearance) - lp
            if not is_reachable(p, eq):
                aligned = False
                break
            goto_pose(p, eq)

    if aligned:
        for step in range(8):
            lp, lq = get_object_pose(leg)
            ep, eq = get_ee_pose()
            h, a = [np.asarray(v) for v in get_holes(top)[0]]
            a /= np.linalg.norm(a)
            delta = np.asarray(lp) - h
            lateral = np.linalg.norm(delta - a * np.dot(delta, a))
            if np.dot(rot(lq)[:, 1], a) < 0.97 or lateral > 0.01:
                print("Alignment lost; retaining grip.", lateral)
                break
            turn = R.from_rotvec(-a * np.pi / 6).as_matrix()
            center = h + a * (min(np.dot(delta, a), extent[1] / 2) - 0.0008)
            p = center + turn @ (np.asarray(ep) - lp)
            q = quat(turn @ rot(eq))
            if not is_reachable(p, q):
                print("Wrist reset needed; retaining grip.")
                break
            goto_pose(p, q)
            print("Threading", step, get_object_pose(leg))

print("Leg:", get_object_pose(leg))
print("Table:", get_object_pose(top))
print("EE:", get_ee_pose())
print("Joints:", get_joint_positions())

# Code block 9
import numpy as np
from scipy.spatial.transform import Rotation as R

def rot(q):
    return R.from_quat(np.asarray(q)[[1, 2, 3, 0]]).as_matrix()

def quat(m):
    return R.from_matrix(m).as_quat()[[3, 0, 1, 2]]

leg, top = "square_table_leg4", "square_table_top"
lp, lq, extent = get_object_pose(leg, return_bbox_extent=True)
holes = get_holes(top)
hole_index = min(range(len(holes)), key=lambda i: np.linalg.norm(
    np.asarray(lp) - (np.asarray(holes[i][0]) + np.asarray(holes[i][1]) * extent[1] / 2)))
h, a = [np.asarray(v) for v in holes[hole_index]]
a = a / np.linalg.norm(a)
print("Nearest hole:", hole_index, h, a)
print("Initial leg axis:", rot(lq)[:, 1])

open_gripper()
# Grip at the centre rather than above the free end.
grip_p = np.asarray(lp).copy()
grip_p[2] = max(grip_p[2], 0.025)
yaw = np.arctan2(rot(lq)[1, 0], rot(lq)[0, 0])
_, eq = get_ee_pose()
options = []
for offset in [0, np.pi / 2, -np.pi / 2, np.pi]:
    gr = R.from_rotvec([0, 0, yaw + offset]).as_matrix() @ rot([0, 1, 0, 0])
    if is_reachable(grip_p, quat(gr)):
        options.append(gr)
if options:
    gr = min(options, key=lambda m: R.from_matrix(m @ rot(eq).T).magnitude())
    goto_pose(grip_p, quat(gr), z_approach=0.06)
    close_gripper()
    before, _ = get_object_pose(leg)
    ep, eq = get_ee_pose()
    goto_pose(np.asarray(ep) + a * 0.035, eq)
    lp, lq = get_object_pose(leg)
    print("Pickup displacement:", np.asarray(lp) - before)
    if np.dot(np.asarray(lp) - before, a) > 0.022:
        for k in range(3):
            lp, lq = get_object_pose(leg)
            ep, eq = get_ee_pose()
            _, tq = get_object_pose(top)
            targets = [R.from_rotvec(a * t).as_matrix() @ rot(tq)
                       for t in [0, np.pi / 2, np.pi, -np.pi / 2]]
            desired = min(targets, key=lambda m: R.from_matrix(m @ rot(lq).T).magnitude())
            correction = desired @ rot(lq).T
            center = h + a * (extent[1] / 2 + 0.035)
            p = center + correction @ (np.asarray(ep) - lp)
            q = quat(correction @ rot(eq))
            if not is_reachable(p, q):
                break
            goto_pose(p, q)
            lp, lq = get_object_pose(leg)
            if R.from_matrix(desired @ rot(lq).T).magnitude() < 0.04:
                break

        lp, lq = get_object_pose(leg)
        if np.dot(rot(lq)[:, 1], a) > 0.995:
            ep, eq = get_ee_pose()
            p = np.asarray(ep) + h + a * (extent[1] / 2) - lp
            if is_reachable(p, eq):
                goto_pose(p, eq)
                for step in range(6):
                    lp, lq = get_object_pose(leg)
                    ep, eq = get_ee_pose()
                    h, a = [np.asarray(v) for v in get_holes(top)[hole_index]]
                    a /= np.linalg.norm(a)
                    d = np.asarray(lp) - h
                    if (np.dot(rot(lq)[:, 1], a) < 0.975 or
                            np.linalg.norm(d - a * np.dot(d, a)) > 0.008):
                        print("Alignment lost.")
                        break
                    turn = R.from_rotvec(-a * np.pi / 6).as_matrix()
                    center = h + a * (min(np.dot(d, a), extent[1] / 2) - 0.0008)
                    p = center + turn @ (np.asarray(ep) - lp)
                    q = quat(turn @ rot(eq))
                    if not is_reachable(p, q):
                        print("Wrist reset required.")
                        break
                    goto_pose(p, q)
                    print("Turn:", step, get_object_pose(leg))
    else:
        print("Centre grasp failed; stopped.")
print("Leg:", get_object_pose(leg))
print("Table:", get_object_pose(top))
print("EE:", get_ee_pose())
print("Joints:", get_joint_positions())

# Code block 10
import numpy as np
from scipy.spatial.transform import Rotation as R

def rot(q):
    return R.from_quat(np.asarray(q)[[1, 2, 3, 0]]).as_matrix()

def quat(m):
    return R.from_matrix(m).as_quat()[[3, 0, 1, 2]]

leg, top = "square_table_leg4", "square_table_top"
lp, lq, extent = get_object_pose(leg, return_bbox_extent=True)
holes = get_holes(top)
idx = min(range(len(holes)), key=lambda i: np.linalg.norm(
    lp - (np.asarray(holes[i][0]) + np.asarray(holes[i][1]) * extent[1] / 2)))

for cycle in range(4):
    ep, eq = get_ee_pose()
    reset_R = rot(eq)
    executed = 0
    for step in range(6):
        lp, lq = get_object_pose(leg)
        ep, eq = get_ee_pose()
        h, a = map(np.asarray, get_holes(top)[idx])
        a = a / np.linalg.norm(a)
        d = np.asarray(lp) - h
        lateral = np.linalg.norm(d - a * np.dot(d, a))
        alignment = np.dot(rot(lq)[:, 1], a)
        if lateral > 0.018 or alignment < 0.92:
            print("Recovery needed:", lateral, alignment, lp, lq)
            break

        # Correct tilt without cancelling the accumulated screw rotation.
        long_axis = rot(lq)[:, 1]
        cross = np.cross(long_axis, a)
        angle = np.arctan2(np.linalg.norm(cross), np.dot(long_axis, a))
        tilt = (R.from_rotvec(cross / np.linalg.norm(cross) * angle).as_matrix()
                if np.linalg.norm(cross) > 1e-8 else np.eye(3))
        turn = R.from_rotvec(-a * np.pi / 6).as_matrix()
        correction = turn @ tilt
        center = h + a * (min(np.dot(d, a), extent[1] / 2) - 0.0007)
        target_p = center + correction @ (np.asarray(ep) - lp)
        target_q = quat(correction @ rot(eq))
        if not is_reachable(target_p, target_q):
            print("Resetting wrist.")
            break
        goto_pose(target_p, target_q)
        executed += 1
        p, q = get_object_pose(leg)
        print("Turn", cycle, step, "center", p, "orientation", q)

    if executed == 0:
        break

    # Leave the leg in place while unwinding the empty wrist.
    open_gripper()
    ep, eq = get_ee_pose()
    goto_pose(np.asarray(ep) + a * 0.075, eq)
    lp, lq = get_object_pose(leg)
    h, a = map(np.asarray, get_holes(top)[idx])
    a = a / np.linalg.norm(a)
    print("Released:", lp, lq)
    if np.dot(rot(lq)[:, 1], a) < 0.95:
        print("Leg tilted after release; stopped.")
        break
    if cycle == 3:
        break
    grip_p = np.asarray(lp) + a * 0.008
    if not is_reachable(grip_p, quat(reset_R)):
        print("Reset grasp unreachable.")
        break
    goto_pose(grip_p + a * 0.075, quat(reset_R))
    goto_pose(grip_p, quat(reset_R))
    close_gripper()

print("Verification leg:", get_object_pose(leg))
print("Verification tabletop:", get_object_pose(top))
print("Verification holes:", get_holes(top))

# Code block 11
import numpy as np
from scipy.spatial.transform import Rotation as R

def rot(q):
    return R.from_quat(np.asarray(q)[[1, 2, 3, 0]]).as_matrix()

def quat(m):
    return R.from_matrix(m).as_quat()[[3, 0, 1, 2]]

leg, top = "square_table_leg4", "square_table_top"
lp, lq = get_object_pose(leg)
_, tq = get_object_pose(top)
h, a = map(np.asarray, get_holes(top)[0])
a = a / np.linalg.norm(a)
lr, tr = rot(lq), rot(tq)
angle = np.arctan2(np.dot(a, np.cross(lr[:, 0], tr[:, 0])),
                   np.dot(lr[:, 0], tr[:, 0]))
angle = (angle + np.pi / 4) % (np.pi / 2) - np.pi / 4
print("Remaining axis alignment angle:", angle)

if abs(angle) > 0.025:
    grip = np.asarray(lp) + a * 0.008
    yaw = np.arctan2(lr[1, 0], lr[0, 0])
    _, eq = get_ee_pose()
    choices = []
    for offset in [0, np.pi / 2, -np.pi / 2, np.pi]:
        gr = R.from_rotvec([0, 0, yaw + offset]).as_matrix() @ rot([0, 1, 0, 0])
        if is_reachable(grip, quat(gr)):
            choices.append(gr)
    if choices:
        gr = min(choices, key=lambda m: R.from_matrix(m @ rot(eq).T).magnitude())
        goto_pose(grip, quat(gr), z_approach=0.06)
        close_gripper()
        lp, lq = get_object_pose(leg)
        ep, eq = get_ee_pose()
        _, tq = get_object_pose(top)
        lr, tr = rot(lq), rot(tq)
        targets = [R.from_rotvec(a * t).as_matrix() @ tr
                   for t in [0, np.pi / 2, -np.pi / 2, np.pi]]
        desired = min(targets, key=lambda m: R.from_matrix(m @ lr.T).magnitude())
        correction = desired @ lr.T
        h, a = map(np.asarray, get_holes(top)[0])
        a = a / np.linalg.norm(a)
        center = h + a * (np.dot(np.asarray(lp) - h, a) - 0.0004)
        p = center + correction @ (np.asarray(ep) - lp)
        q = quat(correction @ rot(eq))
        if is_reachable(p, q):
            goto_pose(p, q)
        open_gripper()
        ep, eq = get_ee_pose()
        goto_pose(np.asarray(ep) + a * 0.08, eq)

print("Final leg:", get_object_pose(leg))
print("Final tabletop:", get_object_pose(top))
print("Final holes:", get_holes(top))