# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os

import numpy as np
import paddle


class MedicalCorpus(paddle.io.Dataset):
    def __init__(self, data_path, tokenizer):
        self.data_path = data_path
        self.tokenizer = tokenizer
        # Add ids for suffixal chinese tokens in tokenized text, e.g. '##度' in '百度'.
        # It should coincide with the vocab dictionary in preprocess.py.
        orig_len = len(self.tokenizer)  # noqa:F841
        suffix_vocab = {}
        for idx, token in enumerate(range(0x4E00, 0x9FA6)):
            suffix_vocab[len(self.tokenizer) + idx] = "##" + chr(token)
        self.tokenizer.added_tokens_decoder.update(suffix_vocab)
        self._samples, self._global_index = self._read_data_files(data_path)

    def _get_data_files(self, data_path):
        # Get all prefix of .npy/.npz files in the current and next-level directories.
        files = [
            os.path.join(data_path, f)
            for f in os.listdir(data_path)
            if (os.path.isfile(os.path.join(data_path, f)) and "_idx.npz" in str(f))
        ]
        files = [x.replace("_idx.npz", "") for x in files]
        return files

    def _read_data_files(self, data_path):
        data_files = self._get_data_files(data_path)
        samples = []
        indexes = []
        for file_id, file_name in enumerate(data_files):

            for suffix in ["_ids.npy", "_idx.npz"]:
                if not os.path.isfile(file_name + suffix):
                    raise ValueError("File Not found, %s" % (file_name + suffix))

            token_ids = np.load(file_name + "_ids.npy", mmap_mode="r", allow_pickle=True)
            samples.append(token_ids)

            split_ids = np.load(file_name + "_idx.npz")
            end_ids = np.cumsum(split_ids["lens"], dtype=np.int64)
            file_ids = np.full(end_ids.shape, file_id)
            split_ids = np.stack([file_ids, end_ids], axis=-1)
            indexes.extend(split_ids)
        indexes = np.stack(indexes, axis=0)
        return samples, indexes

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

    def __getitem__(self, index):
        file_id, end_id = self._global_index[index]
        start_id = 0
        if index > 0:
            pre_file_id, pre_end_id = self._global_index[index - 1]
            if pre_file_id == file_id:
                start_id = pre_end_id
        word_token_ids = self._samples[file_id][start_id:end_id]
        token_ids = []
        is_suffix = np.zeros(word_token_ids.shape)
        for idx, token_id in enumerate(word_token_ids):
            token = self.tokenizer.convert_ids_to_tokens(int(token_id))
            if "##" in token:
                token_id = self.tokenizer.convert_tokens_to_ids(token[-1])
                is_suffix[idx] = 1
            token_ids.append(token_id)

        return token_ids, is_suffix.astype(np.int64)


class DataCollatorForErnieHealth(object):
    def __init__(self, tokenizer, mlm_prob, max_seq_length, return_dict=False):
        self.tokenizer = tokenizer
        self.mlm_prob = mlm_prob
        self.max_seq_len = max_seq_length
        self.return_dict = return_dict
        self._ids = {
            "cls": self.tokenizer.convert_tokens_to_ids(self.tokenizer.cls_token),
            "sep": self.tokenizer.convert_tokens_to_ids(self.tokenizer.sep_token),
            "pad": self.tokenizer.convert_tokens_to_ids(self.tokenizer.pad_token),
            "mask": self.tokenizer.convert_tokens_to_ids(self.tokenizer.mask_token),
        }

    def __call__(self, data):
        masked_input_ids_a, input_ids_a, labels_a = self.mask_tokens(data)
        masked_input_ids_b, input_ids_b, labels_b = self.mask_tokens(data)
        masked_input_ids = paddle.concat([masked_input_ids_a, masked_input_ids_b], axis=0).astype("int64")
        input_ids = paddle.concat([input_ids_a, input_ids_b], axis=0)
        labels = paddle.concat([labels_a, labels_b], axis=0)
        if self.return_dict:
            return {"input_ids": masked_input_ids, "raw_input_ids": input_ids, "generator_labels": labels}

        else:
            return masked_input_ids, input_ids, labels

    def mask_tokens(self, batch_data):

        token_ids = [x[0] for x in batch_data]
        is_suffix = [x[1] for x in batch_data]

        # Create probability matrix where the probability of real tokens is
        # self.mlm_prob, while that of others is zero.
        data = self.add_special_tokens_and_set_maskprob(token_ids, is_suffix)
        token_ids, is_suffix, prob_matrix = data
        token_ids = paddle.to_tensor(token_ids, dtype="int64", stop_gradient=True)
        masked_token_ids = token_ids.clone()
        labels = token_ids.clone()

        # Create masks for words, where '百' must be masked if '度' is masked
        # for the word '百度'.
        prob_matrix = prob_matrix * (1 - is_suffix)
        word_mask_index = np.random.binomial(1, prob_matrix).astype("float")
        is_suffix_mask = is_suffix == 1
        word_mask_index_tmp = word_mask_index
        while word_mask_index_tmp.sum() > 0:
            word_mask_index_tmp = np.concatenate(
                [np.zeros((word_mask_index.shape[0], 1)), word_mask_index_tmp[:, :-1]], axis=1
            )
            word_mask_index_tmp = word_mask_index_tmp * is_suffix_mask
            word_mask_index += word_mask_index_tmp
        word_mask_index = word_mask_index.astype("bool")
        labels[~word_mask_index] = -100

        # 80% replaced with [MASK].
        token_mask_index = paddle.bernoulli(paddle.full(labels.shape, 0.8)).astype("bool").numpy() & word_mask_index
        masked_token_ids[token_mask_index] = self._ids["mask"]

        # 10% replaced with random token ids.
        token_random_index = paddle.to_tensor(
            paddle.bernoulli(paddle.full(labels.shape, 0.5)).astype("bool").numpy()
            & word_mask_index
            & ~token_mask_index
        )
        random_tokens = paddle.randint(low=0, high=self.tokenizer.vocab_size, shape=labels.shape, dtype="int64")
        masked_token_ids = paddle.where(token_random_index, random_tokens, masked_token_ids)

        return masked_token_ids, token_ids, labels

    def add_special_tokens_and_set_maskprob(self, token_ids, is_suffix):
        batch_size = len(token_ids)
        batch_token_ids = np.full((batch_size, self.max_seq_len), self._ids["pad"])
        batch_token_ids[:, 0] = self._ids["cls"]
        batch_is_suffix = np.full_like(batch_token_ids, -1)
        prob_matrix = np.zeros_like(batch_token_ids, dtype="float32")

        for idx in range(batch_size):
            if len(token_ids[idx]) > self.max_seq_len - 2:
                token_ids[idx] = token_ids[idx][: self.max_seq_len - 2]
                is_suffix[idx] = is_suffix[idx][: self.max_seq_len - 2]
            seq_len = len(token_ids[idx])
            batch_token_ids[idx, seq_len + 1] = self._ids["sep"]
            batch_token_ids[idx, 1 : seq_len + 1] = token_ids[idx]
            batch_is_suffix[idx, 1 : seq_len + 1] = is_suffix[idx]
            prob_matrix[idx, 1 : seq_len + 1] = self.mlm_prob

        return batch_token_ids, batch_is_suffix, prob_matrix


def create_dataloader(dataset, mode="train", batch_size=1, use_gpu=True, data_collator=None):
    """
    Creats dataloader.
    Args:
        dataset(obj:`paddle.io.Dataset`):
            Dataset instance.
        mode(obj:`str`, optional, defaults to obj:`train`):
            If mode is 'train', it will shuffle the dataset randomly.
        batch_size(obj:`int`, optional, defaults to 1):
            The sample number of a mini-batch.
        use_gpu(obj:`bool`, optional, defaults to obj:`True`):
            Whether to use gpu to run.
    Returns:
        dataloader(obj:`paddle.io.DataLoader`): The dataloader which generates batches.
    """

    if mode == "train" and use_gpu:
        sampler = paddle.io.DistributedBatchSampler(dataset=dataset, batch_size=batch_size, shuffle=True)
        dataloader = paddle.io.DataLoader(
            dataset, batch_sampler=sampler, return_list=True, collate_fn=data_collator, num_workers=0
        )
    else:
        shuffle = True if mode == "train" else False
        sampler = paddle.io.BatchSampler(dataset=dataset, batch_size=batch_size, shuffle=shuffle)
        dataloader = paddle.io.DataLoader(
            dataset, batch_sampler=sampler, return_list=True, collate_fn=data_collator, num_workers=0
        )

    return dataloader
