# Code block 0
pairs = get_target_assembly()

import numpy as np


def rotation(q):
    q = np.asarray(q, dtype=float)
    w, x, y, z = q / np.linalg.norm(q)
    return np.array([
        [1 - 2 * (y*y + z*z), 2 * (x*y - z*w), 2 * (x*z + y*w)],
        [2 * (x*y + z*w), 1 - 2 * (x*x + z*z), 2 * (y*z - x*w)],
        [2 * (x*z - y*w), 2 * (y*z + x*w), 1 - 2 * (x*x + y*y)],
    ])


def quaternion(r):
    k = np.array([
        [r[0, 0] - r[1, 1] - r[2, 2], r[0, 1] + r[1, 0],
         r[0, 2] + r[2, 0], r[2, 1] - r[1, 2]],
        [r[0, 1] + r[1, 0], r[1, 1] - r[0, 0] - r[2, 2],
         r[1, 2] + r[2, 1], r[0, 2] - r[2, 0]],
        [r[0, 2] + r[2, 0], r[1, 2] + r[2, 1],
         r[2, 2] - r[0, 0] - r[1, 1], r[1, 0] - r[0, 1]],
        [r[2, 1] - r[1, 2], r[0, 2] - r[2, 0],
         r[1, 0] - r[0, 1], np.trace(r)],
    ]) / 3.0
    _, vectors = np.linalg.eigh(k)
    q = vectors[:, -1][[3, 0, 1, 2]]
    return q if q[0] >= 0 else -q


def rz(angle):
    c, s = np.cos(angle), np.sin(angle)
    return np.array([[c, -s, 0.0], [s, c, 0.0], [0.0, 0.0, 1.0]])


def read_part():
    p, q = get_object_pose(leg)
    return np.asarray(p, dtype=float), rotation(q)


def refresh_target():
    slots = get_assembly_slots(base)
    p, q = min(slots, key=lambda slot: np.linalg.norm(np.asarray(slot[0]) - target_p))
    return np.asarray(p, dtype=float), rotation(q)


def assembled(p, r):
    return (
        np.linalg.norm(p - target_p) < 0.006
        and np.min(np.diag(target_r.T @ r)) > np.cos(np.deg2rad(17))
    )


def ready_to_screw(p, r):
    return (
        np.linalg.norm((p - target_p)[:2]) < 0.009
        and np.dot(r[:, long_axis], target_r[:, long_axis]) > np.cos(np.deg2rad(15))
        and -0.012 < p[2] - target_p[2] < max(0.06, length)
    )


def ee_for_part_pose(desired_p, desired_r):
    # Always reconstruct the grasp transform from live measurements.
    p, r = read_part()
    ep, eq = get_ee_pose()
    delta = desired_r @ r.T
    return (
        desired_p + delta @ (np.asarray(ep) - p),
        quaternion(delta @ rotation(eq)),
    )


base, leg = pairs[0]
initial_p, initial_q, extent = get_object_pose(leg, return_bbox_extent=True)
extent = np.asarray(extent, dtype=float)
long_axis = int(np.argmax(extent))
length = float(extent[long_axis])
slots = get_assembly_slots(base)
chosen = min(slots, key=lambda slot: np.linalg.norm(
    np.asarray(slot[0])[:2] - np.asarray(initial_p)[:2]
))
target_p = np.asarray(chosen[0], dtype=float)
target_r = rotation(chosen[1])
up = np.array([0.0, 0.0, 1.0])
down_q = np.array([0.0, 1.0, 0.0, 0.0])
finished = False
turn_budget = 24

for attempt in range(5):
    target_p, target_r = refresh_target()
    p, r = read_part()
    if assembled(p, r):
        break

    if not ready_to_screw(p, r):
        open_gripper()
        gp, gq = sample_grasp_pose(leg)
        goto_pose(np.asarray(gp), np.asarray(gq),
                  z_approach=max(0.065, length), compliant=True)
        close_gripper()

        ep, eq = get_ee_pose()
        goto_pose(np.asarray(ep) + up * max(0.14, length + 0.06),
                  np.asarray(eq))

        # Turn the suspended leg upright, leaving plenty of table clearance.
        p, r = read_part()
        ep, eq = get_ee_pose()
        goto_pose(np.asarray(ep), quaternion(target_r @ r.T @ rotation(eq)))

        # Re-measure after reorientation: the part may have slipped in the jaws.
        target_p, target_r = refresh_target()
        above_p = target_p + up * 0.065
        place_ep, place_eq = ee_for_part_pose(above_p, target_r)
        goto_pose(place_ep, place_eq)

        # Explicit vertical staging works even with a tilted gripper.
        target_p, target_r = refresh_target()
        drop_p = target_p + up * (0.018 + 0.003 * (attempt % 2))
        place_ep, place_eq = ee_for_part_pose(drop_p, target_r)
        goto_pose(place_ep, place_eq, compliant=True)
        open_gripper()

        ep, eq = get_ee_pose()
        goto_pose(np.asarray(ep) + up * max(0.10, length), np.asarray(eq))

        p, r = read_part()
        if assembled(p, r):
            finished = True
            break
        if not ready_to_screw(p, r):
            continue

    # Ratchet with a bounded wrist excursion. Each stroke starts top-down.
    for stroke in range(min(10, turn_budget)):
        target_p, target_r = refresh_target()
        p, r = read_part()
        if assembled(p, r):
            finished = True
            break
        if not ready_to_screw(p, r):
            break

        open_gripper()
        gp, _ = sample_grasp_pose(leg)
        gp = np.asarray(gp, dtype=float)
        goto_pose(gp + up * 0.045, down_q)

        # Refresh the seated leg immediately before the contact descent.
        p, r = read_part()
        ep, eq = get_ee_pose()
        if not ready_to_screw(p, r):
            break
        gp, _ = sample_grasp_pose(leg)
        goto_pose(np.asarray(gp), down_q, z_approach=0.025, compliant=True)
        close_gripper()

        p, r = read_part()
        ep, eq = get_ee_pose()
        ep = np.asarray(ep, dtype=float)
        if not ready_to_screw(p, r):
            open_gripper()
            goto_pose(ep + up * 0.08, np.asarray(eq))
            break

        # Clockwise threading, shortening the stroke near the target yaw.
        transverse = [i for i in range(3) if i != long_axis]
        axis = max(transverse, key=lambda i: np.linalg.norm(r[:2, i]))
        current = r[:2, axis]
        desired = target_r[:2, axis]
        yaw_error = np.arctan2(
            current[0] * desired[1] - current[1] * desired[0],
            np.dot(current, desired),
        )
        clockwise_remaining = (-yaw_error) % (2 * np.pi)
        angle = -min(np.pi / 2, clockwise_remaining)
        if abs(angle) < np.deg2rad(8):
            angle = -np.pi / 2

        delta = rz(angle)
        screw_ep = p + delta @ (ep - p) - up * 0.003
        screw_eq = quaternion(delta @ rotation(eq))
        goto_pose(screw_ep, screw_eq, compliant=True)
        turn_budget -= 1

        open_gripper()
        ep, eq = get_ee_pose()
        goto_pose(np.asarray(ep) + up * 0.045, np.asarray(eq))

        p, r = read_part()
        if assembled(p, r):
            finished = True
            break
        if not ready_to_screw(p, r):
            break

    if finished or turn_budget <= 0:
        break