import torch


def to_object_centric(X, Y):
    """
    X : (obj_x, obj_y, obj_θ, width, height, mass, friction)
    Y : (obj_x', obj_y', cos(obj_θ'), sin(obj_θ'))
    returns (width, height, mass, friction, x', y', θ'), all in object frame
    """
    res = torch.zeros(X.size(0), 7)
    if X.is_cuda:
        res = res.cuda()
    res[:, 0] = X[:, 3]  # width
    res[:, 1] = X[:, 4]  # height
    res[:, 2] = X[:, 5]  # mass
    res[:, 3] = X[:, 6]  # friction
    θ = X[:, 2]
    θ_ = torch.atan2(Y[:, 3], Y[:, 2])
    # translate object successor state
    Y_t = Y[:, :2] - X[:, :2]
    res[:, 4] = Y_t[:, 0] * torch.cos(θ) + Y_t[:, 1] * torch.sin(θ)
    res[:, 5] = -Y_t[:, 0] * torch.sin(θ) + Y_t[:, 1] * torch.cos(θ)
    res[:, 6] = torch.atan2(torch.sin(θ_ - θ), torch.cos(θ_ - θ))
    return res