# Code block 0
pairs = get_target_assembly()

import numpy as np

base, leg = pairs[0]


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 - w*z), 2 * (x*z + w*y)],
        [2 * (x*y + w*z), 1 - 2 * (x*x + z*z), 2 * (y*z - w*x)],
        [2 * (x*z - w*y), 2 * (y*z + w*x), 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)
    x, y, z, w = vectors[:, -1]
    q = np.array([w, x, y, z])
    return q if w >= 0 else -q


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


initial_p, _, 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])
cross_axis = (long_axis + 1) % 3

slots = get_assembly_slots(base)
slot_index = min(
    range(len(slots)),
    key=lambda i: np.linalg.norm(
        np.asarray(slots[i][0])[:2] - np.asarray(initial_p)[:2]
    ),
)


def target():
    p, q = get_assembly_slots(base)[slot_index]
    return np.asarray(p, dtype=float), rotation(q)


def matches(p, r, tp, tr):
    return (
        np.linalg.norm(p - tp) < 0.006
        and np.all(np.diag(tr.T @ r) > np.cos(np.deg2rad(17)))
    )


def standing(p, r, tp, tr):
    return (
        np.linalg.norm((p - tp)[:2]) < 0.010
        and -0.012 < p[2] - tp[2] < max(0.035, 0.5 * length)
        and np.dot(r[:, long_axis], tr[:, long_axis])
        > np.cos(np.deg2rad(15))
    )


def measured_placement(part_position, part_rotation, compliant=False):
    lp, lq = get_object_pose(leg)
    ep, eq = get_ee_pose()
    lp, ep = np.asarray(lp), np.asarray(ep)
    lr, er = rotation(lq), rotation(eq)
    relative_position = er.T @ (lp - ep)
    relative_rotation = er.T @ lr
    desired_er = part_rotation @ relative_rotation.T
    desired_ep = part_position - desired_er @ relative_position
    goto_pose(desired_ep, quaternion(desired_er), compliant=compliant)


assembled = False

for attempt in range(5):
    tp, tr = target()
    lp, lq = get_object_pose(leg)
    lp, lr = np.asarray(lp), rotation(lq)

    if matches(lp, lr, tp, tr):
        assembled = True
        break

    if not standing(lp, lr, tp, tr):
        open_gripper()
        gp, gq = sample_grasp_pose(leg)
        goto_pose(np.asarray(gp), np.asarray(gq),
                  z_approach=max(0.07, length), compliant=True)
        close_gripper()

        ep, eq = get_ee_pose()
        lifted = np.asarray(ep).copy()
        lifted[2] += length + 0.07
        goto_pose(lifted, np.asarray(eq))

        # Rotate in free space while preserving the measured leg origin.
        lp, _ = get_object_pose(leg)
        measured_placement(np.asarray(lp), tr)

        # Refresh both poses after the leg has shifted during reorientation.
        tp, tr = target()
        high = tp + np.array([0.0, 0.0, length + 0.06])
        measured_placement(high, tr)

        # This explicit waypoint is vertical in the world, not the tilted wrist.
        tp, tr = target()
        release_pose = tp + np.array([0.0, 0.0, 0.015])
        measured_placement(release_pose, tr, compliant=True)
        open_gripper()

        ep, eq = get_ee_pose()
        retreat = np.asarray(ep).copy()
        retreat[2] += length + 0.05
        goto_pose(retreat, np.asarray(eq))

        lp, lq = get_object_pose(leg)
        lp, lr = np.asarray(lp), rotation(lq)
        tp, tr = target()
        if matches(lp, lr, tp, tr):
            assembled = True
            break
        if not standing(lp, lr, tp, tr):
            continue

    # Regrasp vertically only after verifying that the released leg is seated.
    open_gripper()
    gp, reset_q = sample_grasp_pose(leg)
    reset_q = np.asarray(reset_q, dtype=float)
    goto_pose(np.asarray(gp), reset_q, z_approach=0.06, compliant=True)
    close_gripper()

    for stroke in range(20):
        tp, tr = target()
        lp, lq = get_object_pose(leg)
        lp, lr = np.asarray(lp), rotation(lq)

        if matches(lp, lr, tp, tr):
            assembled = True
            break
        if not standing(lp, lr, tp, tr):
            break

        # Clockwise threading, shortened when approaching target yaw.
        current_axis = lr[:2, cross_axis]
        target_axis = tr[:2, cross_axis]
        signed_error = np.arctan2(
            current_axis[0] * target_axis[1]
            - current_axis[1] * target_axis[0],
            np.dot(current_axis, target_axis),
        )
        clockwise_remaining = (-signed_error) % (2 * np.pi)
        if clockwise_remaining < np.deg2rad(6):
            turn = -np.pi / 2
        else:
            turn = -min(np.pi / 2, clockwise_remaining)

        # Two modest contact moves per stroke; refresh the pivot each time.
        for _ in range(2):
            lp, lq = get_object_pose(leg)
            ep, eq = get_ee_pose()
            lp, ep = np.asarray(lp), np.asarray(ep)
            delta = rz(turn / 2)
            next_ep = lp + delta @ (ep - lp)
            next_ep[2] -= 0.0015
            goto_pose(next_ep, quaternion(delta @ rotation(eq)),
                      compliant=True)

            tp, tr = target()
            lp, lq = get_object_pose(leg)
            if matches(np.asarray(lp), rotation(lq), tp, tr):
                assembled = True
                break

        if assembled:
            break

        open_gripper()
        ep, eq = get_ee_pose()
        reset_position = np.asarray(ep).copy()
        reset_position[2] += 0.035
        goto_pose(reset_position, np.asarray(eq))
        goto_pose(reset_position, reset_q)

        tp, tr = target()
        lp, lq = get_object_pose(leg)
        if not standing(np.asarray(lp), rotation(lq), tp, tr):
            break

        gp, _ = sample_grasp_pose(leg)
        goto_pose(np.asarray(gp), reset_q,
                  z_approach=0.025, compliant=True)
        close_gripper()

    open_gripper()
    ep, eq = get_ee_pose()
    retreat = np.asarray(ep).copy()
    retreat[2] += length + 0.05
    goto_pose(retreat, np.asarray(eq))

    tp, tr = target()
    lp, lq = get_object_pose(leg)
    assembled = matches(np.asarray(lp), rotation(lq), tp, tr)
    if assembled:
        break