# -*- coding:utf-8 -*-
# Author: hankcs
# Date: 2020-05-08 15:30
from abc import abstractmethod
from collections import Counter
from typing import Union, Tuple, Iterable, Any, Generator

import numpy as np
import tensorflow as tf
from transformers import PreTrainedTokenizer, PretrainedConfig

from hanlp_common.constant import ROOT
from hanlp_common.structure import SerializableDict
from hanlp.common.transform_tf import Transform
from hanlp.common.vocab_tf import VocabTF
from hanlp.components.parsers.alg_tf import tolist, kmeans, randperm, arange
from hanlp.components.parsers.conll import read_conll
from hanlp_common.conll import CoNLLWord, CoNLLUWord, CoNLLSentence
from hanlp.layers.transformers.utils_tf import config_is, adjust_tokens_for_transformers, convert_examples_to_features
from hanlp.utils.log_util import logger
from hanlp.utils.string_util import ispunct
from hanlp_common.util import merge_locals_kwargs


class CoNLLTransform(Transform):

    def __init__(self, config: SerializableDict = None, map_x=True, map_y=True, lower=True, n_buckets=32, min_freq=2,
                 use_pos=True, **kwargs) -> None:
        super().__init__(**merge_locals_kwargs(locals(), kwargs))
        self.form_vocab: VocabTF = None
        if use_pos:
            self.cpos_vocab: VocabTF = None
        self.rel_vocab: VocabTF = None
        self.puncts: tf.Tensor = None

    @property
    def use_pos(self):
        return self.config.get('use_pos', True)

    def x_to_idx(self, x) -> Union[tf.Tensor, Tuple]:
        form, cpos = x
        return self.form_vocab.token_to_idx_table.lookup(form), self.cpos_vocab.token_to_idx_table.lookup(cpos)

    def y_to_idx(self, y):
        head, rel = y
        return head, self.rel_vocab.token_to_idx_table.lookup(rel)

    def X_to_inputs(self, X: Union[tf.Tensor, Tuple[tf.Tensor]]) -> Iterable:
        if len(X) == 2:
            form_batch, cposes_batch = X
            mask = tf.not_equal(form_batch, 0)
        elif len(X) == 3:
            form_batch, cposes_batch, mask = X
        else:
            raise ValueError(f'Expect X to be 2 or 3 elements but got {repr(X)}')
        sents = []

        for form_sent, cposes_sent, length in zip(form_batch, cposes_batch,
                                                  tf.math.count_nonzero(mask, axis=-1)):
            forms = tolist(form_sent)[1:length + 1]
            cposes = tolist(cposes_sent)[1:length + 1]
            sents.append([(self.form_vocab.idx_to_token[f],
                           self.cpos_vocab.idx_to_token[c]) for f, c in zip(forms, cposes)])

        return sents

    def lock_vocabs(self):
        super().lock_vocabs()
        self.puncts = tf.constant([i for s, i in self.form_vocab.token_to_idx.items()
                                   if ispunct(s)], dtype=tf.int64)

    def file_to_inputs(self, filepath: str, gold=True):
        assert gold, 'only support gold file for now'
        use_pos = self.use_pos
        conllu = filepath.endswith('.conllu')
        for sent in read_conll(filepath):
            for i, cell in enumerate(sent):
                form = cell[1]
                cpos = cell[3]
                head = cell[6]
                deprel = cell[7]
                # if conllu:
                #     deps = cell[8]
                #     deps = [x.split(':', 1) for x in deps.split('|')]
                #     heads = [int(x[0]) for x in deps if '_' not in x[0] and '.' not in x[0]]
                #     rels = [x[1] for x in deps if '_' not in x[0] and '.' not in x[0]]
                #     if head in heads:
                #         offset = heads.index(head)
                #         if not self.rel_vocab or rels[offset] in self.rel_vocab:
                #             deprel = rels[offset]
                sent[i] = [form, cpos, head, deprel] if use_pos else [form, head, deprel]
            yield sent

    @property
    def bos(self):
        if self.form_vocab.idx_to_token is None:
            return ROOT
        return self.form_vocab.idx_to_token[2]

    def input_is_single_sample(self, input: Any) -> bool:
        if self.use_pos:
            return isinstance(input[0][0], str) if len(input[0]) else False
        else:
            return isinstance(input[0], str) if len(input[0]) else False

    @abstractmethod
    def batched_inputs_to_batches(self, corpus, indices, shuffle):
        pass

    def len_of_sent(self, sent):
        return 1 + len(sent)  # take ROOT into account

    def samples_to_dataset(self, samples: Generator, map_x=None, map_y=None, batch_size=5000, shuffle=None, repeat=None,
                           drop_remainder=False, prefetch=1, cache=True) -> tf.data.Dataset:
        if shuffle:
            def generator():
                # custom bucketing, load corpus into memory
                corpus = list(x for x in (samples() if callable(samples) else samples))
                lengths = [self.len_of_sent(i) for i in corpus]
                if len(corpus) < 32:
                    n_buckets = 1
                else:
                    n_buckets = min(self.config.n_buckets, len(corpus))
                buckets = dict(zip(*kmeans(lengths, n_buckets)))
                sizes, buckets = zip(*[
                    (size, bucket) for size, bucket in buckets.items()
                ])
                # the number of chunks in each bucket, which is clipped by
                # range [1, len(bucket)]
                chunks = [min(len(bucket), max(round(size * len(bucket) / batch_size), 1)) for size, bucket in
                          zip(sizes, buckets)]
                range_fn = randperm if shuffle else arange
                max_samples_per_batch = self.config.get('max_samples_per_batch', None)
                for i in tolist(range_fn(len(buckets))):
                    split_sizes = [(len(buckets[i]) - j - 1) // chunks[i] + 1
                                   for j in range(chunks[i])]  # how many sentences in each batch
                    for batch_indices in tf.split(range_fn(len(buckets[i])), split_sizes):
                        indices = [buckets[i][j] for j in tolist(batch_indices)]
                        if max_samples_per_batch:
                            for j in range(0, len(indices), max_samples_per_batch):
                                yield from self.batched_inputs_to_batches(corpus, indices[j:j + max_samples_per_batch],
                                                                          shuffle)
                        else:
                            yield from self.batched_inputs_to_batches(corpus, indices, shuffle)

        else:
            def generator():
                # custom bucketing, load corpus into memory
                corpus = list(x for x in (samples() if callable(samples) else samples))
                n_tokens = 0
                batch = []
                for idx, sent in enumerate(corpus):
                    sent_len = self.len_of_sent(sent)
                    if n_tokens + sent_len > batch_size and batch:
                        yield from self.batched_inputs_to_batches(corpus, batch, shuffle)
                        n_tokens = 0
                        batch = []
                    n_tokens += sent_len
                    batch.append(idx)
                if batch:
                    yield from self.batched_inputs_to_batches(corpus, batch, shuffle)

        # next(generator())
        return Transform.samples_to_dataset(self, generator, False, False, 0, False, repeat, drop_remainder, prefetch,
                                            cache)


class CoNLL_DEP_Transform(CoNLLTransform):

    def __init__(self, config: SerializableDict = None, map_x=True, map_y=True, lower=True, n_buckets=32,
                 min_freq=2, **kwargs) -> None:
        super().__init__(config, map_x, map_y, lower, n_buckets, min_freq, **kwargs)

    def batched_inputs_to_batches(self, corpus, indices, shuffle):
        """Convert batched inputs to batches of samples

        Args:
          corpus(list): A list of inputs
          indices(list): A list of indices, each list belongs to a batch
          shuffle:

        Returns:


        """
        raw_batch = [[], [], [], []]
        for idx in indices:
            for b in raw_batch:
                b.append([])
            for cells in corpus[idx]:
                for b, c, v in zip(raw_batch, cells,
                                   [self.form_vocab, self.cpos_vocab, None, self.rel_vocab]):
                    b[-1].append(v.get_idx_without_add(c) if v else c)
        batch = []
        for b, v in zip(raw_batch, [self.form_vocab, self.cpos_vocab, None, self.rel_vocab]):
            b = tf.keras.preprocessing.sequence.pad_sequences(b, padding='post',
                                                              value=v.safe_pad_token_idx if v else 0,
                                                              dtype='int64')
            batch.append(b)
        assert len(batch) == 4
        yield (batch[0], batch[1]), (batch[2], batch[3])

    def create_types_shapes_values(self) -> Tuple[Tuple, Tuple, Tuple]:
        types = (tf.int64, tf.int64), (tf.int64, tf.int64)
        shapes = ([None, None], [None, None]), ([None, None], [None, None])
        values = (self.form_vocab.safe_pad_token_idx, self.cpos_vocab.safe_pad_token_idx), (
            0, self.rel_vocab.safe_pad_token_idx)
        return types, shapes, values

    def inputs_to_samples(self, inputs, gold=False):
        token_mapping: dict = self.config.get('token_mapping', None)
        use_pos = self.config.get('use_pos', True)
        for sent in inputs:
            sample = []
            for i, cell in enumerate(sent):
                if isinstance(cell, tuple):
                    cell = list(cell)
                elif isinstance(cell, str):
                    cell = [cell]
                if token_mapping:
                    cell[0] = token_mapping.get(cell[0], cell[0])
                if self.config['lower']:
                    cell[0] = cell[0].lower()
                if not gold:
                    cell += [0, self.rel_vocab.safe_pad_token]
                sample.append(cell)
            # insert root word with arbitrary fields, anyway it will be masked
            # form, cpos, head, deprel = sample[0]
            sample.insert(0, [self.bos, self.bos, 0, self.bos] if use_pos else [self.bos, 0, self.bos])
            yield sample

    def XY_to_inputs_outputs(self, X: Union[tf.Tensor, Tuple[tf.Tensor]], Y: Union[tf.Tensor, Tuple[tf.Tensor]],
                             gold=False, inputs=None, conll=True, arc_scores=None, rel_scores=None) -> Iterable:
        (words, feats, mask), (arc_preds, rel_preds) = X, Y
        if inputs is None:
            inputs = self.X_to_inputs(X)
        ys = self.Y_to_outputs((arc_preds, rel_preds, mask), inputs=inputs)
        sents = []
        for x, y in zip(inputs, ys):
            sent = CoNLLSentence()
            for idx, (cell, (head, deprel)) in enumerate(zip(x, y)):
                if self.use_pos and not self.config.get('joint_pos', None):
                    form, cpos = cell
                else:
                    form, cpos = cell, None
                if conll:
                    sent.append(
                        CoNLLWord(id=idx + 1, form=form, cpos=cpos, head=head, deprel=deprel) if conll == '.conll'
                        else CoNLLUWord(id=idx + 1, form=form, upos=cpos, head=head, deprel=deprel))
                else:
                    sent.append([head, deprel])
            sents.append(sent)
        return sents

    def fit(self, trn_path: str, **kwargs) -> int:
        use_pos = self.config.use_pos
        self.form_vocab = VocabTF()
        self.form_vocab.add(ROOT)  # make root the 2ed elements while 0th is pad, 1st is unk
        if self.use_pos:
            self.cpos_vocab = VocabTF(pad_token=None, unk_token=None)
        self.rel_vocab = VocabTF(pad_token=None, unk_token=None)
        num_samples = 0
        counter = Counter()
        for sent in self.file_to_samples(trn_path, gold=True):
            num_samples += 1
            for idx, cell in enumerate(sent):
                if use_pos:
                    form, cpos, head, deprel = cell
                else:
                    form, head, deprel = cell
                if idx == 0:
                    root = form
                else:
                    counter[form] += 1
                if use_pos:
                    self.cpos_vocab.add(cpos)
                self.rel_vocab.add(deprel)

        for token in [token for token, freq in counter.items() if freq >= self.config.min_freq]:
            self.form_vocab.add(token)
        return num_samples

    @property
    def root_rel_idx(self):
        root_rel_idx = self.config.get('root_rel_idx', None)
        if root_rel_idx is None:
            for idx, rel in enumerate(self.rel_vocab.idx_to_token):
                if 'root' in rel.lower() and rel != self.bos:
                    self.config['root_rel_idx'] = root_rel_idx = idx
                    break
        return root_rel_idx

    def Y_to_outputs(self, Y: Union[tf.Tensor, Tuple[tf.Tensor]], gold=False, inputs=None, X=None) -> Iterable:
        arc_preds, rel_preds, mask = Y
        sents = []

        for arc_sent, rel_sent, length in zip(arc_preds, rel_preds,
                                              tf.math.count_nonzero(mask, axis=-1)):
            arcs = tolist(arc_sent)[1:length + 1]
            rels = tolist(rel_sent)[1:length + 1]
            sents.append([(a, self.rel_vocab.idx_to_token[r]) for a, r in zip(arcs, rels)])

        return sents


class CoNLL_Transformer_Transform(CoNLL_DEP_Transform):

    def __init__(self, config: SerializableDict = None, map_x=True, map_y=True,
                 lower=True, n_buckets=32, min_freq=0, max_seq_length=256, use_pos=False,
                 mask_p=None, graph=False, topk=None,
                 **kwargs) -> None:
        super().__init__(**merge_locals_kwargs(locals(), kwargs))
        self.tokenizer: PreTrainedTokenizer = None
        self.transformer_config: PretrainedConfig = None
        if graph:
            self.orphan_relation = ROOT

    def lock_vocabs(self):
        super().lock_vocabs()
        if self.graph:
            CoNLL_SDP_Transform._find_orphan_relation(self)

    def fit(self, trn_path: str, **kwargs) -> int:
        if self.config.get('joint_pos', None):
            self.config.use_pos = True
        if self.graph:
            # noinspection PyCallByClass
            num = CoNLL_SDP_Transform.fit(self, trn_path, **kwargs)
        else:
            num = super().fit(trn_path, **kwargs)
        if self.config.get('topk', None):
            counter = Counter()
            for sent in self.file_to_samples(trn_path, gold=True):
                for idx, cell in enumerate(sent):
                    form, head, deprel = cell
                    counter[form] += 1
            self.topk_vocab = VocabTF()
            for k, v in counter.most_common(self.config.topk):
                self.topk_vocab.add(k)
        return num

    def inputs_to_samples(self, inputs, gold=False):
        if self.graph:
            yield from CoNLL_SDP_Transform.inputs_to_samples(self, inputs, gold)
        else:
            yield from super().inputs_to_samples(inputs, gold)

    def file_to_inputs(self, filepath: str, gold=True):
        if self.graph:
            yield from CoNLL_SDP_Transform.file_to_inputs(self, filepath, gold)
        else:
            yield from super().file_to_inputs(filepath, gold)

    @property
    def mask_p(self) -> float:
        return self.config.get('mask_p', None)

    @property
    def graph(self):
        return self.config.get('graph', None)

    def create_types_shapes_values(self) -> Tuple[Tuple, Tuple, Tuple]:
        mask_p = self.mask_p
        types = (tf.int64, (tf.int64, tf.int64, tf.int64)), (tf.bool if self.graph else tf.int64, tf.int64, tf.int64) if mask_p else (
            tf.bool if self.graph else tf.int64, tf.int64)
        if self.graph:
            shapes = ([None, None], ([None, None], [None, None], [None, None])), (
                [None, None, None], [None, None, None], [None, None]) if mask_p else (
                [None, None, None], [None, None, None])
        else:
            shapes = ([None, None], ([None, None], [None, None], [None, None])), (
                [None, None], [None, None], [None, None]) if mask_p else ([None, None], [None, None])

        values = (self.form_vocab.safe_pad_token_idx, (0, 0, 0)), \
                 (0, self.rel_vocab.safe_pad_token_idx, 0) if mask_p else (0, self.rel_vocab.safe_pad_token_idx)
        types_shapes_values = types, shapes, values
        if self.use_pos:
            types_shapes_values = [((shapes[0][0], shapes[0][1] + (shapes[0][0],)), shapes[1]) for shapes in
                                   types_shapes_values]
        return types_shapes_values

    def X_to_inputs(self, X: Union[tf.Tensor, Tuple[tf.Tensor]]) -> Iterable:
        form_batch, feat, prefix_mask = X
        sents = []

        for form_sent, length in zip(form_batch, tf.math.count_nonzero(prefix_mask, axis=-1)):
            forms = tolist(form_sent)[1:length + 1]
            sents.append([self.form_vocab.idx_to_token[f] for f in forms])

        return sents

    def batched_inputs_to_batches(self, corpus, indices, shuffle):
        use_pos = self.use_pos
        if use_pos:
            raw_batch = [[], [], [], []]
        else:
            raw_batch = [[], [], []]
        if self.graph:
            max_len = len(max([corpus[i] for i in indices], key=len))
            for idx in indices:
                arc = np.zeros((max_len, max_len), dtype=np.bool)
                rel = np.zeros((max_len, max_len), dtype=np.int64)
                for b in raw_batch[:2 if use_pos else 1]:
                    b.append([])
                for m, cells in enumerate(corpus[idx]):
                    if use_pos:
                        for b, c, v in zip(raw_batch, cells, [None, self.cpos_vocab]):
                            b[-1].append(v.get_idx_without_add(c) if v else c)
                    else:
                        for b, c, v in zip(raw_batch, cells, [None]):
                            b[-1].append(c)
                    for n, r in zip(cells[-2], cells[-1]):
                        arc[m, n] = True
                        rid = self.rel_vocab.get_idx_without_add(r)
                        if rid is None:
                            logger.warning(f'Relation OOV: {r} not exists in train')
                            continue
                        rel[m, n] = rid
                raw_batch[-2].append(arc)
                raw_batch[-1].append(rel)
        else:
            for idx in indices:
                for s in raw_batch:
                    s.append([])
                for cells in corpus[idx]:
                    if use_pos:
                        for s, c, v in zip(raw_batch, cells, [None, self.cpos_vocab, None, self.rel_vocab]):
                            s[-1].append(v.get_idx_without_add(c) if v else c)
                    else:
                        for s, c, v in zip(raw_batch, cells, [None, None, self.rel_vocab]):
                            s[-1].append(v.get_idx_without_add(c) if v else c)

        # Transformer tokenizing
        config = self.transformer_config
        tokenizer = self.tokenizer
        xlnet = config_is(config, 'xlnet')
        roberta = config_is(config, 'roberta')
        pad_token = tokenizer.pad_token
        pad_token_id = tokenizer.convert_tokens_to_ids([pad_token])[0]
        cls_token = tokenizer.cls_token
        sep_token = tokenizer.sep_token
        max_seq_length = self.config.max_seq_length
        batch_forms = []
        batch_input_ids = []
        batch_input_mask = []
        batch_prefix_offset = []
        mask_p = self.mask_p
        if mask_p:
            batch_masked_offsets = []
            mask_token_id = tokenizer.mask_token_id
        for sent_idx, sent in enumerate(raw_batch[0]):
            batch_forms.append([self.form_vocab.get_idx_without_add(token) for token in sent])
            sent = adjust_tokens_for_transformers(sent)
            sent = sent[1:]  # remove <root> use [CLS] instead
            pad_label_idx = self.form_vocab.pad_idx
            input_ids, input_mask, segment_ids, prefix_mask = \
                convert_examples_to_features(sent,
                                             max_seq_length,
                                             tokenizer,
                                             cls_token_at_end=xlnet,
                                             # xlnet has a cls token at the end
                                             cls_token=cls_token,
                                             cls_token_segment_id=2 if xlnet else 0,
                                             sep_token=sep_token,
                                             sep_token_extra=roberta,
                                             # roberta uses an extra separator b/w pairs of sentences, cf. github.com/pytorch/fairseq/commit/1684e166e3da03f5b600dbb7855cb98ddfcd0805
                                             pad_on_left=xlnet,
                                             # pad on the left for xlnet
                                             pad_token_id=pad_token_id,
                                             pad_token_segment_id=4 if xlnet else 0,
                                             pad_token_label_id=pad_label_idx,
                                             do_padding=False)
            num_masks = sum(prefix_mask)
            # assert len(sent) == num_masks  # each token has a True subtoken
            if num_masks < len(sent):  # long sent gets truncated, +1 for root
                batch_forms[-1] = batch_forms[-1][:num_masks + 1]  # form
                raw_batch[-1][sent_idx] = raw_batch[-1][sent_idx][:num_masks + 1]  # head
                raw_batch[-2][sent_idx] = raw_batch[-2][sent_idx][:num_masks + 1]  # rel
                raw_batch[-3][sent_idx] = raw_batch[-3][sent_idx][:num_masks + 1]  # pos
            prefix_mask[0] = True  # <root> is now [CLS]
            prefix_offset = [idx for idx, m in enumerate(prefix_mask) if m]
            batch_input_ids.append(input_ids)
            batch_input_mask.append(input_mask)
            batch_prefix_offset.append(prefix_offset)
            if mask_p:
                if shuffle:
                    size = int(np.ceil(mask_p * len(prefix_offset[1:])))  # never mask [CLS]
                    mask_offsets = np.random.choice(np.arange(1, len(prefix_offset)), size, replace=False)
                    for offset in sorted(mask_offsets):
                        assert 0 < offset < len(input_ids)
                        # mask_word = raw_batch[0][sent_idx][offset]
                        # mask_prefix = tokenizer.convert_ids_to_tokens([input_ids[prefix_offset[offset]]])[0]
                        # assert mask_word.startswith(mask_prefix) or mask_prefix.startswith(
                        #     mask_word) or mask_prefix == "'", \
                        #     f'word {mask_word} prefix {mask_prefix} not match'  # could vs couldn
                        # mask_offsets.append(input_ids[offset]) # subword token
                        # mask_offsets.append(offset)  # form token
                        input_ids[prefix_offset[offset]] = mask_token_id  # mask prefix
                        # whole word masking, mask the rest of the word
                        for i in range(prefix_offset[offset] + 1, len(input_ids) - 1):
                            if prefix_mask[i]:
                                break
                            input_ids[i] = mask_token_id

                    batch_masked_offsets.append(sorted(mask_offsets))
                else:
                    batch_masked_offsets.append([0])  # No masking in prediction

        batch_forms = tf.keras.preprocessing.sequence.pad_sequences(batch_forms, padding='post',
                                                                    value=self.form_vocab.safe_pad_token_idx,
                                                                    dtype='int64')
        batch_input_ids = tf.keras.preprocessing.sequence.pad_sequences(batch_input_ids, padding='post',
                                                                        value=pad_token_id,
                                                                        dtype='int64')
        batch_input_mask = tf.keras.preprocessing.sequence.pad_sequences(batch_input_mask, padding='post',
                                                                         value=0,
                                                                         dtype='int64')
        batch_prefix_offset = tf.keras.preprocessing.sequence.pad_sequences(batch_prefix_offset, padding='post',
                                                                            value=0,
                                                                            dtype='int64')
        batch_heads = tf.keras.preprocessing.sequence.pad_sequences(raw_batch[-2], padding='post',
                                                                    value=0,
                                                                    dtype='int64')
        batch_rels = tf.keras.preprocessing.sequence.pad_sequences(raw_batch[-1], padding='post',
                                                                   value=self.rel_vocab.safe_pad_token_idx,
                                                                   dtype='int64')
        if mask_p:
            batch_masked_offsets = tf.keras.preprocessing.sequence.pad_sequences(batch_masked_offsets, padding='post',
                                                                                 value=pad_token_id,
                                                                                 dtype='int64')
        feats = (tf.constant(batch_input_ids, dtype='int64'), tf.constant(batch_input_mask, dtype='int64'),
                 tf.constant(batch_prefix_offset))
        if use_pos:
            batch_pos = tf.keras.preprocessing.sequence.pad_sequences(raw_batch[1], padding='post',
                                                                      value=self.cpos_vocab.safe_pad_token_idx,
                                                                      dtype='int64')
            feats += (batch_pos,)
        yield (batch_forms, feats), \
              (batch_heads, batch_rels, batch_masked_offsets) if mask_p else (batch_heads, batch_rels)

    def len_of_sent(self, sent):
        # Transformer tokenizing
        config = self.transformer_config
        tokenizer = self.tokenizer
        xlnet = config_is(config, 'xlnet')
        roberta = config_is(config, 'roberta')
        pad_token = tokenizer.pad_token
        pad_token_id = tokenizer.convert_tokens_to_ids([pad_token])[0]
        cls_token = tokenizer.cls_token
        sep_token = tokenizer.sep_token
        max_seq_length = self.config.max_seq_length
        sent = sent[1:]  # remove <root> use [CLS] instead
        pad_label_idx = self.form_vocab.pad_idx
        sent = [x[0] for x in sent]
        sent = adjust_tokens_for_transformers(sent)
        input_ids, input_mask, segment_ids, prefix_mask = \
            convert_examples_to_features(sent,
                                         max_seq_length,
                                         tokenizer,
                                         cls_token_at_end=xlnet,
                                         # xlnet has a cls token at the end
                                         cls_token=cls_token,
                                         cls_token_segment_id=2 if xlnet else 0,
                                         sep_token=sep_token,
                                         sep_token_extra=roberta,
                                         # roberta uses an extra separator b/w pairs of sentences, cf. github.com/pytorch/fairseq/commit/1684e166e3da03f5b600dbb7855cb98ddfcd0805
                                         pad_on_left=xlnet,
                                         # pad on the left for xlnet
                                         pad_token_id=pad_token_id,
                                         pad_token_segment_id=4 if xlnet else 0,
                                         pad_token_label_id=pad_label_idx,
                                         do_padding=False)
        return len(input_ids)

    def samples_to_dataset(self, samples: Generator, map_x=None, map_y=None, batch_size=5000, shuffle=None, repeat=None,
                           drop_remainder=False, prefetch=1, cache=True) -> tf.data.Dataset:
        if shuffle:
            return CoNLL_DEP_Transform.samples_to_dataset(self, samples, map_x, map_y, batch_size, shuffle, repeat,
                                                          drop_remainder, prefetch, cache)

        def generator():
            # custom bucketing, load corpus into memory
            corpus = list(x for x in (samples() if callable(samples) else samples))
            n_tokens = 0
            batch = []
            for idx, sent in enumerate(corpus):
                sent_len = self.len_of_sent(sent)
                if n_tokens + sent_len > batch_size and batch:
                    yield from self.batched_inputs_to_batches(corpus, batch, shuffle)
                    n_tokens = 0
                    batch = []
                n_tokens += sent_len
                batch.append(idx)
            if batch:
                yield from self.batched_inputs_to_batches(corpus, batch, shuffle)

        # debug for transformer
        # next(generator())
        return Transform.samples_to_dataset(self, generator, False, False, 0, False, repeat, drop_remainder, prefetch,
                                            cache)

    def Y_to_outputs(self, Y: Union[tf.Tensor, Tuple[tf.Tensor]], gold=False, inputs=None, X=None) -> Iterable:
        if self.graph:
            ys = CoNLL_SDP_Transform.Y_to_outputs(self, Y, gold, inputs, X)
            ys = [[([t[0] for t in l], [t[1] for t in l]) for l in y] for y in ys]
            return ys
        return super().Y_to_outputs(Y, gold, inputs, X)


class CoNLL_SDP_Transform(CoNLLTransform):

    def __init__(self, config: SerializableDict = None, map_x=True, map_y=True, lower=True, n_buckets=32, min_freq=2,
                 use_pos=True, **kwargs) -> None:
        super().__init__(config, map_x, map_y, lower, n_buckets, min_freq, use_pos, **kwargs)
        self.orphan_relation = ROOT

    def lock_vocabs(self):
        super().lock_vocabs()
        # heuristic to find the orphan relation
        self._find_orphan_relation()

    def _find_orphan_relation(self):
        for rel in self.rel_vocab.idx_to_token:
            if 'root' in rel.lower():
                self.orphan_relation = rel
                break

    def file_to_inputs(self, filepath: str, gold=True):
        assert gold, 'only support gold file for now'
        use_pos = self.use_pos
        conllu = filepath.endswith('.conllu')
        enhanced_only = self.config.get('enhanced_only', None)
        for i, sent in enumerate(read_conll(filepath)):
            parsed_sent = []
            if conllu:
                for cell in sent:
                    ID = cell[0]
                    form = cell[1]
                    cpos = cell[3]
                    head = cell[6]
                    deprel = cell[7]
                    deps = cell[8]
                    deps = [x.split(':', 1) for x in deps.split('|')]
                    heads = [int(x[0]) for x in deps if x[0].isdigit()]
                    rels = [x[1] for x in deps if x[0].isdigit()]
                    if enhanced_only:
                        if head in heads:
                            offset = heads.index(head)
                            heads.pop(offset)
                            rels.pop(offset)
                    else:
                        if head not in heads:
                            heads.append(head)
                            rels.append(deprel)
                    parsed_sent.append([form, cpos, heads, rels] if use_pos else [form, heads, rels])
            else:
                prev_cells = None
                heads = []
                rels = []
                for j, cell in enumerate(sent):
                    ID = cell[0]
                    form = cell[1]
                    cpos = cell[3]
                    head = cell[6]
                    deprel = cell[7]
                    if prev_cells and ID != prev_cells[0]:  # found end of token
                        parsed_sent.append(
                            [prev_cells[1], prev_cells[2], heads, rels] if use_pos else [prev_cells[1], heads, rels])
                        heads = []
                        rels = []
                    heads.append(head)
                    rels.append(deprel)
                    prev_cells = [ID, form, cpos, head, deprel] if use_pos else [ID, form, head, deprel]
                parsed_sent.append(
                    [prev_cells[1], prev_cells[2], heads, rels] if use_pos else [prev_cells[1], heads, rels])
            yield parsed_sent

    def fit(self, trn_path: str, **kwargs) -> int:
        self.form_vocab = VocabTF()
        self.form_vocab.add(ROOT)  # make root the 2ed elements while 0th is pad, 1st is unk
        if self.use_pos:
            self.cpos_vocab = VocabTF(pad_token=None, unk_token=None)
        self.rel_vocab = VocabTF(pad_token=None, unk_token=None)
        num_samples = 0
        counter = Counter()
        for sent in self.file_to_samples(trn_path, gold=True):
            num_samples += 1
            for idx, cell in enumerate(sent):
                if len(cell) == 4:
                    form, cpos, head, deprel = cell
                elif len(cell) == 3:
                    if self.use_pos:
                        form, cpos = cell[0]
                    else:
                        form = cell[0]
                    head, deprel = cell[1:]
                else:
                    raise ValueError('Unknown data arrangement')
                if idx == 0:
                    root = form
                else:
                    counter[form] += 1
                if self.use_pos:
                    self.cpos_vocab.add(cpos)
                self.rel_vocab.update(deprel)

        for token in [token for token, freq in counter.items() if freq >= self.config.min_freq]:
            self.form_vocab.add(token)
        return num_samples

    def inputs_to_samples(self, inputs, gold=False):
        use_pos = self.use_pos
        for sent in inputs:
            sample = []
            for i, cell in enumerate(sent):
                if isinstance(cell, tuple):
                    cell = list(cell)
                elif isinstance(cell, str):
                    cell = [cell]
                if self.config['lower']:
                    cell[0] = cell[0].lower()
                if not gold:
                    cell += [[0], [self.rel_vocab.safe_pad_token]]
                sample.append(cell)
            # insert root word with arbitrary fields, anyway it will be masked
            if use_pos:
                form, cpos, head, deprel = sample[0]
                sample.insert(0, [self.bos, self.bos, [0], deprel])
            else:
                form, head, deprel = sample[0]
                sample.insert(0, [self.bos, [0], deprel])
            yield sample

    def batched_inputs_to_batches(self, corpus, indices, shuffle):
        use_pos = self.use_pos
        raw_batch = [[], [], [], []] if use_pos else [[], [], []]
        max_len = len(max([corpus[i] for i in indices], key=len))
        for idx in indices:
            arc = np.zeros((max_len, max_len), dtype=bool)
            rel = np.zeros((max_len, max_len), dtype=np.int64)
            for b in raw_batch[:2]:
                b.append([])
            for m, cells in enumerate(corpus[idx]):
                if use_pos:
                    for b, c, v in zip(raw_batch, cells,
                                       [self.form_vocab, self.cpos_vocab]):
                        b[-1].append(v.get_idx_without_add(c))
                else:
                    for b, c, v in zip(raw_batch, cells,
                                       [self.form_vocab]):
                        b[-1].append(v.get_idx_without_add(c))
                for n, r in zip(cells[-2], cells[-1]):
                    arc[m, n] = True
                    rid = self.rel_vocab.get_idx_without_add(r)
                    if rid is None:
                        logger.warning(f'Relation OOV: {r} not exists in train')
                        continue
                    rel[m, n] = rid
            raw_batch[-2].append(arc)
            raw_batch[-1].append(rel)
        batch = []
        for b, v in zip(raw_batch, [self.form_vocab, self.cpos_vocab]):
            b = tf.keras.preprocessing.sequence.pad_sequences(b, padding='post',
                                                              value=v.safe_pad_token_idx,
                                                              dtype='int64')
            batch.append(b)
        batch += raw_batch[2:]
        assert len(batch) == 4
        yield (batch[0], batch[1]), (batch[2], batch[3])

    def create_types_shapes_values(self) -> Tuple[Tuple, Tuple, Tuple]:
        types = (tf.int64, tf.int64), (tf.bool, tf.int64)
        shapes = ([None, None], [None, None]), ([None, None, None], [None, None, None])
        values = (self.form_vocab.safe_pad_token_idx, self.cpos_vocab.safe_pad_token_idx), (
            False, self.rel_vocab.safe_pad_token_idx)
        return types, shapes, values

    def Y_to_outputs(self, Y: Union[tf.Tensor, Tuple[tf.Tensor]], gold=False, inputs=None, X=None) -> Iterable:
        arc_preds, rel_preds, mask = Y
        sents = []

        for arc_sent, rel_sent, length in zip(arc_preds, rel_preds,
                                              tf.math.count_nonzero(mask, axis=-1)):
            sent = []
            for arc, rel in zip(tolist(arc_sent[1:, 1:]), tolist(rel_sent[1:, 1:])):
                ar = []
                for idx, (a, r) in enumerate(zip(arc, rel)):
                    if a:
                        ar.append((idx + 1, self.rel_vocab.idx_to_token[r]))
                if not ar:
                    # orphan
                    ar.append((0, self.orphan_relation))
                sent.append(ar)
            sents.append(sent)

        return sents

    def XY_to_inputs_outputs(self, X: Union[tf.Tensor, Tuple[tf.Tensor]], Y: Union[tf.Tensor, Tuple[tf.Tensor]],
                             gold=False, inputs=None, conll=True) -> Iterable:
        (words, feats, mask), (arc_preds, rel_preds) = X, Y
        xs = inputs
        ys = self.Y_to_outputs((arc_preds, rel_preds, mask))
        sents = []
        for x, y in zip(xs, ys):
            sent = CoNLLSentence()
            for idx, ((form, cpos), pred) in enumerate(zip(x, y)):
                head = [p[0] for p in pred]
                deprel = [p[1] for p in pred]
                if conll:
                    sent.append(CoNLLWord(id=idx + 1, form=form, cpos=cpos, head=head, deprel=deprel))
                else:
                    sent.append([head, deprel])
            sents.append(sent)
        return sents