# 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.

from glob import glob

import numpy as np
import paddle.distributed as dist
from paddle.io import IterableDataset

from paddlenlp.transformers.tokenizer_utils import convert_to_unicode


class DialogueDataset(IterableDataset):
    def __init__(
        self,
        filepattern,
        batch_size,
        pad_token_id,
        bos_token_id,
        sort_pool_size=2**16,
        seed=1,
        n_procs=None,
        rank=None,
        mode="test",
    ):
        super(DialogueDataset, self).__init__()

        self.file_list = glob(filepattern)
        self.sort_pool_size = 0 if mode == "test" else sort_pool_size
        self.n_procs = n_procs if n_procs else dist.get_world_size()
        self.rank = rank if rank else dist.get_rank()
        self.batch_size = batch_size * self.n_procs
        self.shuffle = True if mode == "train" else False
        self.mode = mode
        self.pad_id = pad_token_id
        self.bos_id = bos_token_id
        self.global_rng = np.random.RandomState(seed)

        assert len(self.file_list) > 0, "There is no files in %s." % filepattern

    def load_file(self, file_path):
        with open(file_path, "r", encoding="utf-8") as fin:
            for i, line in enumerate(fin):
                cols = convert_to_unicode(line).strip().split(";")
                cols = list(map(lambda x: list(map(int, x.split(" "))), cols))
                if len(cols) > 3:
                    cols = cols[:3]
                token_ids, type_ids, pos_ids = cols
                if self.mode == "test":
                    tgt_start_idx = len(cols[0])
                else:
                    tgt_start_idx = token_ids.index(self.bos_id, 1)
                sample = [token_ids, type_ids, pos_ids, tgt_start_idx]
                yield sample

    def get_sorted_batch(self, pool):
        """Generate sorted batches from pool."""
        pool = sorted(pool, key=lambda sample: len(sample[0]))
        batches = []
        batch, max_len = [], 0
        for sample in pool:
            max_len = max(max_len, len(sample[0]))
            if self.mode == "test":
                to_append = len(batch) < self.batch_size
            else:
                to_append = (len(batch) + 1) * max_len <= self.batch_size
            if to_append:
                batch.append(sample)
            else:
                batches.append(batch)
                batch, max_len = [sample], len(sample[0])
        if len(batch) > 0:
            batches.append(batch)
        if self.shuffle:
            self.global_rng.shuffle(batches)
        for batch in batches:
            yield batch

    @property
    def get_batch(self):
        all_files = list(self.file_list)
        if self.shuffle:
            self.global_rng.shuffle(all_files)
        if self.sort_pool_size > 0:
            pool = []
            for file_path in all_files:
                for sample in self.load_file(file_path):
                    pool.append(sample)
                    if len(pool) == self.sort_pool_size:
                        for batch in self.get_sorted_batch(pool):
                            yield batch
                        pool = []
                if len(pool) > 0:
                    for batch in self.get_sorted_batch(pool):
                        yield batch
        else:
            batch, max_len = [], 0
            for file_path in all_files:
                for sample in self.load_file(file_path):
                    max_len = max(max_len, len(sample[0]))
                    if self.mode == "test":
                        to_append = len(batch) < self.batch_size
                    else:
                        to_append = (len(batch) + 1) * max_len <= self.batch_size
                    if to_append:
                        batch.append(sample)
                    else:
                        yield batch
                        batch, max_len = [sample], len(sample[0])
            if len(batch) > 0:
                yield batch

    def pad_batch_data(self, batch):
        """Pad the instances to the max sequence length in batch."""
        max_len = max(map(len, batch))
        batch_data = np.array([list(data) + [self.pad_id] * (max_len - len(data)) for data in batch], dtype="int64")
        return batch_data

    def gen_tgt_label_and_pos(self, batch_token_ids, batch_tgt_start_idx):
        max_len = max(map(len, batch_token_ids))
        tgt_label = []
        tgt_pos = []
        for sent_index, sent in enumerate(batch_token_ids):
            sent_b_index = batch_tgt_start_idx[sent_index]
            tgt_label.extend(sent[sent_b_index + 1 :])
            tgt_pos.extend([sent_index * max_len + i for i in range(sent_b_index, len(sent) - 1)])
        tgt_label = np.array(tgt_label).astype("int64")
        tgt_pos = np.array(tgt_pos).astype("int64")

        return tgt_label, tgt_pos

    def gen_self_attn_mask(self, batch_token_ids, batch_tgt_start_idx):
        max_len = max(map(len, batch_token_ids))
        input_mask_data = np.zeros((len(batch_token_ids), max_len, max_len))
        for index, mask_data in enumerate(input_mask_data):
            start = batch_tgt_start_idx[index]
            end = len(batch_token_ids[index])
            mask_data[:end, :start] = 1.0
            # Generate the lower triangular matrix using the slice of matrix
            b = np.tril(np.ones([end - start, end - start]), 0)
            mask_data[start:end, start:end] = b
        return input_mask_data.astype("float32")

    def __iter__(self):
        for batch_data in self.get_batch:
            # sample [token_ids, type_ids, pos_ids, tgt_start_idx]
            # raw_batch [sample0, sample1, ...]
            if self.n_procs > 1:
                batch_data = batch_data[self.rank :: self.n_procs]
            batch_data = zip(*batch_data)
            token_ids, type_ids, pos_ids, tgt_start_idx = batch_data

            pad_token_ids = self.pad_batch_data(token_ids)
            pad_type_ids = self.pad_batch_data(type_ids)
            pad_pos_ids = self.pad_batch_data(pos_ids)

            generation_mask = self.gen_self_attn_mask(token_ids, tgt_start_idx)

            if self.mode == "test":
                # [batch_size, 1]
                tgt_ids = np.array([[self.bos_id]] * len(token_ids), dtype="int64")
                tgt_type = np.ones((len(token_ids), 1), dtype="int64")
                tgt_pos = np.array(tgt_start_idx, dtype="int64").reshape(-1, 1)
                tgt_generation_mask = generation_mask[:, 0:1, :].astype("float32")

                pad_token_ids = np.concatenate((pad_token_ids, tgt_ids), axis=1)
                pad_type_ids = np.concatenate((pad_type_ids, tgt_type), axis=1)
                pad_pos_ids = np.concatenate((pad_pos_ids, tgt_pos), axis=1)
                generation_mask = np.concatenate((generation_mask, tgt_generation_mask), axis=1)

                append_mask = np.zeros((generation_mask.shape[0], generation_mask.shape[1], 1), dtype="float32")
                append_mask[:, -1, :] = 1.0
                generation_mask = np.concatenate((generation_mask, append_mask), axis=2)
                generation_mask = (generation_mask - 1.0) * 1e9
                generation_mask = np.expand_dims(generation_mask, axis=1)
                yield (pad_token_ids, pad_type_ids, pad_pos_ids, generation_mask)
            else:
                tgt_label, tgt_pos = self.gen_tgt_label_and_pos(token_ids, tgt_start_idx)
                generation_mask = (generation_mask - 1.0) * 1e9
                generation_mask = np.expand_dims(generation_mask, axis=1)
                yield (pad_token_ids, pad_type_ids, pad_pos_ids, generation_mask, tgt_label, tgt_pos)


def post_process_response(token_ids, tokenizer):
    """
    Post-process the decoded sequence. Truncate from the first
    <eos> and remove the <bos> and <eos> tokens currently.
    """
    eos_pos = len(token_ids)
    for i, tok_id in enumerate(token_ids):
        if tok_id == tokenizer.sep_token_id:
            eos_pos = i
            break
    token_ids = token_ids[:eos_pos]
    tokens = tokenizer.convert_ids_to_tokens(token_ids)
    response = tokenizer.merge_subword(tokens)
    return token_ids, response


def get_in_turn_repetition(pred, is_cn=False):
    """Get in-turn repetition."""
    if len(pred) == 0:
        return 1.0
    if isinstance(pred[0], str):
        pred = [tok.lower() for tok in pred]
        if is_cn:
            pred = "".join(pred)
    tri_grams = set()
    for i in range(len(pred) - 2):
        tri_gram = tuple(pred[i : i + 3])
        if tri_gram in tri_grams:
            return True
        tri_grams.add(tri_gram)
    return False


def select_response(ids, scores, tokenizer, max_dec_len=None, num_samples=1):
    ids = ids.numpy().tolist()
    scores = scores.numpy()

    if len(ids) != len(scores) or (len(ids) % num_samples) != 0:
        raise ValueError("the length of `ids` is {}, but the `num_samples` is {}".format(len(ids), num_samples))

    group = []
    tmp = []
    for pred, score in zip(ids, scores):
        pred_token_ids, pred_tokens = post_process_response(pred, tokenizer)
        num_token = len(pred_token_ids)
        response = " ".join(pred_tokens)

        in_turn_repetition = get_in_turn_repetition(pred_tokens, True) or get_in_turn_repetition(pred_token_ids)
        # not ending
        if max_dec_len is not None and num_token >= max_dec_len:
            score -= 1e3
        elif in_turn_repetition:
            score -= 1e3

        tmp.append([response, score])
        if len(tmp) == num_samples:
            group.append(tmp)
            tmp = []

    results = []
    for preds in group:
        preds = sorted(preds, key=lambda x: -x[1])
        results.append(preds[0][0])
    return results
