import torch
from data.field.mini_torchtext.field import RawField
from data.field.mini_torchtext.vocab import Vocab
from collections import Counter


class LabelField(RawField):
    def __self__(self, preprocessing):
        super(LabelField, self).__init__(preprocessing=preprocessing)
        self.vocab = None

    def build_vocab(self, *args, **kwargs):
        sources = []
        for arg in args:
            if isinstance(arg, torch.utils.data.Dataset):
                sources += [arg.get_examples(name) for name, field in arg.fields.items() if field is self]
            else:
                sources.append(arg)

        counter = Counter()
        for data in sources:
            for x in data:
                counter.update(x)

        self.vocab = Vocab(counter, specials=[])

    def process(self, example, device=None):
        tensor, lengths = self.numericalize(example, device=device)
        return tensor, lengths

    def numericalize(self, example, device=None):
        example = [self.vocab.stoi[x] + 1 for x in example]
        length = torch.LongTensor([len(example)], device=device).squeeze(0)
        tensor = torch.LongTensor(example, device=device)

        return tensor, length