# Code block 0
import numpy as np

LEG = "square_table_leg4"
TOP = "square_table_top"


def unit(v):
    v = np.asarray(v, dtype=float)
    return v / np.linalg.norm(v)


def qmat(q):
    w, x, y, z = unit(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 matq(R):
    # Eigenvector conversion remains stable for half-turns.
    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
    _, V = np.linalg.eigh(K)
    x, y, z, w = V[:, -1]
    return unit([w, x, y, z])


def rotation(axis, angle):
    x, y, z = unit(axis)
    K = np.array([[0, -z, y], [z, 0, -x], [-y, x, 0]])
    return np.eye(3) + np.sin(angle) * K + (1 - np.cos(angle)) * (K @ K)


def pose(name):
    p, q = get_object_pose(name)
    return np.asarray(p, dtype=float), qmat(q)


def move_part(center, orientation):
    # Measure the grasp transform again: the part may have slipped.
    p, R = pose(LEG)
    e, eq = get_ee_pose()
    e = np.asarray(e, dtype=float)
    E = qmat(eq)
    local_offset = R.T @ (e - p)
    relative_rotation = R.T @ E
    target_e = center + orientation @ local_offset
    target_q = matq(orientation @ relative_rotation)
    if not is_reachable(target_e, target_q):
        return False
    goto_pose(target_e, target_q)
    return True


def mounted_frame(table_R, axis, long_axis):
    # Match the tabletop axes, with the leg's positive long axis pointing out.
    other = (long_axis + 1) % 3
    remaining = (long_axis + 2) % 3
    v = table_R[:, other]
    v = v - axis * np.dot(v, axis)
    if np.linalg.norm(v) < 1e-6:
        v = table_R[:, remaining]
        v = v - axis * np.dot(v, axis)
    R = np.zeros((3, 3))
    R[:, long_axis] = axis
    R[:, other] = unit(v)
    R[:, remaining] = np.cross(R[:, long_axis], R[:, other])
    return R


def main():
    p0, q0, extent = get_object_pose(LEG, return_bbox_extent=True)
    p0 = np.asarray(p0, dtype=float)
    R0 = qmat(q0)
    extent = np.asarray(extent, dtype=float)
    long_axis = int(np.argmax(extent))
    half_length = extent[long_axis] / 2
    _, table_R = pose(TOP)
    holes = get_holes(TOP)
    gp, gq = sample_grasp_pose(LEG)
    gp = np.asarray(gp, dtype=float)
    G = qmat(gq)

    candidates = []
    for hole_index, (opening, direction) in enumerate(holes):
        h, a = np.asarray(opening, dtype=float), unit(direction)
        if a[2] < 0.5:
            continue
        target_R = mounted_frame(table_R, a, long_axis)
        target_center = h + a * (half_length + 0.003)
        for mirror in (0.0, np.pi):
            grasp_R = G @ rotation([0, 0, 1], mirror)
            local_offset = R0.T @ (gp - p0)
            relative_R = R0.T @ grasp_R
            ee_target = target_center + target_R @ local_offset
            ee_R = target_R @ relative_R
            twist_R = rotation(a, -np.pi / 2) @ target_R
            twist_e = h + a * half_length + twist_R @ local_offset
            if (is_reachable(gp, matq(grasp_R))
                    and is_reachable(ee_target + a * 0.06, matq(ee_R))
                    and is_reachable(ee_target, matq(ee_R))
                    and is_reachable(twist_e, matq(twist_R @ relative_R))):
                score = np.linalg.norm(ee_target - gp)
                candidates.append((score, hole_index, grasp_R, target_R))

    if not candidates:
        raise RuntimeError("No reachable grasp and upright threading approach found.")

    _, hole_index, grasp_R, target_R = min(candidates, key=lambda c: c[0])
    open_gripper()
    goto_pose(gp, matq(grasp_R), z_approach=0.07)
    close_gripper()

    e, eq = get_ee_pose()
    goto_pose(np.asarray(e) + np.array([0.0, 0.0, 0.15]), eq)
    lifted_p, _ = pose(LEG)
    if lifted_p[2] < p0[2] + 0.04:
        raise RuntimeError("The leg did not follow the gripper during pickup.")

    h, a = get_holes(TOP)[hole_index]
    h, a = np.asarray(h, dtype=float), unit(a)
    _, table_R = pose(TOP)
    target_R = mounted_frame(table_R, a, long_axis)

    if not move_part(h + a * (half_length + 0.065), target_R):
        raise RuntimeError("Cannot reach the upright pre-insertion pose.")
    if not move_part(h + a * (half_length + 0.001), target_R):
        raise RuntimeError("Cannot align the threaded end with the hole.")

    p, R = pose(LEG)
    end = p - R[:, long_axis] * half_length
    lateral = end - h - a * np.dot(end - h, a)
    if np.linalg.norm(lateral) > 0.004 or np.dot(R[:, long_axis], a) < 0.98:
        if not move_part(h + a * (half_length + 0.001), target_R):
            raise RuntimeError("Thread alignment could not be corrected.")

    # Engage the thread before releasing the sideways pickup grasp.
    engaged_R = rotation(a, -np.pi / 2) @ target_R
    if not move_part(h + a * (half_length - 0.0015), engaged_R):
        raise RuntimeError("The initial threading turn is unreachable.")

    # Regrip near the free end with the approach along the leg.
    # Subsequent turns are wrist rotations rather than an orbiting side grasp.
    open_gripper()
    e, eq = get_ee_pose()
    E = qmat(eq)
    retreat = np.asarray(e) - E[:, 2] * 0.055
    if is_reachable(retreat, eq):
        goto_pose(retreat, eq)

    p, R = pose(LEG)
    grip_offset = max(0.0, half_length - 0.022)
    grip_center = p + a * grip_offset
    down_R = np.column_stack((target_R[:, (long_axis + 1) % 3],
                              np.zeros(3), -a))
    down_R[:, 0] = unit(down_R[:, 0])
    down_R[:, 1] = np.cross(down_R[:, 2], down_R[:, 0])

    choices = []
    _, current_q = get_ee_pose()
    for yaw in np.arange(8) * np.pi / 4:
        start_R = rotation(a, yaw) @ down_R
        finish_R = rotation(a, -np.pi / 2) @ start_R
        if (is_reachable(grip_center, matq(start_R))
                and is_reachable(grip_center - a * 0.002, matq(finish_R))
                and is_reachable(grip_center + a * 0.07, matq(start_R))):
            score = abs(np.dot(matq(start_R), unit(current_q)))
            choices.append((score, start_R))
    if not choices:
        raise RuntimeError("No reachable axial ratcheting grasp found.")

    base_R = max(choices, key=lambda c: c[0])[1]
    goto_pose(grip_center, matq(base_R), z_approach=0.07)
    close_gripper()

    total_turn = np.pi / 2
    recent_heights = []
    verified = False

    for turn in range(31):
        h, a = get_holes(TOP)[hole_index]
        h, a = np.asarray(h, dtype=float), unit(a)
        p, R = pose(LEG)

        if np.dot(R[:, long_axis], a) < 0.97:
            raise RuntimeError("The leg lost axial alignment during threading.")

        # Correct lateral slip and maintain a small axial preload while turning.
        height = np.dot(p - h, a)
        lateral = p - h - height * a
        desired_center = p - lateral - a * 0.0015
        desired_center += a * max(
            0.0, half_length - 0.003 - np.dot(desired_center - h, a)
        )
        desired_R = rotation(a, -np.pi / 2) @ R

        if abs(get_joint_positions()[6]) > 2.72:
            break
        if not move_part(desired_center, desired_R):
            break

        total_turn += np.pi / 2
        p_after, R_after = pose(LEG)
        h_now, a_now = get_holes(TOP)[hole_index]
        h_now, a_now = np.asarray(h_now), unit(a_now)
        axial_height = np.dot(p_after - h_now, a_now)
        lateral_error = np.linalg.norm(
            p_after - h_now - axial_height * a_now
        )
        recent_heights.append(axial_height)
        _, current_table_R = pose(TOP)
        axes_parallel = np.min(
            np.abs(np.diag(current_table_R.T @ R_after))
        ) > 0.985
        settled = (len(recent_heights) >= 4
                   and np.ptp(recent_heights[-4:]) < 0.0008)

        if (total_turn >= 4 * np.pi
                and abs(axial_height - half_length) < 0.003
                and lateral_error < 0.003
                and axes_parallel and settled):
            verified = True
            break

        # Release before unwinding: never reverse the engaged leg.
        open_gripper()
        e, eq = get_ee_pose()
        clearance = np.asarray(e) + a_now * 0.075
        goto_pose(clearance, eq)
        if not is_reachable(clearance, matq(base_R)):
            home_pose()
        else:
            goto_pose(clearance, matq(base_R))

        p, R = pose(LEG)
        grip_center = p + a_now * grip_offset
        if not is_reachable(grip_center, matq(base_R)):
            raise RuntimeError("The leg slipped outside the regrasp workspace.")
        goto_pose(grip_center, matq(base_R), z_approach=0.055)
        close_gripper()

    if verified:
        open_gripper()
        e, eq = get_ee_pose()
        goto_pose(np.asarray(e) + a * 0.08, eq)
        final_p, final_R = pose(LEG)
        final_h, final_a = get_holes(TOP)[hole_index]
        final_a = unit(final_a)
        final_delta = final_p - np.asarray(final_h)
        if (np.linalg.norm(final_delta - final_a * half_length) > 0.005
                or np.dot(final_R[:, long_axis], final_a) < 0.98):
            raise RuntimeError("Final verification failed after releasing the leg.")
    else:
        raise RuntimeError("Thread seating was not verified within the motion budget.")


main()