class InitQATemplate:
    """
    Takes a prompt template and provides methods for filling in blanks.
    The format is as follows:
    [APE] is where text will be generated by the LLM.
    [full_DEMO] is where the full demo will be inserted.
    [INPUT] is where the input to the first demo will be inserted.
    [OUTPUT] is where the output from the first demo will be inserted.
    """

    def __init__(self, template):
        self.template = template
        # Check that the template is valid
        # There should be exactly one [APE] token
        # assert self.template.count('[APE]') == 1

    def fill(self, full_demo=''):
        """
        Fills in the template with the given values.
        """
        return self.template.replace('[full_DEMO]', full_demo)


class GenerationTemplate:
    """
    Takes a prompt template and provides methods for filling in blanks.
    The format is as follows:
    [APE] is where text will be generated by the LLM.
    [full_DEMO] is where the full demo will be inserted.
    [INPUT] is where the input to the first demo will be inserted.
    [OUTPUT] is where the output from the first demo will be inserted.
    """

    def __init__(self, template):
        self.template = template
        # Check that the template is valid
        # There should be exactly one [APE] token
        # assert self.template.count('[APE]') == 1

    def fill(self, full_demo='', input='', output=''):
        """
        Fills in the template with the given values.
        """
        return self.template.replace('[full_DEMO]', full_demo).replace(
            '[INPUT]', input).replace('[OUTPUT]', output)


class EvalTemplate:
    """
    Takes a prompt template and provides methods for filling in blanks.
    The format is as follows:
    [PROMPT] is where the prompt will be inserted.
    [full_DEMO] is where the full demo will be inserted.
    [INPUT] is where the input to the first demo will be inserted.
    [OUTPUT] is where the output from the first demo will be inserted.
    """

    def __init__(self, template):
        self.template = template

    def fill(self, prompt='', full_demo='', input='', output=''):
        """
        Fills in the template with the given values.
        """
        return self.template.replace('[PROMPT]', prompt).replace(
            '[full_DEMO]', full_demo).replace('[INPUT]', input).replace('[OUTPUT]', output)

    def convert_to_generation_template(self):
        """
        Converts the evaluation template to a generation template.
        """
        return GenerationTemplate(self.template.replace('[PROMPT]', '[APE]'))


class DemosTemplate:
    """
    Takes a template for the full demo and provides methods for filling in blanks.
    The format is as follows:
    [INPUT], [OUTPUT]

    """

    def __init__(self, template, delimiter='\n\n'):
        self.template = template
        self.delimiter = delimiter

    def fill(self, data):
        """
        Fills in the template with the given values. Data is a tuple of lists.
        """
        demos = ''
        for i, (input_, output_) in enumerate(zip(*data)):
            demos += self.template.replace('[INPUT]', input_).replace(
                '[OUTPUT]', output_)

            if i != len(data[0]) - 1:
                demos += self.delimiter

        return demos


class GAEvolutionTemplate:
    """
    Template for querying a GPT model to perform evolutionary operations such as
    crossover and mutation to generate better prompts.
    """

    def __init__(self, template):
        self.template = template

    def fill(self, prompt1='', prompt2=''):
        crossover_prompt = f"Cross over the following prompts and generate a new prompt:\nPrompt 1: {prompt1}\nPrompt 2: {prompt2}"
        mutation_prompt = f"Mutate the prompt generated in Step 1 and generate a final prompt bracketed with <prompt> and </prompt>."
        filled_template = self.template.replace(
            '[CROSSOVER_PROMPTS]', crossover_prompt).replace('[MUTATION_PROMPT]', mutation_prompt)
        return filled_template


class DEEvolutionTemplate:
    """
    Template for querying a GPT model to perform differential evolution operations such as
    mutation and crossover to generate better prompts.
    """

    def __init__(self, template):
        self.template = template

    def fill(self, prompt1='', prompt2='', prompt3='', prompt4=''):
        crossover_prompt = f"Identify the different parts between Prompt 1 and Prompt 2:\nPrompt 1: {prompt1}\nPrompt 2: {prompt2}"
        mutation_prompt = f"Randomly mutate the different parts."
        combine_prompt = f"Combine the different parts with Prompt 3, selectively replace it with the different parts in Step 2 and generate a new prompt.\nPrompt 3: {prompt3}"
        crossover_final = f"Cross over the prompt in Step 3 with the following basic prompt and generate a final prompt bracketed with <prompt> and </prompt>:\nBasic Prompt: {prompt4} "

        filled_template = self.template.replace('[CROSSOVER_PROMPTS]', crossover_prompt).replace(
            '[MUTATION_PROMPT]', mutation_prompt).replace('[COMBINE_PROMPT]', combine_prompt).replace(
                '[CROSSOVER_FINAL]', crossover_final)
        return filled_template
