from . import *

# for sents
instance_fields = [
    'sent_index',
    'words', 'word_num', 'word_ids', 'word_span_idxs',
    'piece_idxs', 'attention_masks', 'word_lens',
    'edit_type_idxs', 'upos_type_idxs', 'xpos_type_idxs', 'feats_type_idxs',
    'head_idxs', 'deprel_idxs', 'word_mask',
    'headg_idxs', 'deprelg_idxs'

]

batch_fields = [
    'sent_index',
    'words', 'word_num', 'word_ids', 'word_span_idxs',
    'piece_idxs', 'attention_masks', 'word_lens',
    'edit_type_idxs', 'upos_type_idxs', 'xpos_type_idxs', 'feats_type_idxs',
    'upos_ids', 'xpos_ids', 'feats_ids',
    'head_idxs', 'deprel_idxs', 'word_mask',
    'headg_idxs', 'deprelg_idxs'

]

Instance = namedtuple('Instance', field_names=instance_fields)

Batch = namedtuple('Batch', field_names=batch_fields)



class TaggerDataset(Dataset):
    def __init__(self, config, input_conllu, gold_conllu, evaluate=False):
        self.config = config
        self.input_conllu = input_conllu
        self.gold_conllu = gold_conllu
        self.evaluate = evaluate

        self.treebank_name = config.treebank_name

        self.conllu_doc = []

        self.vocabs_fpath = os.path.join(self.config._save_dir, '{}.vocabs.json'.format(self.config.lang))
        self.vocabs = {}
        self.data = []
        self.load_data()

    def __len__(self):
        return len(self.data)

    def __getitem__(self, item):
        return self.data[item]

    def load_data(self):
        if not self.evaluate:
            self.vocabs, self.data, self.conllu_doc = tget_examples_from_conllu(
                self.config.wordpiece_splitter,
                self.config.max_input_length,
                self.gold_conllu, get_vocab=True)
            with open(self.vocabs_fpath, 'w') as f:
                json.dump(self.vocabs, f)
        else:
            with open(self.vocabs_fpath) as f:
                self.vocabs = json.load(f)
            self.data, self.conllu_doc = tget_examples_from_conllu(
                self.config.wordpiece_splitter,
                self.config.max_input_length,
                self.input_conllu
            )

        print('Loaded {} entries from {}'.format(len(self), self.input_conllu))

    def numberize(self):
        wordpiece_splitter = self.config.wordpiece_splitter
        data = []
        for inst in self.data:
            words = inst['words']
            pieces = [[p for p in wordpiece_splitter.tokenize(w) if p != '▁'] for w in words]
            for ps in pieces:
                if len(ps) == 0:
                    ps += ['-']
            word_lens = [len(x) for x in pieces]
            assert 0 not in word_lens
            flat_pieces = [p for ps in pieces for p in ps]
            assert len(flat_pieces) > 0

            word_span_idxs = []
            start = 1
            for l in word_lens:
                word_span_idxs.append([start, start + l])
                start += l

            # Pad word pieces with special tokens
            piece_idxs = wordpiece_splitter.encode(
                flat_pieces,
                add_special_tokens=True,
                max_length=self.config.max_input_length,
                truncation=True
            )
            assert len(piece_idxs) <= self.config.max_input_length

            attn_masks = [1] * len(piece_idxs)
            piece_idxs = piece_idxs
            assert len(piece_idxs) > 0

            edit_type_idxs = [self.vocabs[LEMMA][edit] for edit in inst[LEMMA]]
            upos_type_idxs = [self.vocabs[UPOS][upos] for upos in inst[UPOS]]
            xpos_type_idxs = [self.vocabs[XPOS][xpos] for xpos in inst[XPOS]]
            feats_type_idxs = [self.vocabs[FEATS][feats] for feats in inst[FEATS]]

            assert len(edit_type_idxs) == len(inst['words'])

            # head, deprel, word_mask
            head_idxs = [head for head in inst[HEAD]]
            deprel_idxs = [self.vocabs[DEPREL][deprel] for deprel in inst[DEPREL]]

            headg_idxs =[head for head in inst[HEADG]]
            # print(headg_idxs,inst[DEPRELG],self.vocabs[DEPRELG])
            deprelg_idxs =[self.vocabs[DEPRELG][deprel1] for deprel1 in inst[DEPRELG]]


            word_mask = [0] * (len(inst['words']) + 1)

            instance = Instance(
                sent_index=inst['sent_index'],
                word_ids=inst['word_ids'],
                words=inst['words'],
                word_num=len(inst['words']),
                piece_idxs=piece_idxs,
                attention_masks=attn_masks,
                word_span_idxs=word_span_idxs,
                word_lens=word_lens,
                edit_type_idxs=edit_type_idxs,
                upos_type_idxs=upos_type_idxs,
                xpos_type_idxs=xpos_type_idxs,
                feats_type_idxs=feats_type_idxs,
                head_idxs=head_idxs,
                deprel_idxs=deprel_idxs,
                headg_idxs=headg_idxs,
                deprelg_idxs=deprelg_idxs,
                word_mask=word_mask
            )
            data.append(instance)
        self.data = data

    def collate_fn(self, batch):
        batch_sent_index = [inst.sent_index for inst in batch]

        batch_words = [inst.words for inst in batch]
        batch_word_num = [inst.word_num for inst in batch]
        batch_word_ids = [inst.word_ids for inst in batch]

        batch_piece_idxs = []
        batch_attention_masks = []
        batch_word_lens = []
        batch_word_span_idxs = []

        batch_edit_type_idxs = []
        batch_upos_type_idxs = []
        batch_xpos_type_idxs = []
        batch_feats_type_idxs = []

        batch_upos_ids, batch_xpos_ids, batch_feats_ids = [], [], []

        batch_head_ids, batch_deprel_ids, batch_word_mask = [], [], []
        
        batch_headg_ids, batch_deprelg_ids, batch_word_mask = [], [], []


        max_word_num = max(batch_word_num)
        max_wordpiece_num = max([len(inst.piece_idxs) for inst in batch])

        for inst in batch:
            batch_piece_idxs.append(inst.piece_idxs + [0] * (max_wordpiece_num - len(inst.piece_idxs)))
            batch_attention_masks.append(inst.attention_masks + [0] * (max_wordpiece_num - len(inst.piece_idxs)))
            batch_word_lens.append(inst.word_lens)
            batch_word_span_idxs.append(inst.word_span_idxs + [[1, 2]] * (max_word_num - inst.word_num))

            # lemmatization
            batch_edit_type_idxs.extend(inst.edit_type_idxs +
                                        [-100] * (max_word_num - inst.word_num))
            # upos, xpos, feats
            batch_upos_type_idxs.extend(inst.upos_type_idxs +
                                        [-100] * (max_word_num - inst.word_num))
            batch_xpos_type_idxs.extend(inst.xpos_type_idxs +
                                        [-100] * (max_word_num - inst.word_num))
            batch_feats_type_idxs.extend(inst.feats_type_idxs +
                                         [-100] * (max_word_num - inst.word_num))
            # head, deprel
            batch_head_ids.append(inst.head_idxs + [0] * (max_word_num - inst.word_num))
            batch_deprel_ids.append(inst.deprel_idxs + [0] * (max_word_num - inst.word_num))
            # print(batch_head_ids)
            # exit()
            batch_headg_ids.append(inst.headg_idxs + [0] * (max_word_num - inst.word_num))
            batch_deprelg_ids.append(inst.deprelg_idxs + [0] * (max_word_num - inst.word_num))
            
            batch_word_mask.append(inst.word_mask + [1] * (max_word_num - inst.word_num))
            # ids for feature building
            batch_upos_ids.append(inst.upos_type_idxs + [0] * (max_word_num - inst.word_num))
            batch_xpos_ids.append(inst.xpos_type_idxs + [0] * (max_word_num - inst.word_num))
            batch_feats_ids.append(inst.feats_type_idxs + [0] * (max_word_num - inst.word_num))

        batch_piece_idxs = torch.tensor(batch_piece_idxs, dtype=torch.long, device=self.config.device)
        batch_attention_masks = torch.tensor(batch_attention_masks, dtype=torch.float16, device=self.config.device)
        batch_edit_type_idxs = torch.tensor(batch_edit_type_idxs, dtype=torch.long, device=self.config.device)
        batch_word_span_idxs = torch.tensor(batch_word_span_idxs, dtype=torch.long, device=self.config.device)

        batch_upos_type_idxs = torch.tensor(batch_upos_type_idxs, dtype=torch.long, device=self.config.device)
        batch_xpos_type_idxs = torch.tensor(batch_xpos_type_idxs, dtype=torch.long, device=self.config.device)
        batch_feats_type_idxs = torch.tensor(batch_feats_type_idxs, dtype=torch.long, device=self.config.device)

        batch_upos_ids = torch.tensor(batch_upos_ids, dtype=torch.long, device=self.config.device)
        batch_xpos_ids = torch.tensor(batch_xpos_ids, dtype=torch.long, device=self.config.device)

        # print("id as  ",batch_head_ids)
        #batch_head_ids = torch.zeros((1,2))#torch.tensor(batch_head_ids, dtype=torch.long, device=self.config.device)

        batch_head_ids = torch.tensor(batch_head_ids, dtype=torch.long, device=self.config.device)
        batch_deprel_ids = torch.tensor(batch_deprel_ids, dtype=torch.long, device=self.config.device)

        batch_headg_ids = torch.tensor(batch_headg_ids, dtype=torch.long, device=self.config.device)
        batch_deprelg_ids = torch.tensor(batch_deprelg_ids, dtype=torch.long, device=self.config.device)

        batch_word_mask = torch.tensor(batch_word_mask, dtype=torch.bool, device=self.config.device)

        return Batch(
            sent_index=batch_sent_index,
            word_ids=batch_word_ids,
            words=batch_words,
            word_num=batch_word_num,
            piece_idxs=batch_piece_idxs,
            attention_masks=batch_attention_masks,
            word_lens=batch_word_lens,
            word_span_idxs=batch_word_span_idxs,
            edit_type_idxs=batch_edit_type_idxs,
            upos_type_idxs=batch_upos_type_idxs,
            xpos_type_idxs=batch_xpos_type_idxs,
            feats_type_idxs=batch_feats_type_idxs,
            upos_ids=batch_upos_ids,
            xpos_ids=batch_xpos_ids,
            feats_ids=batch_xpos_ids,
            head_idxs=batch_head_ids,
            deprel_idxs=batch_deprel_ids,
            headg_idxs=batch_headg_ids,
            deprelg_idxs=batch_deprelg_ids,
            word_mask=batch_word_mask
        )
