# -*- coding: utf-8 -*-

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import json
import os
import random

from absl import flags
import absl.logging as _logging  # pylint: disable=unused-import

import numpy as np


import tensorflow as tf

from prepro_utils import preprocess_text, encode_ids
import sentencepiece as spm


special_symbols = {
    '<unk>': 0,
    '<s>': 1,
    '</s>': 2,
    '<cls>': 3,
    '<sep>': 4,
    '<pad>': 5,
    '<mask>': 6,
    '<eod>': 7,
    '<eop>': 8,
}

VOCAB_SIZE = 32000
UNK_ID = special_symbols['<unk>']
CLS_ID = special_symbols['<cls>']
SEP_ID = special_symbols['<sep>']
MASK_ID = special_symbols['<mask>']
EOD_ID = special_symbols['<eod>']


def _int64_feature(values):
    return tf.train.Feature(int64_list = tf.train.Int64List(value = values))


def _float_feature(values):
    return tf.train.Feature(float_list = tf.train.FloatList(value = values))


def format_filename(
    prefix,
    bsz_per_host,
    seq_len,
    bi_data,
    suffix,
    mask_alpha = 5,
    mask_beta = 1,
    reuse_len = None,
    uncased = False,
    fixed_num_predict = None,
):
    """docs."""
    if reuse_len is None:
        reuse_len_str = ''
    else:
        reuse_len_str = 'reuse-{}.'.format(reuse_len)
    if not uncased:
        uncased_str = ''
    else:
        uncased_str = 'uncased.'
    if bi_data:
        bi_data_str = 'bi'
    else:
        bi_data_str = 'uni'
    if fixed_num_predict is not None:
        fnp_str = 'fnp-{}.'.format(fixed_num_predict)
    else:
        fnp_str = ''

    file_name = '{}.bsz-{}.seqlen-{}.{}{}{}.alpha-{}.beta-{}.{}{}'.format(
        prefix,
        bsz_per_host,
        seq_len,
        reuse_len_str,
        uncased_str,
        bi_data_str,
        mask_alpha,
        mask_beta,
        fnp_str,
        suffix,
    )

    return file_name


def _create_data(idx, input_paths):
    # Load sentence-piece model
    sp = spm.SentencePieceProcessor()
    sp.Load(FLAGS.sp_path)

    input_shards = []
    total_line_cnt = 0
    for input_path in input_paths:
        input_data, sent_ids = [], []
        sent_id, line_cnt = True, 0
        tf.logging.info('Processing %s', input_path)
        for line in tf.gfile.Open(input_path):
            if line_cnt % 100000 == 0:
                tf.logging.info('Loading line %d', line_cnt)
            line_cnt += 1

            if not line.strip():
                if FLAGS.use_eod:
                    sent_id = not sent_id
                    cur_sent = [EOD_ID]
                else:
                    continue
            else:
                if FLAGS.from_raw_text:
                    cur_sent = preprocess_text(
                        line.strip(), lower = FLAGS.uncased
                    )
                    cur_sent = encode_ids(sp, cur_sent)
                else:
                    cur_sent = list(map(int, line.strip().split()))

            input_data.extend(cur_sent)
            sent_ids.extend([sent_id] * len(cur_sent))
            sent_id = not sent_id

        tf.logging.info('Finish with line %d', line_cnt)
        if line_cnt == 0:
            continue

        input_data = np.array(input_data, dtype = np.int64)
        sent_ids = np.array(sent_ids, dtype = np.bool)

        total_line_cnt += line_cnt
        input_shards.append((input_data, sent_ids))

    tf.logging.info('[Task %d] Total number line: %d', idx, total_line_cnt)

    tfrecord_dir = os.path.join(FLAGS.save_dir, 'tfrecords')

    filenames, num_batch = [], 0

    # Randomly shuffle input shards (with a fixed but distinct random seed)
    np.random.seed(100 * FLAGS.task + FLAGS.pass_id)

    perm_indices = np.random.permutation(len(input_shards))
    tf.logging.info(
        'Using perm indices %s for pass %d',
        perm_indices.tolist(),
        FLAGS.pass_id,
    )

    input_data_list, sent_ids_list = [], []
    prev_sent_id = None
    for perm_idx in perm_indices:
        input_data, sent_ids = input_shards[perm_idx]
        # make sure the `send_ids[0] == not prev_sent_id`
        if prev_sent_id is not None and sent_ids[0] == prev_sent_id:
            sent_ids = np.logical_not(sent_ids)

        # append to temporary list
        input_data_list.append(input_data)
        sent_ids_list.append(sent_ids)

        # update `prev_sent_id`
        prev_sent_id = sent_ids[-1]

    input_data = np.concatenate(input_data_list)
    sent_ids = np.concatenate(sent_ids_list)

    file_name, cur_num_batch = create_tfrecords(
        save_dir = tfrecord_dir,
        basename = '{}-{}-{}'.format(FLAGS.split, idx, FLAGS.pass_id),
        data = [input_data, sent_ids],
        bsz_per_host = FLAGS.bsz_per_host,
        seq_len = FLAGS.seq_len,
        bi_data = FLAGS.bi_data,
        sp = sp,
    )

    filenames.append(file_name)
    num_batch += cur_num_batch

    record_info = {'filenames': filenames, 'num_batch': num_batch}

    return record_info


def create_data(_):
    # Validate FLAGS
    assert FLAGS.bsz_per_host % FLAGS.num_core_per_host == 0
    if not FLAGS.use_tpu:
        FLAGS.num_core_per_host = 1  # forced to be one

    # Make workdirs
    if not tf.gfile.Exists(FLAGS.save_dir):
        tf.gfile.MakeDirs(FLAGS.save_dir)

    tfrecord_dir = os.path.join(FLAGS.save_dir, 'tfrecords')
    if not tf.gfile.Exists(tfrecord_dir):
        tf.gfile.MakeDirs(tfrecord_dir)

    # Create and dump corpus_info from task 0
    if FLAGS.task == 0:
        corpus_info = {
            'vocab_size': VOCAB_SIZE,
            'bsz_per_host': FLAGS.bsz_per_host,
            'num_core_per_host': FLAGS.num_core_per_host,
            'seq_len': FLAGS.seq_len,
            'reuse_len': FLAGS.reuse_len,
            'uncased': FLAGS.uncased,
            'bi_data': FLAGS.bi_data,
            'mask_alpha': FLAGS.mask_alpha,
            'mask_beta': FLAGS.mask_beta,
            'num_predict': FLAGS.num_predict,
            'use_eod': FLAGS.use_eod,
            'sp_path': FLAGS.sp_path,
            'input_glob': FLAGS.input_glob,
        }
        corpus_info_path = os.path.join(FLAGS.save_dir, 'corpus_info.json')
        with tf.gfile.Open(corpus_info_path, 'w') as fp:
            json.dump(corpus_info, fp)

    # Interleavely split the work into FLAGS.num_task splits
    file_paths = sorted(tf.gfile.Glob(FLAGS.input_glob))
    tf.logging.info('Use glob: %s', FLAGS.input_glob)
    tf.logging.info('Find %d files: %s', len(file_paths), file_paths)

    task_file_paths = file_paths[FLAGS.task :: FLAGS.num_task]
    if not task_file_paths:
        tf.logging.info('Exit: task %d has no file to process.', FLAGS.task)
        return

    tf.logging.info(
        'Task %d process %d files: %s',
        FLAGS.task,
        len(task_file_paths),
        task_file_paths,
    )
    record_info = _create_data(FLAGS.task, task_file_paths)

    record_prefix = 'record_info-{}-{}-{}'.format(
        FLAGS.split, FLAGS.task, FLAGS.pass_id
    )
    record_name = format_filename(
        prefix = record_prefix,
        bsz_per_host = FLAGS.bsz_per_host,
        seq_len = FLAGS.seq_len,
        mask_alpha = FLAGS.mask_alpha,
        mask_beta = FLAGS.mask_beta,
        reuse_len = FLAGS.reuse_len,
        bi_data = FLAGS.bi_data,
        suffix = 'json',
        uncased = FLAGS.uncased,
        fixed_num_predict = FLAGS.num_predict,
    )
    record_info_path = os.path.join(tfrecord_dir, record_name)

    with tf.gfile.Open(record_info_path, 'w') as fp:
        json.dump(record_info, fp)


def batchify(data, bsz_per_host, sent_ids = None):
    num_step = len(data) // bsz_per_host
    data = data[: bsz_per_host * num_step]
    data = data.reshape(bsz_per_host, num_step)
    if sent_ids is not None:
        sent_ids = sent_ids[: bsz_per_host * num_step]
        sent_ids = sent_ids.reshape(bsz_per_host, num_step)

    if sent_ids is not None:
        return data, sent_ids
    return data


def _split_a_and_b(data, sent_ids, begin_idx, tot_len, extend_target = False):
    """Split two segments from `data` starting from the index `begin_idx`."""

    data_len = data.shape[0]
    if begin_idx + tot_len >= data_len:
        tf.logging.info(
            '[_split_a_and_b] returns None: '
            'begin_idx %d + tot_len %d >= data_len %d',
            begin_idx,
            tot_len,
            data_len,
        )
        return None

    end_idx = begin_idx + 1
    cut_points = []
    while end_idx < data_len:
        if sent_ids[end_idx] != sent_ids[end_idx - 1]:
            if end_idx - begin_idx >= tot_len:
                break
            cut_points.append(end_idx)
        end_idx += 1

    a_begin = begin_idx
    if len(cut_points) == 0 or random.random() < 0.5:
        label = 0
        if len(cut_points) == 0:
            a_end = end_idx
        else:
            a_end = random.choice(cut_points)

        b_len = max(1, tot_len - (a_end - a_begin))
        # (zihang): `data_len - 1` to account for extend_target
        b_begin = random.randint(0, data_len - 1 - b_len)
        b_end = b_begin + b_len
        while b_begin > 0 and sent_ids[b_begin - 1] == sent_ids[b_begin]:
            b_begin -= 1
        # (zihang): `data_len - 1` to account for extend_target
        while b_end < data_len - 1 and sent_ids[b_end - 1] == sent_ids[b_end]:
            b_end += 1

        new_begin = a_end
    else:
        label = 1
        a_end = random.choice(cut_points)
        b_begin = a_end
        b_end = end_idx

        new_begin = b_end

    while a_end - a_begin + b_end - b_begin > tot_len:
        if a_end - a_begin > b_end - b_begin:
            # delete the right side only for the LM objective
            a_end -= 1
        else:
            b_end -= 1

    ret = [data[a_begin:a_end], data[b_begin:b_end], label, new_begin]

    if extend_target:
        if a_end >= data_len or b_end >= data_len:
            tf.logging.info(
                '[_split_a_and_b] returns None: '
                'a_end %d or b_end %d >= data_len %d',
                a_end,
                b_end,
                data_len,
            )
            return None
        a_target = data[a_begin + 1 : a_end + 1]
        b_target = data[b_begin : b_end + 1]
        ret.extend([a_target, b_target])

    return ret


def _is_start_piece(piece):
    special_pieces = set(list('!"#$%&"()*+,-./:;?@[\\]^_`{|}~'))
    if (
        piece.startswith('▁')
        or piece.startswith('<')
        or piece in special_pieces
    ):
        return True
    else:
        return False


def _sample_mask(
    sp, seg, reverse = False, max_gram = 5, goal_num_predict = None
):
    """Sample `goal_num_predict` tokens for partial prediction.
  About `mask_beta` tokens are chosen in a context of `mask_alpha` tokens."""

    seg_len = len(seg)
    mask = np.array([False] * seg_len, dtype = np.bool)

    num_predict = 0

    ngrams = np.arange(1, max_gram + 1, dtype = np.int64)
    pvals = 1.0 / np.arange(1, max_gram + 1)
    pvals /= pvals.sum(keepdims = True)

    if reverse:
        seg = np.flip(seg, 0)

    cur_len = 0
    while cur_len < seg_len:
        if goal_num_predict is not None and num_predict >= goal_num_predict:
            break

        n = np.random.choice(ngrams, p = pvals)
        if goal_num_predict is not None:
            n = min(n, goal_num_predict - num_predict)
        ctx_size = (n * FLAGS.mask_alpha) // FLAGS.mask_beta
        l_ctx = np.random.choice(ctx_size)
        r_ctx = ctx_size - l_ctx

        # Find the start position of a complete token
        beg = cur_len + l_ctx
        while beg < seg_len and not _is_start_piece(
            sp.IdToPiece(seg[beg].item())
        ):
            beg += 1
        if beg >= seg_len:
            break

        # Find the end position of the n-gram (start pos of the n+1-th gram)
        end = beg + 1
        cnt_ngram = 1
        while end < seg_len:
            if _is_start_piece(sp.IdToPiece(seg[beg].item())):
                cnt_ngram += 1
                if cnt_ngram > n:
                    break
            end += 1
        if end >= seg_len:
            break

        # Update
        mask[beg:end] = True
        num_predict += end - beg

        cur_len = end + r_ctx

    while goal_num_predict is not None and num_predict < goal_num_predict:
        i = np.random.randint(seg_len)
        if not mask[i]:
            mask[i] = True
            num_predict += 1

    if reverse:
        mask = np.flip(mask, 0)

    return mask


def create_tfrecords(
    save_dir, basename, data, bsz_per_host, seq_len, bi_data, sp
):
    data, sent_ids = data[0], data[1]

    num_core = FLAGS.num_core_per_host
    bsz_per_core = bsz_per_host // num_core

    if bi_data:
        assert bsz_per_host % (2 * FLAGS.num_core_per_host) == 0
        fwd_data, fwd_sent_ids = batchify(data, bsz_per_host // 2, sent_ids)

        fwd_data = fwd_data.reshape(num_core, 1, bsz_per_core // 2, -1)
        fwd_sent_ids = fwd_sent_ids.reshape(num_core, 1, bsz_per_core // 2, -1)

        bwd_data = fwd_data[:, :, :, ::-1]
        bwd_sent_ids = fwd_sent_ids[:, :, :, ::-1]

        data = np.concatenate([fwd_data, bwd_data], 1).reshape(bsz_per_host, -1)
        sent_ids = np.concatenate([fwd_sent_ids, bwd_sent_ids], 1).reshape(
            bsz_per_host, -1
        )
    else:
        data, sent_ids = batchify(data, bsz_per_host, sent_ids)

    tf.logging.info('Raw data shape %s.', data.shape)

    file_name = format_filename(
        prefix = basename,
        bsz_per_host = bsz_per_host,
        seq_len = seq_len,
        bi_data = bi_data,
        suffix = 'tfrecords',
        mask_alpha = FLAGS.mask_alpha,
        mask_beta = FLAGS.mask_beta,
        reuse_len = FLAGS.reuse_len,
        uncased = FLAGS.uncased,
        fixed_num_predict = FLAGS.num_predict,
    )
    save_path = os.path.join(save_dir, file_name)
    record_writer = tf.python_io.TFRecordWriter(save_path)
    tf.logging.info('Start writing %s.', save_path)

    num_batch = 0
    reuse_len = FLAGS.reuse_len

    # [sep] x 2 + [cls]
    assert reuse_len < seq_len - 3

    data_len = data.shape[1]
    sep_array = np.array([SEP_ID], dtype = np.int64)
    cls_array = np.array([CLS_ID], dtype = np.int64)

    i = 0
    while i + seq_len <= data_len:
        if num_batch % 500 == 0:
            tf.logging.info('Processing batch %d', num_batch)

        all_ok = True
        features = []
        for idx in range(bsz_per_host):
            inp = data[idx, i : i + reuse_len]
            tgt = data[idx, i + 1 : i + reuse_len + 1]

            results = _split_a_and_b(
                data[idx],
                sent_ids[idx],
                begin_idx = i + reuse_len,
                tot_len = seq_len - reuse_len - 3,
                extend_target = True,
            )
            if results is None:
                tf.logging.info('Break out with seq idx %d', i)
                all_ok = False
                break

            # unpack the results
            (a_data, b_data, label, _, a_target, b_target) = tuple(results)

            # sample ngram spans to predict
            reverse = bi_data and (idx // (bsz_per_core // 2)) % 2 == 1
            if FLAGS.num_predict is None:
                num_predict_0 = num_predict_1 = None
            else:
                num_predict_1 = FLAGS.num_predict // 2
                num_predict_0 = FLAGS.num_predict - num_predict_1
            mask_0 = _sample_mask(
                sp, inp, reverse = reverse, goal_num_predict = num_predict_0
            )
            mask_1 = _sample_mask(
                sp,
                np.concatenate(
                    [a_data, sep_array, b_data, sep_array, cls_array]
                ),
                reverse = reverse,
                goal_num_predict = num_predict_1,
            )

            # concatenate data
            cat_data = np.concatenate(
                [inp, a_data, sep_array, b_data, sep_array, cls_array]
            )
            seg_id = (
                [0] * (reuse_len + a_data.shape[0])
                + [0]
                + [1] * b_data.shape[0]
                + [1]
                + [2]
            )
            assert cat_data.shape[0] == seq_len
            assert mask_0.shape[0] == seq_len // 2
            assert mask_1.shape[0] == seq_len // 2

            # the last two CLS's are not used, just for padding purposes
            tgt = np.concatenate(
                [tgt, a_target, b_target, cls_array, cls_array]
            )
            assert tgt.shape[0] == seq_len

            is_masked = np.concatenate([mask_0, mask_1], 0)
            if FLAGS.num_predict is not None:
                assert np.sum(is_masked) == FLAGS.num_predict

            feature = {
                'input': _int64_feature(cat_data),
                'is_masked': _int64_feature(is_masked),
                'target': _int64_feature(tgt),
                'seg_id': _int64_feature(seg_id),
                'label': _int64_feature([label]),
            }
            features.append(feature)

        if all_ok:
            assert len(features) == bsz_per_host
            for feature in features:
                example = tf.train.Example(
                    features = tf.train.Features(feature = feature)
                )
                record_writer.write(example.SerializeToString())
            num_batch += 1
        else:
            break

        i += reuse_len

    record_writer.close()
    tf.logging.info('Done writing %s. Num of batches: %d', save_path, num_batch)

    return save_path, num_batch


################
# get_input_fn #
################
def _convert_example(example, use_bfloat16):
    """Cast int64 into int32 and float32 to bfloat16 if use_bfloat16."""
    for key in list(example.keys()):
        val = example[key]
        if tf.keras.backend.is_sparse(val):
            val = tf.sparse.to_dense(val)
        if val.dtype == tf.int64:
            val = tf.cast(val, tf.int32)
        if use_bfloat16 and val.dtype == tf.float32:
            val = tf.cast(val, tf.bfloat16)

        example[key] = val


def parse_files_to_dataset(
    parser,
    file_names,
    split,
    num_batch,
    num_hosts,
    host_id,
    num_core_per_host,
    bsz_per_core,
):
    # list of file pathes
    num_files = len(file_names)
    num_files_per_host = num_files // num_hosts
    my_start_file_id = host_id * num_files_per_host
    my_end_file_id = (host_id + 1) * num_files_per_host
    if host_id == num_hosts - 1:
        my_end_file_id = num_files
    file_paths = file_names[my_start_file_id:my_end_file_id]
    tf.logging.info('Host %d handles %d files', host_id, len(file_paths))

    assert split == 'train'
    dataset = tf.data.Dataset.from_tensor_slices(file_paths)

    # file-level shuffle
    if len(file_paths) > 1:
        dataset = dataset.shuffle(len(file_paths))

    # Note: we cannot perform sample-level shuffle here because this will violate
    # the consecutive requirement of data stream.
    dataset = tf.data.TFRecordDataset(dataset)

    # (zihang): since we are doing online preprocessing, the parsed result of
    # the same input at each time will be different. Thus, cache processed data
    # is not helpful. It will use a lot of memory and lead to contrainer OOM.
    # So, change to cache non-parsed raw data instead.
    dataset = dataset.map(parser).repeat()
    dataset = dataset.batch(bsz_per_core, drop_remainder = True)
    dataset = dataset.prefetch(num_core_per_host * bsz_per_core)

    return dataset


def _local_perm(inputs, targets, is_masked, perm_size, seq_len):
    """
  Sample a permutation of the factorization order, and create an
  attention mask accordingly.

  Args:
    inputs: int64 Tensor in shape [seq_len], input ids.
    targets: int64 Tensor in shape [seq_len], target ids.
    is_masked: bool Tensor in shape [seq_len]. True means being selected
      for partial prediction.
    perm_size: the length of longest permutation. Could be set to be reuse_len.
      Should not be larger than reuse_len or there will be data leaks.
    seq_len: int, sequence length.
  """

    # Generate permutation indices
    index = tf.range(seq_len, dtype = tf.int64)
    index = tf.transpose(tf.reshape(index, [-1, perm_size]))
    index = tf.random_shuffle(index)
    index = tf.reshape(tf.transpose(index), [-1])

    # `perm_mask` and `target_mask`
    # non-functional tokens
    non_func_tokens = tf.logical_not(
        tf.logical_or(tf.equal(inputs, SEP_ID), tf.equal(inputs, CLS_ID))
    )

    non_mask_tokens = tf.logical_and(tf.logical_not(is_masked), non_func_tokens)
    masked_or_func_tokens = tf.logical_not(non_mask_tokens)

    # Set the permutation indices of non-masked (& non-funcional) tokens to the
    # smallest index (-1):
    # (1) they can be seen by all other positions
    # (2) they cannot see masked positions, so there won"t be information leak
    smallest_index = -tf.ones([seq_len], dtype = tf.int64)
    rev_index = tf.where(non_mask_tokens, smallest_index, index)

    # Create `target_mask`: non-funcional and maksed tokens
    # 1: use mask as input and have loss
    # 0: use token (or [SEP], [CLS]) as input and do not have loss
    target_tokens = tf.logical_and(masked_or_func_tokens, non_func_tokens)
    target_mask = tf.cast(target_tokens, tf.float32)

    # Create `perm_mask`
    # `target_tokens` cannot see themselves
    self_rev_index = tf.where(target_tokens, rev_index, rev_index + 1)

    # 1: cannot attend if i <= j and j is not non-masked (masked_or_func_tokens)
    # 0: can attend if i > j or j is non-masked
    perm_mask = tf.logical_and(
        self_rev_index[:, None] <= rev_index[None, :], masked_or_func_tokens
    )
    perm_mask = tf.cast(perm_mask, tf.float32)

    # new target: [next token] for LM and [curr token] (self) for PLM
    new_targets = tf.concat([inputs[0:1], targets[:-1]], axis = 0)

    # construct inputs_k
    inputs_k = inputs

    # construct inputs_q
    inputs_q = target_mask

    return perm_mask, new_targets, target_mask, inputs_k, inputs_q


def get_dataset(
    params,
    num_hosts,
    num_core_per_host,
    split,
    file_names,
    num_batch,
    seq_len,
    reuse_len,
    perm_size,
    mask_alpha,
    mask_beta,
    use_bfloat16 = False,
    num_predict = None,
):

    bsz_per_core = params['batch_size']
    if num_hosts > 1:
        host_id = params['context'].current_host
    else:
        host_id = 0

        #### Function used to parse tfrecord

    def parser(record):
        """function used to parse tfrecord."""

        record_spec = {
            'input': tf.FixedLenFeature([seq_len], tf.int64),
            'target': tf.FixedLenFeature([seq_len], tf.int64),
            'seg_id': tf.FixedLenFeature([seq_len], tf.int64),
            'label': tf.FixedLenFeature([1], tf.int64),
            'is_masked': tf.FixedLenFeature([seq_len], tf.int64),
        }

        # retrieve serialized example
        example = tf.parse_single_example(
            serialized = record, features = record_spec
        )

        inputs = example.pop('input')
        target = example.pop('target')
        is_masked = tf.cast(example.pop('is_masked'), tf.bool)

        non_reuse_len = seq_len - reuse_len
        assert perm_size <= reuse_len and perm_size <= non_reuse_len

        perm_mask_0, target_0, target_mask_0, input_k_0, input_q_0 = _local_perm(
            inputs[:reuse_len],
            target[:reuse_len],
            is_masked[:reuse_len],
            perm_size,
            reuse_len,
        )

        perm_mask_1, target_1, target_mask_1, input_k_1, input_q_1 = _local_perm(
            inputs[reuse_len:],
            target[reuse_len:],
            is_masked[reuse_len:],
            perm_size,
            non_reuse_len,
        )

        perm_mask_0 = tf.concat(
            [perm_mask_0, tf.ones([reuse_len, non_reuse_len])], axis = 1
        )
        perm_mask_1 = tf.concat(
            [tf.zeros([non_reuse_len, reuse_len]), perm_mask_1], axis = 1
        )
        perm_mask = tf.concat([perm_mask_0, perm_mask_1], axis = 0)
        target = tf.concat([target_0, target_1], axis = 0)
        target_mask = tf.concat([target_mask_0, target_mask_1], axis = 0)
        input_k = tf.concat([input_k_0, input_k_1], axis = 0)
        input_q = tf.concat([input_q_0, input_q_1], axis = 0)

        if num_predict is not None:
            indices = tf.range(seq_len, dtype = tf.int64)
            bool_target_mask = tf.cast(target_mask, tf.bool)
            indices = tf.boolean_mask(indices, bool_target_mask)

            ##### extra padding due to CLS/SEP introduced after prepro
            actual_num_predict = tf.shape(indices)[0]
            pad_len = num_predict - actual_num_predict

            ##### target_mapping
            target_mapping = tf.one_hot(indices, seq_len, dtype = tf.float32)
            paddings = tf.zeros(
                [pad_len, seq_len], dtype = target_mapping.dtype
            )
            target_mapping = tf.concat([target_mapping, paddings], axis = 0)
            example['target_mapping'] = tf.reshape(
                target_mapping, [num_predict, seq_len]
            )

            ##### target
            target = tf.boolean_mask(target, bool_target_mask)
            paddings = tf.zeros([pad_len], dtype = target.dtype)
            target = tf.concat([target, paddings], axis = 0)
            example['target'] = tf.reshape(target, [num_predict])

            ##### target mask
            target_mask = tf.concat(
                [
                    tf.ones([actual_num_predict], dtype = tf.float32),
                    tf.zeros([pad_len], dtype = tf.float32),
                ],
                axis = 0,
            )
            example['target_mask'] = tf.reshape(target_mask, [num_predict])
        else:
            example['target'] = tf.reshape(target, [seq_len])
            example['target_mask'] = tf.reshape(target_mask, [seq_len])

        # reshape back to fixed shape
        example['perm_mask'] = tf.reshape(perm_mask, [seq_len, seq_len])
        example['input_k'] = tf.reshape(input_k, [seq_len])
        example['input_q'] = tf.reshape(input_q, [seq_len])

        _convert_example(example, use_bfloat16)

        for k, v in example.items():
            tf.logging.info('%s: %s', k, v)

        return example

    # Get dataset
    dataset = parse_files_to_dataset(
        parser = parser,
        file_names = file_names,
        split = split,
        num_batch = num_batch,
        num_hosts = num_hosts,
        host_id = host_id,
        num_core_per_host = num_core_per_host,
        bsz_per_core = bsz_per_core,
    )

    return dataset


def get_input_fn(
    tfrecord_dir,
    split,
    bsz_per_host,
    seq_len,
    reuse_len,
    bi_data,
    num_hosts = 1,
    num_core_per_host = 1,
    perm_size = None,
    mask_alpha = None,
    mask_beta = None,
    uncased = False,
    num_passes = None,
    use_bfloat16 = False,
    num_predict = None,
):

    # Merge all record infos into a single one
    record_glob_base = format_filename(
        prefix = 'record_info-{}-*'.format(split),
        bsz_per_host = bsz_per_host,
        seq_len = seq_len,
        bi_data = bi_data,
        suffix = 'json',
        mask_alpha = mask_alpha,
        mask_beta = mask_beta,
        reuse_len = reuse_len,
        uncased = uncased,
        fixed_num_predict = num_predict,
    )

    record_info = {'num_batch': 0, 'filenames': []}

    tfrecord_dirs = tfrecord_dir.split(',')
    tf.logging.info('Use the following tfrecord dirs: %s', tfrecord_dirs)

    for idx, record_dir in enumerate(tfrecord_dirs):
        record_glob = os.path.join(record_dir, record_glob_base)
        tf.logging.info('[%d] Record glob: %s', idx, record_glob)

        record_paths = sorted(tf.gfile.Glob(record_glob))
        tf.logging.info(
            '[%d] Num of record info path: %d', idx, len(record_paths)
        )

        cur_record_info = {'num_batch': 0, 'filenames': []}

        for record_info_path in record_paths:
            if num_passes is not None:
                record_info_name = os.path.basename(record_info_path)
                fields = record_info_name.split('.')[0].split('-')
                pass_id = int(fields[-1])
                if len(fields) == 5 and pass_id >= num_passes:
                    tf.logging.info(
                        'Skip pass %d: %s', pass_id, record_info_name
                    )
                    continue

            with tf.gfile.Open(record_info_path, 'r') as fp:
                info = json.load(fp)
                if num_passes is not None:
                    eff_num_passes = min(num_passes, len(info['filenames']))
                    ratio = eff_num_passes / len(info['filenames'])
                    cur_record_info['num_batch'] += int(
                        info['num_batch'] * ratio
                    )
                    cur_record_info['filenames'] += info['filenames'][
                        :eff_num_passes
                    ]
                else:
                    cur_record_info['num_batch'] += info['num_batch']
                    cur_record_info['filenames'] += info['filenames']

        # overwrite directory for `cur_record_info`
        new_filenames = []
        for filename in cur_record_info['filenames']:
            basename = os.path.basename(filename)
            new_filename = os.path.join(record_dir, basename)
            new_filenames.append(new_filename)
        cur_record_info['filenames'] = new_filenames

        tf.logging.info(
            '[Dir %d] Number of chosen batches: %s',
            idx,
            cur_record_info['num_batch'],
        )
        tf.logging.info(
            '[Dir %d] Number of chosen files: %s',
            idx,
            len(cur_record_info['filenames']),
        )
        tf.logging.info(cur_record_info['filenames'])

        # add `cur_record_info` to global `record_info`
        record_info['num_batch'] += cur_record_info['num_batch']
        record_info['filenames'] += cur_record_info['filenames']

    tf.logging.info('Total number of batches: %d', record_info['num_batch'])
    tf.logging.info('Total number of files: %d', len(record_info['filenames']))
    tf.logging.info(record_info['filenames'])

    def input_fn(params):
        """docs."""
        assert params['batch_size'] * num_core_per_host == bsz_per_host

        dataset = get_dataset(
            params = params,
            num_hosts = num_hosts,
            num_core_per_host = num_core_per_host,
            split = split,
            file_names = record_info['filenames'],
            num_batch = record_info['num_batch'],
            seq_len = seq_len,
            reuse_len = reuse_len,
            perm_size = perm_size,
            mask_alpha = mask_alpha,
            mask_beta = mask_beta,
            use_bfloat16 = use_bfloat16,
            num_predict = num_predict,
        )

        return dataset

    return input_fn, record_info


if __name__ == '__main__':
    FLAGS = flags.FLAGS
    flags.DEFINE_bool('use_tpu', True, help = 'whether to use TPUs')
    flags.DEFINE_integer('bsz_per_host', 32, help = 'batch size per host.')
    flags.DEFINE_integer(
        'num_core_per_host', 8, help = 'num TPU cores per host.'
    )

    flags.DEFINE_integer('seq_len', 512, help = 'Sequence length.')
    flags.DEFINE_integer(
        'reuse_len',
        256,
        help = 'Number of token that can be reused as memory. '
        'Could be half of `seq_len`.',
    )
    flags.DEFINE_bool('uncased', True, help = 'Use uncased inputs or not.')
    flags.DEFINE_bool(
        'bi_data', True, help = 'whether to create bidirectional data'
    )
    flags.DEFINE_integer(
        'mask_alpha', default = 6, help = 'How many tokens to form a group.'
    )
    flags.DEFINE_integer(
        'mask_beta',
        default = 1,
        help = 'How many tokens to mask within each group.',
    )
    flags.DEFINE_bool(
        'use_eod', True, help = 'whether to append EOD at the end of a doc.'
    )
    flags.DEFINE_bool(
        'from_raw_text',
        True,
        help = 'Whether the input is raw text or encoded ids.',
    )
    flags.DEFINE_integer(
        'num_predict', default = 85, help = 'Num of tokens to predict.'
    )

    flags.DEFINE_string(
        'input_glob', 'data/example/*.txt', help = 'Input file glob.'
    )
    flags.DEFINE_string(
        'sp_path', '', help = 'Path to the sentence piece model.'
    )
    flags.DEFINE_string(
        'save_dir',
        'proc_data/example',
        help = 'Directory for saving the processed data.',
    )
    flags.DEFINE_enum(
        'split',
        'train',
        ['train', 'dev', 'test'],
        help = 'Save the data as which split.',
    )

    flags.DEFINE_integer(
        'pass_id',
        0,
        help = 'ID of the current pass.'
        'Different passes sample different negative segment.',
    )
    flags.DEFINE_integer('num_task', 1, help = 'Number of total tasks.')
    flags.DEFINE_integer(
        'task',
        0,
        help = 'The Task ID. This value is used when '
        'using multiple workers to identify each worker.',
    )

    tf.logging.set_verbosity(tf.logging.INFO)
    tf.app.run(create_data)
