# Code block 0
pairs = get_target_assembly()

import numpy as np

base, leg = pairs[0]


def rotation(q):
    w, x, y, z = np.asarray(q, dtype=float) / 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):
    # Eigenvector conversion also handles rotations near 180 degrees.
    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 np.array([w, x, y, z])


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


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


def ee_pose():
    p, q = get_ee_pose()
    return np.asarray(p, dtype=float), rotation(q)


slots = [(np.asarray(p, dtype=float), rotation(q))
         for p, q in get_assembly_slots(base)]
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])
slot_index = min(
    range(len(slots)),
    key=lambda i: np.linalg.norm(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 assembled(p, R, tp, T):
    return (
        np.linalg.norm(p - tp) < 0.006
        and np.min(np.diag(T.T @ R)) > np.cos(np.deg2rad(17))
    )


def standing(p, R, tp, T):
    return (
        np.linalg.norm(p[:2] - tp[:2]) < 0.010
        and np.dot(R[:, long_axis], T[:, long_axis])
        > np.cos(np.deg2rad(17))
        and abs(p[2] - tp[2]) < max(0.045, length * 0.8)
    )


def retreat(distance=0.075):
    p, R = ee_pose()
    goto_pose(p + np.array([0.0, 0.0, distance]), quaternion(R))


def place_upright():
    open_gripper()
    gp, gq = sample_grasp_pose(leg)
    goto_pose(np.asarray(gp), np.asarray(gq),
              z_approach=max(0.075, length), compliant=True)
    close_gripper()

    ep, E = ee_pose()
    goto_pose(ep + np.array([0.0, 0.0, max(0.13, 1.5 * length)]),
              quaternion(E))

    tp, T = target()
    lp, L = pose(leg)
    ep, E = ee_pose()
    # Preserve the measured rigid grasp while rotating the leg upright.
    relative_p = E.T @ (lp - ep)
    relative_R = E.T @ L
    desired_E = T @ relative_R.T
    elevated_leg = lp.copy()
    elevated_leg[2] = max(lp[2], tp[2] + length + 0.09)
    goto_pose(elevated_leg - desired_E @ relative_p,
              quaternion(desired_E))

    # The grasp may slip during the large rotation. Measure it again.
    lp, L = pose(leg)
    ep, E = ee_pose()
    relative_p = E.T @ (lp - ep)
    relative_R = E.T @ L
    desired_E = T @ relative_R.T
    hover_leg = tp + np.array([0.0, 0.0, max(0.065, length)])
    goto_pose(hover_leg - desired_E @ relative_p,
              quaternion(desired_E))

    # Refresh once more before descending; no leg-length offset is added
    # to the assembly pose, only a small temporary drop clearance.
    lp, L = pose(leg)
    ep, E = ee_pose()
    tp, T = target()
    relative_p = E.T @ (lp - ep)
    relative_R = E.T @ L
    desired_E = T @ relative_R.T
    drop_leg = tp + np.array([0.0, 0.0, 0.018])
    goto_pose(drop_leg - desired_E @ relative_p,
              quaternion(desired_E), compliant=True)
    open_gripper()
    retreat()


done = False
turns = 0

for placement_attempt in range(5):
    tp, T = target()
    lp, L = pose(leg)
    if assembled(lp, L, tp, T):
        done = True
        break

    if not standing(lp, L, tp, T):
        place_upright()

    tp, T = target()
    lp, L = pose(leg)
    if not standing(lp, L, tp, T):
        continue

    for _ in range(12):
        tp, T = target()
        lp, L = pose(leg)
        if assembled(lp, L, tp, T):
            done = True
            break
        if not standing(lp, L, tp, T) or turns >= 24:
            break

        # Regrasp from above with a neutral wrist for each bounded turn.
        open_gripper()
        gp, gq = sample_grasp_pose(leg)
        goto_pose(np.asarray(gp), np.asarray(gq),
                  z_approach=0.060, compliant=True)
        close_gripper()

        lp, L = pose(leg)
        ep, E = ee_pose()
        tp, T = target()
        if not standing(lp, L, tp, T):
            open_gripper()
            retreat()
            break

        # Standard right-hand thread: clockwise viewed from above.
        # Near the final depth, finish at the exact target yaw rather
        # than relying on square-cross-section symmetry.
        correction = T @ L.T
        yaw_error = np.arctan2(correction[1, 0], correction[0, 0])
        if np.linalg.norm(lp - tp) < 0.012:
            angle = yaw_error
            if angle > 0.025:
                angle -= 2 * np.pi
            angle = float(np.clip(angle, -np.pi / 2, 0.0))
            if abs(angle) < 0.025 and lp[2] - tp[2] > 0.006:
                angle = -np.pi / 2
        else:
            angle = -np.pi / 2

        # Rotate about the measured leg axis, not the gripper origin.
        # Two compliant substeps keep contact loads and wrist travel small.
        for increment in (angle / 2, angle / 2):
            lp, L = pose(leg)
            ep, E = ee_pose()
            tp, T = target()
            if not standing(lp, L, tp, T):
                break
            D = rz(increment)
            next_leg_p = lp.copy()
            next_leg_p[2] -= min(0.0025, max(0.0, lp[2] - tp[2]))
            next_ee_p = next_leg_p + D @ (ep - lp)
            goto_pose(next_ee_p, quaternion(D @ E), compliant=True)

        open_gripper()
        retreat(0.065)
        turns += 1

    if done or turns >= 24:
        break