import numpy as np

def sample_insertion_pose():
        # Peg
        x_range = [0.1, 0.2]
        y_range = [0.4, 0.6]
        z_range = [0.05, 0.05]

        ranges = np.vstack([x_range, y_range, z_range])
        peg_position = np.random.uniform(ranges[:, 0], ranges[:, 1])

        peg_quat = np.array([1, 0, 0, 0])
        peg_pose = np.concatenate([peg_position, peg_quat])

        # Socket
        x_range = [-0.2, -0.1]
        y_range = [0.4, 0.6]
        z_range = [0.05, 0.05]

        ranges = np.vstack([x_range, y_range, z_range])
        socket_position = np.random.uniform(ranges[:, 0], ranges[:, 1])

        socket_quat = np.array([1, 0, 0, 0])
        socket_pose = np.concatenate([socket_position, socket_quat])

        return peg_pose, socket_pose

def sample_transfer_pose():
        # Box
        x_range = [-0.2, 0.2]
        y_range = [0.7, 1.1]
        z_range = [0.47, 0.47]

        ranges = np.vstack([x_range, y_range, z_range])
        box_position = np.random.uniform(ranges[:, 0], ranges[:, 1])


        return box_position