import torch


def build_design_tensor(item_thetas, individual_assignment):
    """
    Given a sequence of angles representing objects on S^1, and
    assignments of those objects to individuals, creates a design tensor
    of the appropriate dimensions.

    :param torch.Tensor item_thetas: a tensor of dimension batch x n representing
        the item to be used at the nth trial
    :param torch.Tensor individual_assignment: a tensor of dimension batch x n x 2
        with 0-1 rows indicating the individual to assign item n to
    """
    # batch x n x 2
    item_features = torch.stack([item_thetas.cos(), -item_thetas.sin()], dim=-1)
    ind1 = individual_assignment[..., 0].unsqueeze(-1)*item_features
    ind2 = individual_assignment[..., 1].unsqueeze(-1)*item_features
    # batch x n x 6
    return torch.cat([item_features, ind1, ind2], dim=-1)