import os
import logging

import torch 
from torch.utils.data import TensorDataset

class InputExample(object):
    """A single training/test example for simple sequence classification."""

    def __init__(self, guid, text_a, text_b=None, label=None):
        """Constructs a InputExample.

        Args:
            guid: Unique id for the example.
            text_a: string. The untokenized text of the first sequence. For single
            sequence tasks, only this sequence must be specified.
            text_b: (Optional) string. The untokenized text of the second sequence.
            Only must be specified for sequence pair tasks.
            label: (Optional) string. The label of the example. This should be
            specified for train and dev examples, but not for test examples.
        """
        self.guid = guid
        self.text_a = text_a
        self.text_b = text_b
        self.label = label


class InputFeatures(object):
    """A single set of features of data."""

    def __init__(self, input_ids, input_mask, label_id, valid_ids=None, label_mask=None):
        self.input_ids = input_ids
        self.input_mask = input_mask
        self.label_id = label_id
        self.valid_ids = valid_ids
        self.label_mask = label_mask


class NerProcessor:
    """Processor for the CoNLL-2003 data set."""

    def get_train_examples(self, data_dir):
        """See base class."""
        return self._create_examples(
            self._read_file(os.path.join(data_dir, "train.txt")), "train")

    def get_dev_examples(self, data_dir):
        """See base class."""
        return self._create_examples(
            self._read_file(os.path.join(data_dir, "valid.txt")), "valid")

    def get_test_examples(self, data_dir):
        """See base class."""
        return self._create_examples(
            self._read_file(os.path.join(data_dir, "test.txt")), "test")

    def get_labels(self):
        return [
            "B-ALINK_H",
            "B-ALINK_T",
            "B-SLINK_H",
            "B-SLINK_T",
            "O"
        ]
        '''return [
            "B-ACTION",
            "I-ACTION",
            "B-ASPECTUAL",
            "B-DONT_KNOW",
            "B-I_ACTION",
            "B-I_STATE",
            "B-LIGHT_PREDICATE",
            "B-MODALITY",
            "I-MODALITY",
            "O",
            "B-PERCEPTION",
            "B-REPORTING",
            "B-STATE",
            ]'''


        #return ["O", "B-PER", "I-PER", "B-ORG", "I-ORG", "B-LOC", "I-LOC", "B-MISC", "I-MISC"]

    def _read_file(self, filename):
        '''
        read file
        '''
        f = open(filename)
        data = []
        sentence = []
        label = []

        for i, line in enumerate(f, 1):
            if not line.strip() or len(line) == 0 or line.startswith('-DOCSTART') or line[0] == "\n" or line[0] == '.':
                if len(sentence) > 0:
                    data.append((sentence, label))
                    sentence = []
                    label = []
                continue

            splits = line.split()
            assert len(splits) >= 2, "error on line {}. Found {} splits".format(i, len(splits))
            word, tag = splits[0], splits[-1]
            assert tag in self.get_labels(), "unknown tag {} in line {}".format(tag, i)
            sentence.append(word.strip())
            label.append(tag.strip())

        if len(sentence) > 0:
            data.append((sentence, label))
            sentence = []
            label = []
        return data

    def _create_examples(self, lines, set_type):
        examples = []

        for i, (sentence, label) in enumerate(lines):
            guid = "%s-%s" % (set_type, i)
            text_a = ' '.join(sentence)
            text_b = None
            label = label
            examples.append(InputExample(
                guid=guid, text_a=text_a, text_b=text_b, label=label))
        return examples


def convert_examples_to_features(examples, label_list, max_seq_length, encode_method):
    """Converts a set of examples into XLMR compatible format

    * Labels are only assigned to the positions correspoinding to the first BPE token of each word.
    * Other positions are labeled with 0 ("IGNORE")

    """
    ignored_label = "IGNORE"
    label_map = {label: i for i, label in enumerate(label_list, 1)}
    label_map[ignored_label] = 0  # 0 label is to be ignored
    
    features = []
    for (ex_index, example) in enumerate(examples):

        textlist = example.text_a.split(' ')
        labellist = example.label
        labels = []
        valid = []
        label_mask = []
        token_ids = []
       
        for i, word in enumerate(textlist):  
            tokens = encode_method(word.strip())  # word token ids   
            token_ids.extend(tokens)  # all sentence token ids
            label_1 = labellist[i]
            for m in range(len(tokens)):
                if m == 0:  # only label the first BPE token of each work
                    labels.append(label_1)
                    valid.append(1)
                    label_mask.append(1)
                else:
                    labels.append(ignored_label)  # unlabeled BPE token
                    label_mask.append(0)
                    valid.append(0)

        logging.debug("token ids = ")
        logging.debug(token_ids)
        logging.debug("labels = ")
        logging.debug(labels)
        logging.debug("valid = ")
        logging.debug(valid)

        if len(token_ids) >= max_seq_length - 1:  # trim extra tokens
            token_ids = token_ids[0:(max_seq_length-2)]
            labels = labels[0:(max_seq_length-2)]
            valid = valid[0:(max_seq_length-2)]
            label_mask = label_mask[0:(max_seq_length-2)]

        # adding <s>
        token_ids.insert(0, 0)
        labels.insert(0, ignored_label)
        label_mask.insert(0, 0)
        valid.insert(0, 0)

        # adding </s>
        token_ids.append(2)
        labels.append(ignored_label)
        label_mask.append(0)
        valid.append(0)

        assert len(token_ids) == len(labels)
        assert len(valid) == len(labels)

        label_ids = []
        for i, _ in enumerate(token_ids):
            label_ids.append(label_map[labels[i]])

        assert len(token_ids) == len(label_ids)
        assert len(valid) == len(label_ids)

        input_mask = [1] * len(token_ids)

        while len(token_ids) < max_seq_length:
            token_ids.append(1)  # token padding idx
            input_mask.append(0)
            label_ids.append(label_map[ignored_label])  # label ignore idx
            valid.append(0)
            label_mask.append(0)

        while len(label_ids) < max_seq_length:
            label_ids.append(label_map[ignored_label])
            label_mask.append(0)

        assert len(token_ids) == max_seq_length
        assert len(input_mask) == max_seq_length
        assert len(label_ids) == max_seq_length
        assert len(valid) == max_seq_length
        assert len(label_mask) == max_seq_length

        if ex_index < 2:
            logging.info("*** Example ***")
            logging.info("guid: %s" % (example.guid))
            logging.info("tokens: %s" % " ".join(
                [str(x) for x in token_ids]))
            logging.info("input_ids: %s" %
                         " ".join([str(x) for x in token_ids]))
            logging.info("input_mask: %s" %
                         " ".join([str(x) for x in input_mask]))
            logging.info("label: %s (id = %s)" % (example.label, " ".join(map(str, label_ids))))
            logging.info("label_mask: %s" %
                         " ".join([str(x) for x in label_mask]))
            logging.info("valid mask: %s" %
                        " ".join([str(x) for x in valid]))

        features.append(
            InputFeatures(input_ids=token_ids,
                          input_mask=input_mask,
                          label_id=label_ids,
                          valid_ids=valid,
                          label_mask=label_mask))

    return features


def create_dataset(features):
    
    all_input_ids = torch.tensor(
        [f.input_ids for f in features], dtype=torch.long)
    all_label_ids = torch.tensor(
        [f.label_id for f in features], dtype=torch.long)
    all_valid_ids = torch.tensor(
        [f.valid_ids for f in features], dtype=torch.long)
    all_lmask_ids = torch.tensor(
        [f.label_mask for f in features], dtype=torch.long)

    return TensorDataset(
        all_input_ids, all_label_ids, all_lmask_ids, all_valid_ids)
