# 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"
print("leg", get_object_pose(leg, return_bbox_extent=True))
print("top", get_object_pose(top, return_bbox_extent=True))
print("holes", get_holes(top))
print("grasp", sample_grasp_pose(leg))
print("ee", get_ee_pose())
print("joints", get_joint_positions())

# Code block 1
gp, gq = sample_grasp_pose(leg)
open_gripper()
goto_pose(gp, gq, z_approach=0.07)
close_gripper()
goto_pose(np.asarray(gp) + [0, 0, 0.15], gq)
lp, lq, ext = get_object_pose(leg, return_bbox_extent=True)
ep, eq = get_ee_pose()
print("lifted", lp, lq)
relative_R = rot(lq).T @ rot(eq)
relative_p = rot(lq).T @ (np.asarray(ep) - lp)
tp, tq = get_object_pose(top)
target_R = rot(tq)
hole, axis = get_holes(top)[2]
target_center = hole + axis * (ext[1] / 2 + 0.09)
target_ep = target_center + target_R @ relative_p
target_eq = quat(target_R @ relative_R)
print("upright target", target_ep, target_eq, is_reachable(target_ep, target_eq))
if is_reachable(target_ep, target_eq):
    goto_pose(target_ep, target_eq)
print("leg", get_object_pose(leg))
print("ee", get_ee_pose())

# Code block 2
lp, lq, ext = get_object_pose(leg, return_bbox_extent=True)
tp, tq = get_object_pose(top)
holes = get_holes(top)
hole, axis = holes[2]
hole, axis = np.asarray(hole), np.asarray(axis)
print("top", tp, tq, "hole", hole, axis, "leg size", ext)
ep, eq = get_ee_pose()
relative_R = rot(lq).T @ rot(eq)
relative_p = rot(lq).T @ (np.asarray(ep) - lp)
target_R = rot(tq)
target_center = hole + axis * (ext[1] / 2 + 0.025)
target_ep = target_center + target_R @ relative_p
target_eq = quat(target_R @ relative_R)
if is_reachable(target_ep, target_eq):
    goto_pose(target_ep, target_eq)
    lp, lq = get_object_pose(leg)
    ep, eq = get_ee_pose()
    relative_R = rot(lq).T @ rot(eq)
    relative_p = rot(lq).T @ (np.asarray(ep) - lp)
    target_center = hole + axis * (ext[1] / 2 - 0.001)
    target_ep = target_center + target_R @ relative_p
    target_eq = quat(target_R @ relative_R)
    if is_reachable(target_ep, target_eq):
        goto_pose(target_ep, target_eq)
print("contact leg", get_object_pose(leg))
print("ee", get_ee_pose())
print("joints", get_joint_positions())

# Code block 3
hole, axis = map(np.asarray, get_holes(top)[2])
lp, lq, ext = get_object_pose(leg, return_bbox_extent=True)
grasp_p = np.asarray(lp) + axis * 0.012
down_R = rot([0, 1, 0, 0])

def axial_rotation(angle):
    return R.from_rotvec(axis * angle).as_matrix()

choices = []
for yaw in np.linspace(-np.pi, np.pi, 12, endpoint=False):
    start_R = axial_rotation(yaw) @ down_R
    if all(is_reachable(grasp_p, quat(axial_rotation(-a) @ start_R))
           for a in [0, np.pi/4, np.pi/2, 3*np.pi/4, np.pi]):
        choices.append(start_R)

if not choices:
    print("No reachable downward screw grasp", get_object_pose(leg))
else:
    start_R = choices[0]
    open_gripper()
    ep, eq = get_ee_pose()
    goto_pose(np.asarray(ep) + axis * 0.09, eq)
    goto_pose(grasp_p, quat(start_R), z_approach=0.07)
    close_gripper()

    for cycle in range(2):
        for step in range(4):
            lp, lq = get_object_pose(leg)
            ep, eq = get_ee_pose()
            lp, ep = np.asarray(lp), np.asarray(ep)
            turn = axial_rotation(-np.pi/4)
            new_leg_R = turn @ rot(lq)
            offset = rot(lq).T @ (ep - lp)
            grip_R = rot(lq).T @ rot(eq)
            height = np.dot(lp - hole, axis)
            target_height = min(height - 0.0015, ext[1]/2 - 0.001)
            target_center = hole + axis * target_height
            target_p = target_center + new_leg_R @ offset
            target_q = quat(new_leg_R @ grip_R)
            if not is_reachable(target_p, target_q):
                print("Twist unreachable", cycle, step, get_joint_positions())
                break
            goto_pose(target_p, target_q)
            actual_p, actual_q = get_object_pose(leg)
            print("screw", cycle, step, actual_p, actual_q)
        if cycle == 0:
            lp, lq = get_object_pose(leg)
            grasp_p = np.asarray(lp) + axis * 0.012
            open_gripper()
            ep, eq = get_ee_pose()
            goto_pose(np.asarray(ep) + axis * 0.065, eq)
            goto_pose(grasp_p + axis * 0.065, quat(start_R))
            goto_pose(grasp_p, quat(start_R))
            close_gripper()

    print("final leg", get_object_pose(leg))
    print("top", get_object_pose(top))
    print("ee", get_ee_pose())
    print("joints", get_joint_positions())

# Code block 4
import numpy as np
print("leg", get_object_pose("square_table_leg4", return_bbox_extent=True))
print("top", get_object_pose("square_table_top", return_bbox_extent=True))
holes = get_holes("square_table_top")
print("holes", holes)
lp, lq, ext = get_object_pose("square_table_leg4", return_bbox_extent=True)
tp, tq = get_object_pose("square_table_top")
for i, (h, a) in enumerate(holes):
    h, a = np.asarray(h), np.asarray(a)
    delta = np.asarray(lp) - h
    print("hole", i, "radial error", np.linalg.norm(delta - a * np.dot(delta, a)),
          "height", np.dot(delta, a), "half length", ext[1]/2)
print("relative orientation", rot(tq).T @ rot(lq))
print("ee", get_ee_pose())

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

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

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

leg = "square_table_leg4"
top = "square_table_top"
lp, lq = get_object_pose(leg)
tp, tq = get_object_pose(top)
hole, axis = map(np.asarray, get_holes(top)[2])
relative = mat(tq).T @ mat(lq)
remaining = -np.arctan2(relative[0, 2], relative[0, 0])
grasp_p = np.asarray(lp) + axis * 0.015
down = mat([0, 1, 0, 0])
start = None
for yaw in np.linspace(-np.pi, np.pi, 16, endpoint=False):
    candidate = R.from_rotvec(axis * yaw).as_matrix() @ down
    if all(is_reachable(grasp_p, wxyz(R.from_rotvec(axis * a).as_matrix() @ candidate))
           for a in np.linspace(0, remaining, 5)):
        start = candidate
        break

if start is not None:
    open_gripper()
    ep, eq = get_ee_pose()
    goto_pose(np.asarray(ep) + axis * 0.075, eq)
    goto_pose(grasp_p + axis * 0.075, wxyz(start))
    goto_pose(grasp_p, wxyz(start))
    close_gripper()
    for i in range(4):
        lp, lq = get_object_pose(leg)
        ep, eq = get_ee_pose()
        tp, tq = get_object_pose(top)
        relative = mat(tq).T @ mat(lq)
        error = -np.arctan2(relative[0, 2], relative[0, 0])
        if abs(error) < 0.025:
            break
        angle = np.clip(error, -np.pi/6, np.pi/6)
        turn = R.from_rotvec(axis * angle).as_matrix()
        target_p = np.asarray(lp) + turn @ (np.asarray(ep) - lp)
        target_q = wxyz(turn @ mat(eq))
        if not is_reachable(target_p, target_q):
            break
        goto_pose(target_p, target_q)
        print("turn", i, get_object_pose(leg))
    open_gripper()
    ep, eq = get_ee_pose()
    goto_pose(np.asarray(ep) + axis * 0.08, eq)

lp, lq = get_object_pose(leg)
tp, tq = get_object_pose(top)
print("leg", lp, lq)
print("relative axes", mat(tq).T @ mat(lq))
hole, axis = map(np.asarray, get_holes(top)[2])
delta = np.asarray(lp) - hole
print("height", np.dot(delta, axis), "radial error",
      np.linalg.norm(delta - axis * np.dot(delta, axis)))

# Code block 6
lp, lq = get_object_pose(leg)
hole, axis = map(np.asarray, get_holes(top)[2])
grasp_p = np.asarray(lp) + axis * 0.015
tp, tq = get_object_pose(top)
relative = mat(tq).T @ mat(lq)
error = -np.arctan2(relative[0, 2], relative[0, 0])

if abs(error) > 0.035:
    for yaw in np.linspace(-np.pi, np.pi, 16, endpoint=False):
        start = R.from_rotvec(axis * yaw).as_matrix() @ mat([0, 1, 0, 0])
        end = R.from_rotvec(axis * error).as_matrix() @ start
        if not (is_reachable(grasp_p, wxyz(start)) and
                is_reachable(grasp_p, wxyz(end))):
            continue
        goto_pose(grasp_p, wxyz(start), z_approach=0.06)
        close_gripper()
        for _ in range(3):
            lp, lq = get_object_pose(leg)
            ep, eq = get_ee_pose()
            tp, tq = get_object_pose(top)
            relative = mat(tq).T @ mat(lq)
            error = -np.arctan2(relative[0, 2], relative[0, 0])
            if abs(error) < 0.02:
                break
            turn = R.from_rotvec(axis * np.clip(error, -0.25, 0.25)).as_matrix()
            target_p = np.asarray(lp) + turn @ (np.asarray(ep) - lp)
            target_q = wxyz(turn @ mat(eq))
            if not is_reachable(target_p, target_q):
                break
            goto_pose(target_p, target_q)
        open_gripper()
        ep, eq = get_ee_pose()
        goto_pose(np.asarray(ep) + axis * 0.08, eq)
        break

lp, lq = get_object_pose(leg)
tp, tq = get_object_pose(top)
hole, axis = map(np.asarray, get_holes(top)[2])
delta = np.asarray(lp) - hole
print("leg", lp, lq)
print("relative axes", mat(tq).T @ mat(lq))
print("height", np.dot(delta, axis),
      "radial error", np.linalg.norm(delta - axis * np.dot(delta, axis)))

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

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

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

leg = "square_table_leg4"
top = "square_table_top"
lp, lq, extent = get_object_pose(leg, return_bbox_extent=True)
tp, tq = get_object_pose(top)
hole, axis = map(np.asarray, get_holes(top)[2])
axis = axis / np.linalg.norm(axis)
relative = matrix(tq).T @ matrix(lq)
error = -np.arctan2(relative[0, 2], relative[0, 0])
print("Initial angular error:", error)

if abs(error) > 0.04:
    grasp = np.asarray(lp) + axis * 0.022
    down = matrix([0, 1, 0, 0])
    for yaw in np.linspace(-np.pi, np.pi, 16, endpoint=False):
        start = R.from_rotvec(axis * yaw).as_matrix() @ down
        end = R.from_rotvec(axis * error).as_matrix() @ start
        if not all(is_reachable(grasp, quaternion(m)) for m in (start, end)):
            continue
        open_gripper()
        goto_pose(grasp, quaternion(start), z_approach=0.06)
        close_gripper()
        for step in range(3):
            lp, lq = get_object_pose(leg)
            ep, eq = get_ee_pose()
            tp, tq = get_object_pose(top)
            relative = matrix(tq).T @ matrix(lq)
            error = -np.arctan2(relative[0, 2], relative[0, 0])
            if abs(error) < 0.035:
                break
            turn = R.from_rotvec(axis * np.clip(error, -0.4, 0.4)).as_matrix()
            target = np.asarray(lp) + turn @ (np.asarray(ep) - lp) - axis * 0.001
            orientation = quaternion(turn @ matrix(eq))
            if not is_reachable(target, orientation):
                break
            goto_pose(target, orientation)
            print("After turn:", get_object_pose(leg))
        open_gripper()
        ep, eq = get_ee_pose()
        goto_pose(np.asarray(ep) + axis * 0.08, eq)
        break

lp, lq = get_object_pose(leg)
tp, tq = get_object_pose(top)
hole, axis = map(np.asarray, get_holes(top)[2])
delta = np.asarray(lp) - hole
print("Final relative axes:", matrix(tq).T @ matrix(lq))
print("Height:", np.dot(delta, axis))
print("Radial error:", np.linalg.norm(delta - axis * np.dot(delta, axis)))