"""Is Your Code Generated by ChatGPT Really Correct? Rigorous Evaluation of Large Language Models for Code Generation
https://openreview.net/forum?id=1qvx610Cu7

The MBPP+ dataset is created by the EvalPlus framework which extends the original MBPP dataset
by adding more automatically generated test cases to each problem. Note MBPP+ only includes 399
tasks which are a subset of the original MBPP dataset. The subset is selected from the sanitized
MBPP (a subset of manually examined tasks by the original MBPP authors) and EvalPlus further
removes low-quality and ill-formed tasks for benchmark quality control.

Homepage: https://github.com/evalplus/evalplus
"""

import os

from bigcode_eval.tasks.mbpp import MBPP
from bigcode_eval.tasks.custom_metrics.code_eval import compute_code_eval

_CITATION = """
@inproceedings{evalplus,
  title = {Is Your Code Generated by Chat{GPT} Really Correct? Rigorous Evaluation of Large Language Models for Code Generation},
  author = {Liu, Jiawei and Xia, Chunqiu Steven and Wang, Yuyao and Zhang, Lingming},
  booktitle = {Thirty-seventh Conference on Neural Information Processing Systems},
  year = {2023},
  url = {https://openreview.net/forum?id=1qvx610Cu7},
}
"""


class MBPPPlus(MBPP):
    """A task represents an entire benchmark including its dataset, problems,
    answers, generation settings and evaluation methods.
    """

    DATASET_PATH = "evalplus/mbppplus"

    def get_prompt_wo_base(self, doc):
        """Builds the prompt for the LM to generate from.
        MBPP prompt is built following to InCoder (Fried et al.) approach
        prompt = docstring that includes one test
        """
        function_head = doc["code"].split("def")[1].split(":")[0].strip()

        prefixes = [
            "Write a python function to ",
            "Write a python function that ",
            "Write a python function ",
            "Write a function that ",
            "Write a function to ",
            "Write a function which ",
            "Write a function ",
            "Write function to ",
            "write a function that ",
        ]

        description = doc["prompt"]
        for prefix in prefixes:
            if prefix in description:
                description = description.replace(prefix, "")

        description = description.strip().rstrip(".").capitalize()

        test_example = doc["test_list"][0]
        if not test_example.startswith(">>>"):
            test_example = f">>> {test_example}"

        inp = f"""Write a Python function named `{function_head}` to solve the following problem:
{description}
{test_example}"""

        prompt = f'Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n### Instruction:\n{inp}\n\n### Response:\n'
        return prompt.strip()

    def get_prompt(self, doc):
        """Builds the prompt for the LM to generate from.
        MBPP prompt is built following to InCoder (Fried et al.) approach
        prompt = docstring that includes one test
        """
        function_head = doc["code"].split("def")[1].split(":")[0].strip()

        prefixes = [
            "Write a python function to ",
            "Write a python function that ",
            "Write a python function ",
            "Write a function that ",
            "Write a function to ",
            "Write a function which ",
            "Write a function ",
            "Write function to ",
            "write a function that ",
        ]

        description = doc["prompt"]
        for prefix in prefixes:
            if prefix in description:
                description = description.replace(prefix, "")

        description = description.strip().rstrip(".").capitalize()

        test_example = doc["test_list"][0]
        if not test_example.startswith(">>>"):
            test_example = f">>> {test_example}"

        inp = f"""Write a Python function named `{function_head}` to solve the following problem:
{description}
{test_example}"""

        prompt_base = f'''def {function_head}:
    """
    {description}

    {test_example}
    """
'''.strip()

        prompt = f'Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n### Instruction:\n{inp}\n\n### Response:\n{prompt_base}'
        return prompt.strip()

    # NOTE(@ganler): MBPP+ extends the original MBPP jsonl data with a "test" field which
    #                includes the testing code ready for execution. Note the "test" field
    #                is different from HumanEval(+) which further requires a `check` func
    def get_reference(self, doc):
        """Builds the reference solution for the doc (sample from the test dataset)."""
        use_mbpp_tests = os.getenv("MBBPPLUS_USE_MBPP_TESTS", "0")
        if use_mbpp_tests == "1":
            return "\n".join(doc["test_list"])
        return "\n" + doc["test"]

    def get_dataset(self):
        """Returns dataset for the task or an iterable of any object, that get_prompt can handle"""
        dataset = self.dataset["test"]
        return dataset

    def process_results(self, generations, references):
        """Takes the list of LM generations and evaluates them against ground truth references,
        returning the metric for the generations.
        :param generations: list(list(str))
            list of lists containing generations
        :param references: list(str)
            list of str containing refrences
        """
        python_imports = "\n".join([
            "import math",
            "import re",
            "import sys",
            "import copy",
            "import datetime",
            "import itertools",
            "import collections",
            "import heapq",
            "import statistics",
            "import functools",
            "import hashlib",
            "import numpy",
            "import numpy as np",
            "import string",
            "from typing import *",
            "from collections import *",
        ])
        generations = [
            [(python_imports + "\n" + g).strip() for g in gen] for gen in generations
        ]

        results, _ = compute_code_eval(
            references=references,
            predictions=generations,
            num_workers=48,  # 48 workers
            timeout=10.0,  # 10s timeout
        )
        return results

    def remove_last_block(self, code):
        """
        Adapted from https://github.com/THUDM/CodeGeeX/blob/23ee51505a2bcd34d59d2e271b22e5bd91475462/codegeex/benchmark/utils.py#L151
        """
        for w in self.stop_words:
            if w in code:
                code = code[:code.find(w)]

        ### Find the first occassion where a chain of { } is closed
        for i, line in enumerate(code.split("\n")):
            if len(line.strip()) > 0 and line[0] != ' ' and line[0] != '\t':
                return "\n".join(code.split("\n")[:i])
        return code

    def postprocess_generation(self, generation, idx):
        """Defines the postprocessing for a LM generation.
        :param generation: str
            code generation from LM
        :param idx: int
            index of doc in the dataset to which the generation belongs
            (not used for Humaneval-Task)
        """
        doc = self.get_dataset()[idx]
        prompt = self.get_prompt(doc)
        prompt_wo_base = self.get_prompt_wo_base(doc)
        gen = self.remove_last_block(generation[len(prompt):].rstrip())
        return (prompt + gen)[len(prompt_wo_base):]