import torch


def create_coref_tensors(pos_coref_mention_pairs, pos_eds,
                         neg_coref_mention_pairs=None, neg_eds=None):
    """ Combines positive and negative corefeference samples into tensors """
    neg_coref_mention_pairs = neg_coref_mention_pairs if neg_coref_mention_pairs else []
    neg_eds = neg_eds if neg_eds else []

    coref_mention_pairs = pos_coref_mention_pairs + neg_coref_mention_pairs
    coref_ed = pos_eds + neg_eds
    coref_types = [1] * len(pos_coref_mention_pairs) + [0] * len(neg_coref_mention_pairs)

    if coref_mention_pairs:
        coref_mention_pairs = torch.tensor(coref_mention_pairs, dtype=torch.long)
        coref_types = torch.tensor(coref_types, dtype=torch.long)
        coref_ed = torch.tensor(coref_ed, dtype=torch.long)
        coref_sample_masks = torch.ones([coref_mention_pairs.shape[0]], dtype=torch.bool)
    else:
        coref_mention_pairs = torch.zeros([1, 2], dtype=torch.long)
        coref_types = torch.zeros([1], dtype=torch.long)
        coref_ed = torch.zeros([1], dtype=torch.long)
        coref_sample_masks = torch.zeros([1], dtype=torch.long)

    return coref_mention_pairs, coref_types, coref_ed, coref_sample_masks