# coding:utf-8
# Copyright (c) 2021  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 paddle

from ..data import Pad, Stack, Tuple
from ..transformers import GPTChineseTokenizer, GPTForGreedyGeneration, GPTTokenizer
from .task import Task
from .utils import download_file, static_mode_guard

usage = r"""
         """

URLS = {
    "gpt-cpm-large-cn": [
        "https://bj.bcebos.com/paddlenlp/taskflow/text_generation/gpt-cpm/gpt-cpm-large-cn_params.tar",
        "5aad6f81053cfdbba4797f044fcf66d1",
    ],
}


class TextGenerationTask(Task):
    """
    The text generation model to predict the question or chinese  poetry.
    Args:
        task(string): The name of task.
        model(string): The model name in the task.
        kwargs (dict, optional): Additional keyword arguments passed along to the specific task.
    """

    def __init__(self, task, model, **kwargs):
        super().__init__(task=task, model=model, **kwargs)
        # Default to static mode
        self._static_mode = True
        self._usage = usage
        if self._static_mode:
            download_file(self._task_path, "gpt-cpm-large-cn_params.tar", URLS[self.model][0], URLS[self.model][1])
            self._get_inference_model()
        else:
            self._construct_model(model)
        self._construct_tokenizer(model)
        self.kwargs["generation_task"] = task

    def _construct_input_spec(self):
        """
        Construct the input spec for the convert dygraph model to static model.
        """
        self._input_spec = [paddle.static.InputSpec(shape=[None, None], dtype="int64", name="token_ids")]

    def _construct_model(self, model):
        """
        Construct the inference model for the predictor.
        """
        model_instance = GPTForGreedyGeneration.from_pretrained(self.model, max_predict_len=32)
        # Load the model parameter for the predict
        model_instance.eval()
        self._model = model_instance

    def _construct_tokenizer(self, model):
        """
        Construct the tokenizer for the predictor.
        """
        if self.model == "gpt-cpm-large-cn":
            tokenizer_instance = GPTChineseTokenizer.from_pretrained(model)
        else:
            tokenizer_instance = GPTTokenizer.from_pretrained(model)

        self._tokenizer = tokenizer_instance

    def _preprocess(self, inputs, padding=True, add_special_tokens=True):
        """
        Transform the raw text to the model inputs, two steps involved:
           1) Transform the raw text to token ids.
           2) Generate the other model inputs from the raw text and token ids.
        """
        inputs = self._check_input_text(inputs)
        # Get the config from the kwargs
        batch_size = self.kwargs["batch_size"] if "batch_size" in self.kwargs else 1
        generation_task = self.kwargs["generation_task"] if "generation_task" in self.kwargs else "question_answering"

        def select_few_shot_input(model_name, generation_task):
            pre_input = ""
            if generation_task not in ["question_answering", "poetry_generation"]:
                raise ValueError("The generation task must be question or poetry")
            if model_name == "gpt-cpm-large-cn":
                if generation_task == "question_answering":
                    pre_input = "问题：中国的首都是哪里？答案：北京。\n问题：{} 答案："
                else:
                    pre_input = "默写古诗: 大漠孤烟直，长河落日圆。\n{}"
            return pre_input

        pre_input = select_few_shot_input(self.model, generation_task)

        examples = []
        filter_inputs = []
        for input_text in inputs:
            if not (isinstance(input_text, str) and len(input_text) > 0):
                continue
            filter_inputs.append(input_text)
            few_shot_input = pre_input.format(input_text)
            ids = self._tokenizer(few_shot_input)["input_ids"]
            examples.append((ids, len(ids)))

        batchify_fn = lambda samples, fn=Tuple(
            Pad(axis=0, pad_val=0, dtype="int64"),
            Stack(dtype="int64"),  # seq_len
        ): fn(samples)

        batches = [examples[idx : idx + batch_size] for idx in range(0, len(examples), batch_size)]
        outputs = {}
        outputs["text"] = filter_inputs
        outputs["data_loader"] = batches
        self._batchify_fn = batchify_fn
        return outputs

    def _run_model(self, inputs):
        """
        Run the task model from the outputs of the `_tokenize` function.
        """
        results = []
        lens = []
        with static_mode_guard():
            for batch in inputs["data_loader"]:
                ids, seq_len = self._batchify_fn(batch)
                self.input_handles[0].copy_from_cpu(ids)
                self.predictor.run()
                result = self.output_handle[0].copy_to_cpu().tolist()
                results.extend(result)
                lens.extend(seq_len.tolist())
        inputs["results"] = results
        inputs["lens"] = lens
        return inputs

    def _postprocess(self, inputs):
        """
        The model output is tag ids, this function will convert the model output to raw text.
        """
        batch_out = []
        preds = inputs["results"]
        for index in range(0, len(preds)):
            seq_len = inputs["lens"][index]
            single_result = {}
            single_result["text"] = inputs["text"][index]
            single_result["answer"] = self._tokenizer.convert_ids_to_string(preds[index][seq_len:-1])
            batch_out.append(single_result)
        return batch_out