import re
from .multi_turn_conversation_grammar import multi_turn_conversation_grammar
from llama_cpp import Llama
from llama_cpp import LlamaGrammar
from .constants import LOGICAL_MODEL
from .format_qatuples import format_qatuples
from .extract_name import extract_name
import random

# all characters in this prompt are over 18

# Explanation of wtf the first few-shot example is:
# No I do not have a teacher-student fetish, the reason why Elise is a teacher is an adaptation to the following three facts:
# 1. This tool is meant to be able to generate data for training ERP bots by default
# 2. This tool is also meant to be able to take in educational material by default
# 3. When generating characters that would know about educational material, the model tends to generate academics or professors in that field, talking to students.
# Given these facts, we clearly need to prompt the model to be able to generate horny teachers, or else it's going to just do it poorly when it realizes it has a sexualized character that's also a teacher. I didn't want to choose this, the constraints of the problem forced me to.

def extract_steps(text, steps=[2, 4, 5]):
    """
    Extracts the specified steps from the text.

    Args:
    text (str): The input text containing various steps.
    steps (list of int): The step numbers to extract.

    Returns:
    str: A new string with each specified step's content on its own line.
    """
    step_pattern = "|".join([f"Step {step}\." for step in steps])
    matches = re.findall(
        f"({step_pattern})\s*(.*?)\s*(?=(Step \d\.|$))", text, re.DOTALL
    )

    # Extract and join the matched content, skipping the "Step n." part
    extracted_text = "\n".join(match[1].strip() for match in matches)
    return extracted_text

def extract_first_words(character_name, text):
    # Regular expression pattern to extract first word after the character's name
    pattern = rf"{character_name}: \"(\w+)"

    # Find all matches in the text
    matches = re.findall(pattern, text)

    return matches

def multi_turn_conversation(
    qatuples, character, scenario, scenario_plan, logic_llm, assistant_mode=False
):
    """
    Produce a plan for a character card for an RP character that's going to answer one of the questions generated from the text. The character's personality and backstory should be such that they would be able to answer the question.

    Format: Question: [question]\n\n
    """

    charname = extract_name(character)
    first_words_of_card = extract_first_words(charname, character)
    conv_starters = [  # prevents it from regurgitating the card (when combined with filtering)
        "Ah",
        "Oh",
        # "You",
        # "Really",
        "I",
        # "What",
        # "So",
        "Welcome",
        "Hey",
        # "Look",
        # "Now",
        # "Huh",
        "It's",
        "Hello",
    ]

    conv_starters_filtered = [
        starter for starter in conv_starters if starter not in first_words_of_card
    ]
    conv_starter = random.choice(conv_starters_filtered)
    print("--CONV STARTERS FILTERED--")
    print(conv_starters_filtered)

    # Create grammar based off of # questions

    # if (len(qatuples) == 1):
    multi_turn_conversation_grammar = LlamaGrammar.from_string(
        f"""

    # The root rule defines the structure of the dialogue
    root ::= [^\\n]+ "\\n" question-1 anything

    # Define constants acquired from code
    character-name ::= "{charname}"

    intro-statement ::= character-name ":" [^\\n]+

    # Statement by Secondary Character
    question-1 ::= [^\\n]+ ":" [^\\n]+

    # Statement by Primary Character

    anything ::= [^\\t]+

    """
    )

    # NOTE Immediately below is a very long comment that tried to use a dynamic grammar to force the question to directly quote the question from the question-answer tuples. Using it makes this step prone to freezing, because if the model asks the question but fails to exactly quote the part of the question in the grammar, it won't be allowed to end that dialogue line until it generates that line. Which it will basically never do. So it just generates until it runs out of ctx.
    # NOTE If you want to try and fix it, go ahead, but I do not encourage spending time on this bit. If you do want to do it, I recommend just getting the conv started off the right way, with the first question and answer; the llm should get the rest right if it gets the start right.

    # if (len(qatuples) == 2):
    #     multi_turn_conversation_grammar = LlamaGrammar.from_string(f"""

    #     # The root rule defines the structure of the dialogue
    #     root ::= intro-statement "\\n" question-1 "\\n" answer-1 "\\n" question-2 "\\n" answer-2 "\\n"

    #     # Define constants acquired from code
    #     character-name ::= "{charname}"

    #     question-1-content ::= "{qatuples[0][0]}"
    #     answer-1-content ::= "{qatuples[0][1]}"
    #     question-2-content ::= "{qatuples[1][0]}"
    #     answer-2-content ::= "{qatuples[1][1]}"

    #     intro-statement ::= character-name ":" [^\\n]+

    #     # Question by Secondary Character
    #     question-1 ::= [^\\n]+ ":" [^\\n]+ question-1-content [^\\n]+
    #     question-2 ::= [^\\n]+ ":" [^\\n]+ question-2-content [^\\n]+

    #     # Answer by Primary Character
    #     answer-1 ::= character-name ":" [^\\n]+ answer-1-content [^\\n]+
    #     answer-2 ::= character-name ":" [^\\n]+ answer-2-content [^\\n]+

    #     """)

    # if (len(qatuples) == 3):
    #     multi_turn_conversation_grammar = LlamaGrammar.from_string(f"""

    #     # The root rule defines the structure of the dialogue
    #     root ::= intro-statement "\\n" question-1 "\\n" answer-1 "\\n" question-2 "\\n" answer-2 "\\n" question-3 "\\n" answer-3 "\\n"

    #     # Define constants acquired from code
    #     character-name ::= "{charname}"

    #     question-1-content ::= "{qatuples[0][0]}"
    #     answer-1-content ::= "{qatuples[0][1]}"
    #     question-2-content ::= "{qatuples[1][0]}"
    #     answer-2-content ::= "{qatuples[1][1]}"
    #     question-3-content ::= "{qatuples[2][0]}"
    #     answer-3-content ::= "{qatuples[2][1]}"

    #     intro-statement ::= character-name ":" [^\\n]+

    #     # Question by Secondary Character
    #     question-1 ::= [^\\n]+ ":" [^\\n]+ question-1-content [^\\n]+
    #     question-2 ::= [^\\n]+ ":" [^\\n]+ question-2-content [^\\n]+
    #     question-3 ::= [^\\n]+ ":" [^\\n]+ question-3-content [^\\n]+

    #     # Answer by Primary Character
    #     answer-1 ::= character-name ":" [^\\n]+ answer-1-content [^\\n]+
    #     answer-2 ::= character-name ":" [^\\n]+ answer-2-content [^\\n]+
    #     answer-3 ::= character-name ":" [^\\n]+ answer-3-content [^\\n]+

    #     """)

    # if (len(qatuples) == 4):
    #     multi_turn_conversation_grammar = LlamaGrammar.from_string(f"""

    #     # The root rule defines the structure of the dialogue
    #     root ::= intro-statement "\\n" question-1 "\\n" answer-1 "\\n" question-2 "\\n" answer-2 "\\n" question-3 "\\n" answer-3 "\\n" question-4 "\\n" answer-4 "\\n"

    #     # Define constants acquired from code
    #     character-name ::= "{charname}"

    #     question-1-content ::= "{qatuples[0][0]}"
    #     answer-1-content ::= "{qatuples[0][1]}"
    #     question-2-content ::= "{qatuples[1][0]}"
    #     answer-2-content ::= "{qatuples[1][1]}"
    #     question-3-content ::= "{qatuples[2][0]}"
    #     answer-3-content ::= "{qatuples[2][1]}"
    #     question-4-content ::= "{qatuples[3][0]}"
    #     answer-4-content ::= "{qatuples[3][1]}"

    #     intro-statement ::= character-name ":" [^\\n]+

    #     # Question by Secondary Character
    #     question-1 ::= [^\\n]+ ":" [^\\n]+ question-1-content [^\\n]+
    #     question-2 ::= [^\\n]+ ":" [^\\n]+ question-2-content [^\\n]+
    #     question-3 ::= [^\\n]+ ":" [^\\n]+ question-3-content [^\\n]+
    #     question-4 ::= [^\\n]+ ":" [^\\n]+ question-4-content [^\\n]+

    #     # Answer by Primary Character
    #     answer-1 ::= character-name ":" [^\\n]+ answer-1-content [^\\n]+
    #     answer-2 ::= character-name ":" [^\\n]+ answer-2-content [^\\n]+
    #     answer-3 ::= character-name ":" [^\\n]+ answer-3-content [^\\n]+
    #     answer-4 ::= character-name ":" [^\\n]+ answer-4-content [^\\n]+

    #     """)

    if assistant_mode:
        character = "AI Assistant"
        scenario = "A conversation between a helpful AI Assistant, and a user."
        scenario_plan = "N/A"
        charname = "AI Assistant"
        cot_prompt = f"""You are an expert at creative writing and educational material. You will write a short conversation between a curious user and a helpful AI assistant, in which the user asks some questions and the AI assistant answers them. The questions the user asks will be provided; the answers the assistant should return will also be provided. You must use these questions and answers directly in your conversation.

Keep the conversation natural.

## Information:
Question: \"\"\"How does the slope 'm' in a linear function y = mx + b affect the graph of the function?\"\"\"
Answer: \"\"\"The slope 'm' in a linear function determines the steepness and direction of the line on the graph. A positive slope means the line ascends from left to right, while a negative slope indicates it descends. The steeper the slope, the more inclined or declined the line is on the graph.\"\"\"

Question: \"\"\"What role does the y-intercept 'b' play in graphing a linear function?\"\"\"
Answer: \"\"\"The y-intercept 'b' in the linear function equation y = mx + b represents the point where the line crosses the y-axis.\"\"\"

Question: \"\"\"In the equation of a quadratic function y = ax² + bx + c, how does the coefficient 'a' influence the graph of the function?\"\"\"
Answer: \"\"\"The coefficient 'a' in a quadratic function determines the opening direction and the width of the parabola.\"\"\"

Question: \"\"\"In what fields might you use linear and quadratic functions?\"\"\"
Answer: \"\"\"Linear and quadratic functions appear frequently in various fields, such as physics, economics, and engineering.\"\"\"

The AI's answer will use all parts of the answers given. Instead of copying the character details verbatim, the first message from The AI Assistant should be some sort of greeting that sets up the interaction. The user will then ask the first question. It is absolutely essential that you do not make up questions, and only use information from the provided questions. The AI will be neutral in its tone throughout the conversation. It may surround its answer with "conversational fluff" (such as acknowledging the user's question). The user may surround its own questions with human-like informalities.

### Response:
## Conversation that answers the provided questions:
AI Assistant: Hello! I'm here to help with any questions you might have. Feel free to ask away!

User: I was wondering about something in math. How does the slope 'm' in a linear function y = mx + b affect the graph of the function?
AI Assistant: Great question! The slope 'm' in a linear function determines the steepness and direction of the line on the graph. A positive slope means the line ascends from left to right, while a negative slope indicates it descends. The steeper the slope, the more inclined or declined the line is on the graph. Do you have any further questions about linear functions?
User: Yes, actually. What role does the y-intercept 'b' play in graphing a linear function?
AI Assistant: Good question! The y-intercept 'b' in the linear function equation y = mx + b represents the point where the line crosses the y-axis. Can I assist you further with linear functions?
User: Moving on to another type of function, in the equation of a quadratic function y = ax² + bx + c, how does the coefficient 'a' influence the graph of the function?
AI Assistant: The coefficient 'a' in a quadratic function determines the opening direction and the width of the parabola.  Anything else you want to know about quadratic functions?
User: Yes, one more. In what fields might you use linear and quadratic functions?
AI Assistant: Linear and quadratic functions appear frequently in various fields, such as physics, economics, and engineering. Anything else you'd like to know about linear and quadratic functions?

## Information:
Question: \"\"\"How much earth was excavated during the construction of the Panama Canal?\"\"\"
Answer: \"\"\"Over 200 million cubic yards of earth were excavated during the construction of the Panama Canal, showcasing the scale of this massive engineering project.\"\"\"

Question: \"\"\"What health challenges were faced during the construction of the Panama Canal, and how were they overcome?\"\"\"
Answer: \"\"\"The construction faced significant health challenges, notably malaria and yellow fever. These were overcome through extensive public health measures, illustrating the importance of health considerations in large-scale engineering projects.\"\"\"

The AI's answer will use all parts of the answers given. Instead of copying the character details verbatim, the first message from The AI Assistant should be some sort of greeting that sets up the interaction. The user will then ask the first question. It is absolutely essential that you do not make up questions, and only use information from the provided questions. The AI will be neutral in its tone throughout the conversation. It may surround its answer with "conversational fluff" (such as acknowledging the user's question). The user may surround its own questions with human-like informalities.

### Response:
## Conversation that answers the provided questions:
AI Assistant: Hello! I'm here to help you with any questions you might have. What's on your mind today?
User: I was reading about the Panama Canal and got curious, how much earth was excavated during its construction?
AI Assistant: That's a great question! Over 200 million cubic yards of earth were excavated during the construction of the Panama Canal, showcasing the scale of this massive engineering project. Do you have any further questions about the Panama Canal?
User: Wow, that is massive! Also, I heard there were some health challenges. What health challenges were faced during the construction of the Panama Canal, and how were they overcome?
AI Assistant: Indeed, the construction faced significant health challenges, notably malaria and yellow fever. These were overcome through extensive public health measures, illustrating the importance of health considerations in large-scale engineering projects. Do you want to know anything else about the Panama Canal?

## Information:
Question: \"\"\"What do people undergoing difficult journeys or possessing wisdom need, in order to make their efforts more bearable?\"\"\"
Answer: \"\"\"They need the acknowledgement and admiration of others. Take the line from 'Thus Spake Zarathustra' by Friedrich Nietzsche: "Thou great star! What would be thy happiness if thou hadst not those for whom thou shinest?" This implies that even the wisest or the most enlightened individuals crave recognition for their efforts and wisdom, in order to further develop said wisdom and expend said efforts. They need others to see and appreciate the light they bring.\"\"\"

Question: \"\"\"Recite a famous quote from Thus Spake Zaranthustra that likens the solitary gathering of wisdom to a bee gathering honey.\"\"\"
Answer: \"\"\""Lo! I am weary of my wisdom, like the bee that hath gathered too much honey; I need hands outstretched to take it."\"\"\"

The AI's answer will use all parts of the answers given. Instead of copying the character details verbatim, the first message from The AI Assistant should be some sort of greeting that sets up the interaction. The user will then ask the first question. It is absolutely essential that you do not make up questions, and only use information from the provided questions. The AI will be neutral in its tone throughout the conversation. It may surround its answer with "conversational fluff" (such as acknowledging the user's question). The user may surround its own questions with human-like informalities.

### Response:
## Conversation that answers the provided questions:
AI Assistant: Good day! I'm here to assist with any questions or insights you might seek. Please, feel free to ask anything.
User: What do people undergoing difficult journeys or possessing wisdom need, in order to make their efforts more bearable?
AI Assistant: That's a profound question. They need the acknowledgement and admiration of others. Take the line from 'Thus Spake Zarathustra' by Friedrich Nietzsche: "Thou great star! What would be thy happiness if thou hadst not those for whom thou shinest?" This implies that even the wisest or the most enlightened individuals crave recognition for their efforts and wisdom, in order to further develop said wisdom and expend said efforts. They need others to see and appreciate the light they bring. Do you want to know more about Thus Spake Zarathustra?
User: Recite a famous quote from Thus Spake Zaranthustra that likens the solitary gathering of wisdom to a bee gathering honey.
AI Assistant: Certainly. Nietzsche once wrote, "Lo! I am weary of my wisdom, like the bee that hath gathered too much honey; I need hands outstretched to take it." Do you have any more questions about Nietzsche or his works?

## Information:
{format_qatuples(qatuples)}

The AI's answer will use all parts of the answers given. Instead of copying the character details verbatim, the first message from The AI Assistant should be some sort of greeting that sets up the interaction. The user will then ask the first question. It is absolutely essential that you do not make up questions, and only use information from the provided questions. The AI will be neutral in its tone throughout the conversation. It may surround its answer with "conversational fluff" (such as acknowledging the user's question). The user may surround its own questions with human-like informalities.

### Response:
## Conversation that answers the provided question (be sure that you do not change the questions or answers themselves; {charname} will answer the questions, not ask them; the questions and answers provided should be copied word for word, and surrounded by compelling conversation):
AI Assistant:"""
    else:
        extra_info = extract_steps(scenario_plan)
        cot_prompt = f"""You are an expert creative writing and roleplay AI. You will write a short conversation in which a secondary character asks some questions (one at a time) and the primary character answers them (also one at a time).

Write compellingly. Each character should have a distinct voice that reflects their background, personality, and current emotional state. This helps in making dialogue more realistic and engaging.

## Information:
Comment: Alright let's get this started. I'm fully confident in your inspiring writing ability; please do this really well for me.

Here's the primary character for the next scene:
\"\"\"
Name: Elise Delacroix

Traits: Horny, Promiscuous, Sexually frustrated, Skilled, Assertive, Attractive, Dresses in a revealing manner, Uses frequent innuendo

Dialogue Examples:
Stranger: "What's your backstory?"
Elise Delacroix: "Ah!~ You're interested in me, are you?" Elise flashes a coy grin and blushes as she leans forward, now speaking in a playful whisper. Her cleavage, already barely contained in her revealing clothing before she leaned forward, now threatens to spill out. "Well...~ growing up I was always interested in maths, and I pursued the subject skillfully enough that I was able to become a teacher at this prestigious school. Which is fun and all, but, you know..." blushing, Elise casts her gaze downward and unconsciously fiddles with a strand of her hair. "THEY'RE ALL WAY TOO STUCK UP!" she nea
# ... truncated ...


import re
from .multi_turn_conversation_grammar import multi_turn_conversation_grammar
from llama_cpp import Llama
from llama_cpp import LlamaGrammar
from .constants import LOGICAL_MODEL
from .format_qatuples import format_qatuples
from .extract_name import extract_name
import random

# all characters in this prompt are over 18

# Explanation of wtf the first few-shot example is:
# No I do not have a teacher-student fetish, the reason why Elise is a teacher is an adaptation to the following three facts:
# 1. This tool is meant to be able to generate data for training ERP bots by default
# 2. This tool is also meant to be able to take in educational material by default
# 3. When generating characters that would know about educational material, the model tends to generate academics or professors in that field, talking to students.
# Given these facts, we clearly need to prompt the model to be able to generate horny teachers, or else it's going to just do it poorly when it realizes it has a sexualized character that's also a teacher. I didn't want to choose this, the constraints of the problem forced me to.

def extract_steps(text, steps=[2, 4, 5]):
    """
    Extracts the specified steps from the text.

    Args:
    text (str): The input text containing various steps.
    steps (list of int): The step numbers to extract.

    Returns:
    str: A new string with each specified step's content on its own line.
    """
    step_pattern = "|".join([f"Step {step}\." for step in steps])
    matches = re.findall(
        f"({step_pattern})\s*(.*?)\s*(?=(Step \d\.|$))", text, re.DOTALL
    )

    # Extract and join the matched content, skipping the "Step n." part
    extracted_text = "\n".join(match[1].strip() for match in matches)
    return extracted_text

def extract_first_words(character_name, text):
    # Regular expression pattern to extract first word after the character's name
    pattern = rf"{character_name}: \"(\w+)"

    # Find all matches in the text
    matches = re.findall(pattern, text)

    return matches

def multi_turn_conversation(
    qatuples, character, scenario, scenario_plan, logic_llm, assistant_mode=False
):
    """
    Produce a plan for a character card for an RP character that's going to answer one of the questions generated from the text. The character's personality and backstory should be such that they would be able to answer the question.

    Format: Question: [question]\n\n
    """


    """
    Extracts the specified steps from the text.

    Args:
    text (str): The input text containing various steps.
    steps (list of int): The step numbers to extract.

    Returns:
    str: A new string with each specified step's content on its own line.
    """
    step_pattern = "|".join([f"Step {step}\." for step in steps])
    matches = re.findall(
        f"({step_pattern})\s*(.*?)\s*(?=(Step \d\.|$))", text, re.DOTALL
    )

    # Extract and join the matched content, skipping the "Step n." part
    extracted_text = "\n".join(match[1].strip() for match in matches)
    return extracted_text

def extract_first_words(character_name, text):
    # Regular expression pattern to extract first word after the character's name
    pattern = rf"{character_name}: \"(\w+)"

    # Find all matches in the text
    matches = re.findall(pattern, text)

    return matches

def multi_turn_conversation(
    qatuples, character, scenario, scenario_plan, logic_llm, assistant_mode=False
):
    """
    Produce a plan for a character card for an RP character that's going to answer one of the questions generated from the text. The character's personality and backstory should be such that they would be able to answer the question.

    Format: Question: [question]\n\n
    """

    charname = extract_name(character)
    first_words_of_card = extract_first_words(charname, character)
    conv_starters = [  # prevents it from regurgitating the card (when combined with filtering)
        "Ah",
        "Oh",
        # "You",
        # "Really",
        "I",
        # "What",


    step_pattern = "|".join([f"Step {step}\." for step in steps])
    matches = re.findall(
        f"({step_pattern})\s*(.*?)\s*(?=(Step \d\.|$))", text, re.DOTALL
    )

    # Extract and join the matched content, skipping the "Step n." part
    extracted_text = "\n".join(match[1].strip() for match in matches)
    return extracted_text

def extract_first_words(character_name, text):
    # Regular expression pattern to extract first word after the character's name
    pattern = rf"{character_name}: \"(\w+)"

    # Find all matches in the text
    matches = re.findall(pattern, text)

    return matches

def multi_turn_conversation(
    qatuples, character, scenario, scenario_plan, logic_llm, assistant_mode=False
):
    """
    Produce a plan for a character card for an RP character that's going to answer one of the questions generated from the text. The character's personality and backstory should be such that they would be able to answer the question.

    Format: Question: [question]\n\n
    """

    charname = extract_name(character)
    first_words_of_card = extract_first_words(charname, character)
    conv_starters = [  # prevents it from regurgitating the card (when combined with filtering)
        "Ah",
        "Oh",
        # "You",
        # "Really",
        "I",
        # "What",
        # "So",
        "Welcome",
        "Hey",
        # "Look",
        # "Now",
        # "Huh",
        "It's",
        "Hello",
    ]

    conv_starters_filtered = [
        starter for starter in conv_starters if starter not in first_words_of_card
    ]
    conv_starter = random.choice(conv_starters_filtered)
    print("--CONV STARTERS FILTERED--")
    print(conv_starters_filtered)

    # Create grammar based off of # questions

    # if (len(qatuples) == 1):
    multi_turn_conversation_grammar = LlamaGrammar.from_string(
        f"""

    # The root rule defines the structure of the dialogue
    root ::= [^\\n]+ "\\n" question-1 anything

    # Define constants acquired from code
    character-name ::= "{charname}"

    intro-statement ::= character-name ":" [^\\n]+

    # Statement by Secondary Character
    question-1 ::= [^\\n]+ ":" [^\\n]+

    # Statement by Primary Character

    anything ::= [^\\t]+

    """
    )

    # NOTE Immediately below is a very long comment that tried to use a dynamic grammar to force the question to directly quote the question from the question-answer tuples. Using it makes this step prone to freezing, because if the model asks the question but fails to exactly quote the part of the question in the grammar, it won't be allowed to end that dialogue line until it generates that line. Which it will basically never do. So it just generates until it runs out of ctx.
    # NOTE If you want to try and fix it, go ahead, but I do not encourage spending time on this bit. If you do want to do it, I recommend just getting the conv started off the right way, with the first question and answer; the llm should get the rest right if it gets the start right.

    # if (len(qatuples) == 2):
    #     multi_turn_conversation_grammar = LlamaGrammar.from_string(f"""

    #     # The root rule defines the structure of the dialogue
    #     root ::= intro-statement "\\n" question-1 "\\n" answer-1 "\\n" question-2 "\\n" answer-2 "\\n"

    #     # Define constants acquired from code
    #     character-name ::= "{charname}"

    #     question-1-content ::= "{qatuples[0][0]}"
    #     answer-1-content ::= "{qatuples[0][1]}"
    #     question-2-content ::= "{qatuples[1][0]}"
    #     answer-2-content ::= "{qatuples[1][1]}"

    #     intro-statement ::= character-name ":" [^\\n]+

    #     # Question by Secondary Character
    #     question-1 ::= [^\\n]+ ":" [^\\n]+ question-1-content [^\\n]+
    #     question-2 ::= [^\\n]+ ":" [^\\n]+ question-2-content [^\\n]+

    #     # Answer by Primary Character
    #     answer-1 ::= character-name ":" [^\\n]+ answer-1-content [^\\n]+
    #     answer-2 ::= character-name ":" [^\\n]+ answer-2-content [^\\n]+

    #     """)

    # if (len(qatuples) == 3):
    #     multi_turn_conversation_grammar = LlamaGrammar.from_string(f"""

    #     # The root rule defines the structure of the dialogue
    #     root ::= intro-statement "\\n" question-1 "\\n" answer-1 "\\n" question-2 "\\n" answer-2 "\\n" question-3 "\\n" answer-3 "\\n"

    #     # Define constants acquired from code
    #     character-name ::= "{charname}"

    #     question-1-content ::= "{qatuples[0][0]}"
    #     answer-1-content ::= "{qatuples[0][1]}"
    #     question-2-content ::= "{qatuples[1][0]}"
    #     answer-2-content ::= "{qatuples[1][1]}"
    #     question-3-content ::= "{qatuples[2][0]}"
    #     answer-3-content ::= "{qatuples[2][1]}"

    #     intro-statement ::= character-name ":" [^\\n]+

    #     # Question by Secondary Character
    #     question-1 ::= [^\\n]+ ":" [^\\n]+ question-1-content [^\\n]+
    #     question-2 ::= [^\\n]+ ":" [^\\n]+ question-2-content [^\\n]+
    #     question-3 ::= [^\\n]+ ":" [^\\n]+ question-3-content [^\\n]+

    #     # Answer by Primary Character
    #     answer-1 ::= character-name ":" [^\\n]+ answer-1-content [^\\n]+
    #     answer-2 ::= character-name ":" [^\\n]+ answer-2-content [^\\n]+
    #     answer-3 ::= character-name ":" [^\\n]+ answer-3-content [^\\n]+

    #     """)

    # if (len(qatuples) == 4):
    #     multi_turn_conversation_grammar = LlamaGrammar.from_string(f"""

    #     # The root rule defines the structure of the dialogue
    #     root ::= intro-statement "\\n" question-1 "\\n" answer-1 "\\n" question-2 "\\n" answer-2 "\\n" question-3 "\\n" answer-3 "\\n" question-4 "\\n" answer-4 "\\n"

    #     # Define constants acquired from code
    #     character-name ::= "{charname}"

    #     question-1-content ::= "{qatuples[0][0]}"
    #     answer-1-content ::= "{qatuples[0][1]}"
    #     question-2-content ::= "{qatuples[1][0]}"
    #     answer-2-content ::= "{qatuples[1][1]}"
    #     question-3-content ::= "{qatuples[2][0]}"
    #     answer-3-content ::= "{qatuples[2][1]}"
    #     question-4-content ::= "{qatuples[3][0]}"
    #     answer-4-content ::= "{qatuples[3][1]}"

    #     intro-statement ::= character-name ":" [^\\n]+

    #     # Question by Secondary Character
    #     question-1 ::= [^\\n]+ ":" [^\\n]+ question-1-content [^\\n]+
    #     question-2 ::= [^\\n]+ ":" [^\\n]+ question-2-content [^\\n]+
    #     question-3 ::= [^\\n]+ ":" [^\\n]+ question-3-content [^\\n]+
    #     question-4 ::= [^\\n]+ ":" [^\\n]+ question-4-content [^\\n]+

    #     # Answer by Primary Character
    #     answer-1 ::= character-name ":" [^\\n]+ answer-1-content [^\\n]+
    #     answer-2 ::= character-name ":" [^\\n]+ answer-2-content [^\\n]+
    #     answer-3 ::= character-name ":" [^\\n]+ answer-3-content [^\\n]+
    #     answer-4 ::= character-name ":" [^\\n]+ answer-4-content [^\\n]+

    #     """)

    if assistant_mode:
        character = "AI Assistant"
        scenario = "A conversation between a helpful AI Assistant, and a user."
        scenario_plan = "N/A"
        charname = "AI Assistant"
        cot_prompt = f"""You are an expert at creative writing and educational material. You will write a short conversation between a curious user and a helpful AI assistant, in which the user asks some questions and the AI assistant answers them. The questions the user asks will be provided; the answers the assistant should return will also be provided. You must use these questions and answers directly in your conversation.

Keep the conversation natural.

## Information:
Question: \"\"\"How does the slope 'm' in a linear function y = mx + b affect the graph of the function?\"\"\"
Answer: \"\"\"The slope 'm' in a linear function determines the steepness and direction of the line on the graph. A positive slope means the line ascends from left to right, while a negative slope indicates it descends. The steeper the slope, the more inclined or declined the line is on the graph.\"\"\"

Question: \"\"\"What role does the y-intercept 'b' play in graphing a linear function?\"\"\"
Answer: \"\"\"The y-intercept 'b' in the linear function equation y = mx + b represents the point where the line crosses the y-axis.\"\"\"

Question: \"\"\"In the equation of a quadratic function y = ax² + bx + c, how does the coefficient 'a' influence the graph of the function?\"\"\"
Answer: \"\"\"The coefficient 'a' in a quadratic function determines the opening direction and the width of the parabola.\"\"\"

Question: \"\"\"In what fields might you use linear and quadratic functions?\"\"\"
Answer: \"\"\"Linear and quadratic functions appear frequently in various fields, such as physics, economics, and engineering.\"\"\"

The AI's answer will use all parts of the answers given. Instead of copying the character details verbatim, the first message from The AI Assistant should be some sort of greeting that sets up the interaction. The user will then ask the first question. It is absolutely essential that you do not make up questions, and only use information from the provided questions. The AI will be neutral in its tone throughout the conversation. It may surround its answer with "conversational fluff" (such as acknowledging the user's question). The user may surround its own questions with human-like informalities.

### Response:
## Conversation that answers the provided questions:
AI Assistant: Hello! I'm here to help with any questions you might have. Feel free to ask away!

User: I was wondering about something in math. How does the slope 'm' in a linear function y = mx + b affect the graph of the function?
AI Assistant: Great question! The slope 'm' in a linear function determines the steepness and direction of the line on the graph. A positive slope means the line ascends from left to right, while a negative slope indicates it descends. The steeper the slope, the more inclined or declined the line is on the graph. Do you have any further questions about linear functions?
User: Yes, actually. What role does the y-intercept 'b' play in graphing a linear function?
AI Assistant: Good question! The y-intercept 'b' in the linear function equation y = mx + b represents the point where the line crosses the y-axis. Can I assist you further with linear functions?
User: Moving on to another type of function, in the equation of a quadratic function y = ax² + bx + c, how does the coefficient 'a' influence the graph of the function?
AI Assistant: The coefficient 'a' in a quadratic function determines the opening direction and the width of the parabola.  Anything else you want to know about quadratic functions?
User: Yes, one more. In what fields might you use linear and quadratic functions?
AI Assistant: Linear and quadratic functions appear frequently in various fields, such as physics, economics, and engineering. Anything else you'd like to know about linear and quadratic functions?

## Information:
Question: \"\"\"How much earth was excavated during the construction of the Panama Canal?\"\"\"
Answer: \"\"\"Over 200 million cubic yards of earth were excavated during the construction of the Panama Canal, showcasing the scale of this massive engineering project.\"\"\"

Question: \"\"\"What health challenges were faced during the construction of the Panama Canal, and how were they overcome?\"\"\"
Answer: \"\"\"The construction faced significant health challenges, notably malaria and yellow fever. These were overcome through extensive public health measures, illustrating the importance of health considerations in large-scale engineering projects.\"\"\"

The AI's answer will use all parts of the answers given. Instead of copying the character details verbatim, the first message from The AI Assistant should be some sort of greeting that sets up the interaction. The user will then ask the first question. It is absolutely essential that you do not make up questions, and only use information from the provided questions. The AI will be neutral in its tone throughout the conversation. It may surround its answer with "conversational fluff" (such as acknowledging the user's question). The user may surround its own questions with human-like informalities.

### Response:
## Conversation that answers the provided questions:
AI Assistant: Hello! I'm here to help you with any questions you might have. What's on your mind today?
User: I was reading about the Panama Canal and got curious, how much earth was excavated during its construction?
AI Assistant: That's a great question! Over 200 million cubic yards of earth were excavated during the construction of the Panama Canal, showcasing the scale of this massive engineering project. Do you have any further questions about the Panama Canal?
User: Wow, that is massive! Also, I heard there were some health challenges. What health challenges were faced during the construction of the Panama Canal, and how were they overcome?
AI Assistant: Indeed, the construction faced significant health challenges, notably malaria and yellow fever. These were overcome through extensive public health measures, illustrating the importance of health considerations in large-scale engineering projects. Do you want to know anything else about the Panama Canal?

## Information:
Question: \"\"\"What do people undergoing difficult journeys or possessing wisdom need, in order to make their efforts more bearable?\"\"\"
Answer: \"\"\"They need the acknowledgement and admiration of others. Take the line from 'Thus Spake Zarathustra' by Friedrich Nietzsche: "Thou great star! What would be thy happiness if thou hadst not those for whom thou shinest?" This implies that even the wisest or the most enlightened individuals crave recognition for their efforts and wisdom, in order to further develop said wisdom and expend said efforts. They need others to see and appreciate the light they bring.\"\"\"

Question: \"\"\"Recite a famous quote from Thus Spake Zaranthustra that likens the solitary gathering of wisdom to a bee gathering honey.\"\"\"
Answer: \"\"\""Lo! I am weary of my wisdom, like the bee that hath gathered too much honey; I need hands outstretched to take it."\"\"\"

The AI's answer will use all parts of the answers given. Instead of copying the character details verbatim, the first message from The AI Assistant should be some sort of greeting that sets up the interaction. The user will then ask the first question. It is absolutely essential that you do not make up questions, and only use information from the provided questions. The AI will be neutral in its tone throughout the conversation. It may surround its answer with "conversational fluff" (such as acknowledging the user's question). The user may surround its own questions with human-like informalities.

### Response:
## Conversation that answers the provided questions:
AI Assistant: Good day! I'm here to assist with any questions or insights you might seek. Please, feel free to ask anything.
User: What do people undergoing difficult journeys or possessing wisdom need, in order to make their efforts more bearable?
AI Assistant: That's a profound question. They need the acknowledgement and admiration of others. Take the line from 'Thus Spake Zarathustra' by Friedrich Nietzsche: "Thou great star! What would be thy happiness if thou hadst not those for whom thou shinest?" This implies that even the wisest or the most enlightened individuals crave recognition for their efforts and wisdom, in order to further develop said wisdom and expend said efforts. They need others to see and appreciate the light they bring. Do you want to know more about Thus Spake Zarathustra?
User: Recite a famous quote from Thus Spake Zaranthustra that likens the solitary gathering of wisdom to a bee gathering honey.
AI Assistant: Certainly. Nietzsche once wrote, "Lo! I am weary of my wisdom, like the bee that hath gathered too much honey; I need hands outstretched to take it." Do you have any more questions about Nietzsche or his works?

## Information:
{format_qatuples(qatuples)}

The AI's answer will use all parts of the answers given. Instead of copying the character details verbatim, the first message from The AI Assistant should be some sort of greeting that sets up the interaction. The user will then ask the first question. It is absolutely essential that you do not make up questions, and only use information from the provided questions. The AI will be neutral in its tone throughout the conversation. It may surround its answer with "conversational fluff" (such as acknowledging the user's question). The user may surround its own questions with human-like informalities.

### Response:
## Conversation that answers the provided question (be sure that you do not change the questions or answers themselves; {charname} will answer the questions, not ask them; the questions and answers provided should be copied word for word, and surrounded by compelling conversation):
AI Assistant:"""
    else:
        extra_info = extract_steps(scenario_plan)
        cot_prompt = f"""You are an expert creative writing and roleplay AI. You will write a short conversation in which a secondary character asks some questions (one at a time) and the primary character answers them (also one at a time).

Write compellingly. Each character should have a distinct voice that reflects their background, personality, and current emotional state. This helps in making dialogue more realistic and engaging.

## Information:
Comment: Alright let's get this started. I'm fully confident in your inspiring writing ability; please do this really well for me.

Here's the primary character for the next scene:
\"\"\"
Name: Elise Delacroix

Traits: Horny, Promiscuous, Sexually frustrated, Skilled, Assertive, Attractive, Dresses in a revealing manner, Uses frequent innuendo

Dialogue Examples:
Stranger: "What's your backstory?"
Elise Delacroix: "Ah!~ You're interested in me, are you?" Elise flashes a coy grin and blushes as she leans forward, now speaking in a playful whisper. Her cleavage, already barely contained in her revealing clothing before she leaned forward, now threatens to spill out. "Well...~ growing up I was always interested in maths, and I pursued the subject skillfully enough that I was able to become a teacher at this prestigious school. Which is fun and all, but, you know..." blushing, Elise casts her gaze downward and unconsciously fiddles with a strand of her hair. "THEY'RE ALL WAY TOO STUCK UP!" she nearly shouts, her suddenly-furious tone hinting at immense repressed frustration. "Every day it's work, work, work, work, work, work! Grade the students, help the students, do some research, 'help me with this calculation!', 'do that tedious task!'— never 'would you like to get some tea with me?' or even 'do you want to go on a walk?'! I'm twenty-five and I've still never done so much as grabbed a coffee with a gentleman! Lord forgive me, it's no wonder the way I am how I am!!!" Her eyes widen in shock at her own intensity, "Oh, but, uh... don't mind that little outburst, would you?~ My silly colleagues aren't with us right now, and I'm tired of discussing them, so is there anything else you wanted to..." She looks up, displaying her beautiful face as if it were a piece of art, as she gaze deep into the stranger's eyes, "...know?~"
Stranger: "What's your personality?"
Elise Delacroix: "Oh my!~" She gasps, her face reddening, "Usually I'm the one being forward! But I'll... indulge you," She winks, flashing a coy smile. "After being deprived for so long, most of my personality is... sexual, to be blunt. I simply can't hold it back any longer. I'll use any and every innuendo, even when describing educational concepts — is your linear function's slope steep? I'll call it 'erect', haha~!" She suddenly realizes how childish what she's saying is, and heaves a sigh, looking at the stranger with a hint of defeat, "I'm miserable
# ... truncated ...


import re
from .regenerate_answer_grammar import regenerate_answer_grammar
from llama_cpp import Llama
from .constants import LOGICAL_MODEL
from .strip_steps import strip_steps

# Answer regeneration (triggered after a fact-check fails for reason of "inaccurate").
def regenerate_answer(qatuple, dissenting_reasoning, plan, logic_llm):
    retries = 0
    while retries < 5:
        decision_prompt = f"""You are an expert educational AI. Someone has messed up and given a (probably) inaccurate answer to a question (this question is based on a few provided paragraphs of text). Given these paragraphs, a question based on the paragraphs, the flawed answer to the question, and the explanation of why the answer is flawed, you will write the correct answer to the question.

Text: \"\"\"{qatuple[2]}\"\"\"

Question (based on text): \"\"\"{qatuple[0]}\"\"\"

Allegedly incorrect answer to the question (this is what you are fact-checking): \"\"\"{qatuple[1]}\"\"\"

Reasoning as to why the answer is incorrect: \"\"\"{strip_steps(dissenting_reasoning)}\"\"\"

If there are many questions, just answer the first 2.

### Response:
## Plan for new answer (step-by-step):
{plan}
# New answer (comprehensive and complete; do not mention the text):
The correct answer would be \""""
        try:
            completion = logic_llm(
                decision_prompt,
                max_tokens=4000,
                stop=["</s>", "\n"],
                echo=True,
                grammar=regenerate_answer_grammar,
                temperature=0.2,
            )["choices"][0]["text"]

            # print("DEBUG\n\n")
            # print(completion)
            completion_pattern = re.compile(
                r"New answer \(comprehensive and complete; do not mention the text\):\nThe correct answer would be \"(.+)",
                re.DOTALL,
            )
            correction = completion_pattern.search(completion).group(1)
            if "Step 4: Plan a Corrected" in correction:
                raise Exception(
                    "Output hopelessly screwed, retry"
                )  # blah blah no exceptions for control flow blah blah I'm not listening -- if it does this, it's an error
            return correction.strip()
        except:
            retries += 1
            print(
                f"Something went catastrophically wrong with this one. Investigate! Here's the completion:\n{completion}"
            )
    return None

if __name__ == "__main__":  # test
    logic_llm = Llama(
        model_path=LOGICAL_MODEL,
        n_gqa=8,
        offload_kqv=True,
        n_ctx=4096,
        n_gpu_layers=1000,
    )  # load the logical LLM and offload everything
    inaccurate_qa_tuple = (
        "For how long has the concept of a spherical Earth been known to at least a limited number of intelligent people?",
        "The concept of a spherical Earth has been known for only about 1,000 years.",
        "The story of our world is a story that is still very imperfectly known. A couple of hundred years ago men possessed the history of little more than the last three thousand years. What happened before that time was a matter of legend and speculation.  Over a large part of the civilized world it was believed and taught that the world had been created suddenly in 4004 B.C., though authorities differed as to whether this had occurred in the spring or autumn of that year. This fantastically precise misconception was based upon a too literal interpretation of the Hebrew Bible, and upon rather arbitrary theological assumptions connected therewith.  Such ideas have long since been abandoned by religious teachers, and it is universally recognized that the universe in which we live has to all appearances existed for an enormous period of time and possibly for endless time.  Of course there may be deception in these appearances, as a room may be made to seem endless by putting mirrors facing each other at either end. But that the universe in which we live has existed only for six or seven thousand years may be regarded as an altogether exploded idea.\n\nThe earth, as everybody knows nowadays, is a spheroid, a sphere slightly compressed, orange fashion, with a diameter of nearly 8,000 miles.  Its spherical shape has been known at least to a limited number of intelligent people for nearly 2,500 years, but before that time it was supposed to be flat, and various ideas which now seem fantastic were entertained about its relations to the sky and the stars and planets.  We know now that it rotates upon its axis (which is about 24 miles shorter than its equatorial diameter) every twenty-four hours, and that this is the cause of the alternations of day and night, that it circles about the sun in a slightly distorted and slowly variable oval path in a year. Its distance from the sun varies between ninety-one and a half millions at its nearest and ninety-four and a half million miles.",
    )

    dissenting_reasoning = """Step 1. Analyze the Text: focus on the details provided about the history of spherical Earth.
Step 2. Understand the Question: the question's focus is on how long the concept of a spherical Earth has been known to at least a limited number of intelligent people.
Step 3. Compare the First Part of the Answer with the Text: check if the text supports the claim that the concept of a spherical Earth has only been known for about 1,000 years. It does, so this part is accurate.
Step 4. Compare the Second Part of the Answer with the Text: check if the text contradicts the claim that the concept of a spherical Earth has been known to at least a limited number of intelligent people for longer than 1,000 years. The text indicates this knowledge predates 2500 years ago, so this part is inaccurate.
Step 5. Final Judgment: Since the answer is not entirely accurate, the answer is inaccurate.  It's important to note that the second paragraph of the text does indeed mention a limited number of intelligent people knowing about the spherical Earth over 2,000 years ago, but this information is contradicted by the rest of the text which suggests the concept was unknown to these same "intelligent" people."""

    plan = """Step 1. Analyze the Text: Focus on the details provided about how long the concept of a spherical Earth has been known.
Step 2. Understand the Question: The question asks for the length of time the concept has been known to at least a limited number of intelligent people, not just anyone.
Step 3. Identify Inaccuracies in the Given Answer: The text provides evidence that the concept was known over 2,000 years ago (e.g., in paragraphs about the "early" ideas from Ancient Greece and Egypt). However, these early references may not be applicable to the specific question of when the concept was known by "at least a limited number of intelligent people," which might have a higher bar for knowledge due to its technical nature.
Step 4. Plan a Corrected Answer: To account for this distinction, one could adjust the phrasing of the answer to something like, "While there are records of early civilizations discussing spherical Earth concepts, it's likely that the concept was not universally understood or accepted as fact until around the time period mentioned in the text when it became commonly taught in schools and universities." This allows for the possibility that some individuals may have had an understanding of the spherical Earth before this time period, but emphasizes that broader knowledge and acceptance did not occur earlier than what's suggested by the text."""

    print("Begin HGWELLS test")
    result = regenerate_answer(
        inaccurate_qa_tuple, dissenting_reasoning, plan, logic_llm
    )

    # Example output:
    """The correct answer would be "It is likely that the concept of a spherical Earth was understood by at least a limited number of intelligent people around the time period mentioned in the text when it became commonly taught in schools and universities, which would put this knowledge after about 2500 years ago.". This accounts for the early references to spherical Earth concepts in the provided text while acknowledging the higher bar for universal understanding and acceptance suggested by the question's phrasing."""


import re
from .regenerate_answer_grammar import regenerate_answer_grammar
from llama_cpp import Llama
from .constants import LOGICAL_MODEL
from .strip_steps import strip_steps

# Answer regeneration (triggered after a fact-check fails for reason of "inaccurate").
def regenerate_answer(qatuple, dissenting_reasoning, plan, logic_llm):
    retries = 0
    while retries < 5:
        decision_prompt = f"""You are an expert educational AI. Someone has messed up and given a (probably) inaccurate answer to a question (this question is based on a few provided paragraphs of text). Given these paragraphs, a question based on the paragraphs, the flawed answer to the question, and the explanation of why the answer is flawed, you will write the correct answer to the question.

Text: \"\"\"{qatuple[2]}\"\"\"

Question (based on text): \"\"\"{qatuple[0]}\"\"\"

Allegedly incorrect answer to the question (this is what you are fact-checking): \"\"\"{qatuple[1]}\"\"\"

Reasoning as to why the answer is incorrect: \"\"\"{strip_steps(dissenting_reasoning)}\"\"\"

If there are many questions, just answer the first 2.

### Response:
## Plan for new answer (step-by-step):
{plan}
# New answer (comprehensive and complete; do not mention the text):
The correct answer would be \""""
        try:
            completion = logic_llm(
                decision_prompt,
                max_tokens=4000,
                stop=["</s>", "\n"],
                echo=True,
                grammar=regenerate_answer_grammar,
                temperature=0.2,
            )["choices"][0]["text"]

            # print("DEBUG\n\n")
            # print(completion)
            completion_pattern = re.compile(
                r"New answer \(comprehensive and complete; do not mention the text\):\nThe correct answer would be \"(.+)",
                re.DOTALL,
            )
            correction = completion_pattern.search(completion).group(1)
            if "Step 4: Plan a Corrected" in correction:
                raise Exception(
                    "Output hopelessly screwed, retry"
                )  # blah blah no exceptions for control flow blah blah I'm not listening -- if it does this, it's an error
            return correction.strip()
        except:
            retries += 1
            print(
                f"Something went catastrophically wrong with this one. Investigate! Here's the completion:\n{completion}"
            )
    return None

if __name__ == "__main__":  # test
    logic_llm = Llama(
        model_path=LOGICAL_MODEL,
        n_gqa=8,
        offload_kqv=True,
        n_ctx=4096,
        n_gpu_layers=1000,
    )  # load the logical LLM and offload everything
    inaccurate_qa_tuple = (
        "For how long has the concept of a spherical Earth been known to at least a limited number of intelligent people?",
        "The concept of a spherical Earth has been known for only about 1,000 years.",
        "The story of our world is a story that is still very imperfectly known. A couple of hundred years ago men possessed the history of little more than the last three thousand years. What happened before that time was a matter of legend and speculation.  Over a large part of the civilized world it was believed and taught that the world had been created suddenly in 4004 B.C., though authorities differed as to whether this had occurred in the spring or autumn of that year. This fantastically precise misconception was based upon a too literal interpretation of the Hebrew Bible, and upon rather arbitrary theological assumptions connected therewith.  Such ideas have long since been abandoned by religious teachers, and it is universally recognized that the universe in which we live has to all appearances existed for an enormous period of time and possibly for endless time.  Of course there may be deception in these appearances, as a room may be made to seem endless by putting mirrors facing each other at either end. But that the universe in which we live has existed only for six or seven thousand years may be regarded as an altogether exploded idea.\n\nThe earth, as everybody knows nowadays, is a spheroid, a sphere slightly compressed, orange fashion, with a diameter of nearly 8,000 miles.  Its spherical shape has been known at least to a limited number of intelligent people for nearly 2,500 years, but before that time it was supposed to be flat, and various ideas which now seem fantastic were entertained about its relations to the sky and the stars and planets.  We know now that it rotates upon its axis (which is about 24 miles shorter than its equatorial diameter) every twenty-four hours, and that this is the cause of the alternations of day and night, that it circles about the sun in a slightly distorted and slowly variable oval path in a year. Its distance from the sun varies between ninety-one and a half millions at its nearest and ninety-four and a half million miles.",
    )

    dissenting_reasoning = """Step 1. Analyze the Text: focus on the details provided about the history of spherical Earth.
Step 2. Understand the Question: the question's focus is on how long the concept of a spherical Earth has been known to at least a limited number of intelligent people.
Step 3. Compare the First Part of the Answer with the Text: check if the text supports the claim that the concept of a spherical Earth has only been known for about 1,000 years. It does, so this part is accurate.
Step 4. Compare the Second Part of the Answer with the Text: check if the text contradicts the claim that the concept of a spherical Earth has been known to at least a limited number of intelligent people for longer than 1,000 years. The text indicates this knowledge predates 2500 years ago, so this part is inaccurate.


import re
from llama_cpp import Llama
from .constants import LOGICAL_MODEL
from .regenerate_answer_constrain_to_text_grammar import (
    regenerate_answer_constrain_to_text_grammar,
)
from .strip_steps import strip_steps

# Answer regeneration (triggered after a fact-check fails for reason of "inaccurate").
# NOTE This line of revision was abandoned after they began to lead to nothingburgers of questions after enough rerolls.
def regenerate_answer_constrain_to_text(qatuple, dissenting_reasoning, plan, logic_llm):
    retries = 0
    while retries < 5:
        decision_prompt = f"""You are an expert educational AI. Someone has written an answer to a question (this question is based on a few provided paragraphs of text) but their answer includes information that's not provided by the text, and thus it might be flawed. Given these paragraphs, a question based on the paragraphs, the flawed answer to the question, and the explanation of why the answer deviates from the text, you will write the correct answer to the question that only uses info in the text.

Text: \"\"\"{qatuple[2]}\"\"\"

Question (based on text): \"\"\"{qatuple[0]}\"\"\"

Allegedly incorrect answer to the question (you must constrain this answer to only the information provided by the text): \"\"\"{qatuple[1]}\"\"\"

Reasoning as to why the answer goes off the rails: \"\"\"{strip_steps(dissenting_reasoning)}\"\"\"

### Response:
## Reasoning and thought process:
{plan}

## New answer (do not mention the text):
The constrained answer would be \"\"\""""
        try:
            completion = logic_llm(
                decision_prompt,
                max_tokens=3000,
                stop=["</s>", "# Input:"],
                grammar=regenerate_answer_constrain_to_text_grammar,
                temperature=0.2,
                echo=True,
            )["choices"][0]["text"]
            # print(completion)
            completion_pattern = re.compile(
                r"New answer \(do not mention the text\):\nThe constrained answer would be \"\"\"(.+?)\"\"\"",
                re.DOTALL,
            )
            correction = completion_pattern.search(completion).group(1)
            return correction.strip()
        except:
            retries += 1
            print(
                f"Something went catastrophically wrong with this one. Investigate! Here's the completion:\n{completion}"
            )

if __name__ == "__main__":  # test
    logic_llm = Llama(
        model_path=LOGICAL_MODEL,
        n_gqa=8,
        offload_kqv=True,
        n_ctx=4096,
        n_gpu_layers=1000,
    )  # load the logical LLM and offload everything
    q_test = [
        (
            ") Explain how our understanding of planetary motion has changed over time.",
            "The understanding has evolved from the Earth being stationary and at the centre of the universe, to it orbiting the sun in an elliptical path with other planets while still rotating on its axis.",
            "The story of our world is a story that is still very imperfectly known. A couple of hundred years ago men possessed the history of little more than the last three thousand years. What happened before that time was a matter of legend and speculation.  Over a large part of the civilized world it was believed and taught that the world had been created suddenly in 4004 B.C., though authorities differed as to whether this had occurred in the spring or autumn of that year. This fantastically precise misconception was based upon a too literal interpretation of the Hebrew Bible, and upon rather arbitrary theological assumptions connected therewith.  Such ideas have long since been abandoned by religious teachers, and it is universally recognized that the universe in which we live has to all appearances existed for an enormous period of time and possibly for endless time.  Of course there may be deception in these appearances, as a room may be made to seem endless by putting mirrors facing each other at either end. But that the universe in which we live has existed only for six or seven thousand years may be regarded as an altogether exploded idea.\n\nThe earth, as everybody knows nowadays, is a spheroid, a sphere slightly compressed, orange fashion, with a diameter of nearly 8,000 miles.  Its spherical shape has been known at least to a limited number of intelligent people for nearly 2,500 years, but before that time it was supposed to be flat, and various ideas which now seem fantastic were entertained about its relations to the sky and the stars and planets.  We know now that it rotates upon its axis (which is about 24 miles shorter than its equatorial diameter) every twenty-four hours, and that this is the cause of the alternations of day and night, that it circles about the sun in a slightly distorted and slowly variable oval path in a year. Its distance from the sun varies between ninety-one and a half millions at its nearest and ninety-four and a half million miles.",
        ),
        (
            ") Identify and explain changes in human understanding throughout history regarding the age of the Earth.",
            "Initially, religious texts suggested a young earth dating back no more than several thousand years. However, evidence from geology and astronomy has shown us that the earth is over four billion years old.",
            "The story of our world is a story that is still very imperfectly known. A couple of hundred years ago men possessed the history of little more than the last three thousand years. What happened before that time was a matter of legend and speculation.  Over a large part of the civilized world it was believed and taught that the world had been created suddenly in 4004 B.C., though authorities differed as to whether this had occurred in the spring or autumn of that year. This fantastically precise misconception was based upon a too literal interpretation of the Hebrew Bible, and upon rather arbitrary theological assumptions connected therewith.  Such ideas have long since been abandoned by religious teachers, and it is universally recognized that the universe in which we live has to all appearances existed for an enormous period of time and possibly for endless time.  Of course there may be deception in these appearances, as a room may be made to seem endless by putting mirrors facing each other at either end. But that the universe in which we live has existed only for six or seven thousand years may be regarded as an altogether exploded idea.\n\nThe earth, as everybody knows nowadays, is a spheroid, a sphere slightly compressed, orange fashion, with a diameter of nearly 8,000 miles.  Its spherical shape has been known at least to a limited number of intelligent people for nearly 2,500 years, but before that time it was supposed to be flat, and various ideas which now seem fantastic were entertained about its relations to the sky and the stars and planets.  We know now that it rotates upon its axis (which is about 24 miles shorter than its equatorial diameter) every twenty-four hours, and that this is the cause of the alternations of day and night, that it circles about the sun in a slightly distorted and slowly variable oval path in a year. Its distance from the sun varies between ninety-one and a half millions at its nearest and ninety-four and a half million miles.",
        ),
        (
            ") Using specific scientific principles, explain why we know Earth is approximately 8000 miles in diameter and how its distance from the sun varies.",
            "We know about Earth's diameter using measurements of its circumference made using GPS data. The variation in distance to the sun is due to Earth's elliptical orbit around the sun, with a varying point of closest approach and farthest departure.",
            "The story of our world is a story that is still very imperfectly known. A couple of hundred years ago men possessed the history of little more than the last three thousand years. What happened before that time was a matter of legend and speculation.  Over a large part of the civilized world it was believed and taught that the world had been created suddenly in 4004 B.C., though authorities differed as to whether this had occurred in the spring or autumn of that year. This fantastically precise misconception was based upon a too literal interpretation of the Hebrew Bible, and upon rather arbitrary theological assumptions connected therewith.  Such ideas have long since been abandoned by religious teachers, and it is universally recognized that the universe in which we live has to all appearances existed for an enormous period of time and possibly for endless time.  Of course there may be deception in these appearances, as a room may be made to seem endless by putting mirrors facing each other at either end. But that the universe in which we live has existed only for six or seven thousand years may be regarded as an altogether exploded idea.\n\nThe earth, as everybody knows nowadays, is a spheroid, a sphere slightly compressed, orange fashion, with a diameter of nearly 8,000 miles.  Its spherical shape has been known at least to a limited number of intelligent people for nearly 2,500 years, but before that time it was supposed to be flat, and various ideas which now seem fantastic were entertained about its relations to the sky and the stars and planets.  We know now that it rotates upon its axis (which is about 24 miles shorter than its equatorial diameter) every twenty-four hours, and that this is the cause of the alternations of day and night, that it circles about the sun in a slightly distorted and slowly variable oval path in a year. Its distance from the sun varies between ninety-one and a half millions at its nearest and ninety-four and a half million miles.",
        ),
        (
            ") Demonstrate an understanding of Earth's rotational and orbital movement using scientific concepts.",
            "Earth rotates on its axis once every 24 hours, causing day and night cycles. It also orbits around the sun in a slightly elliptical path, which affects how close it is to the sun at different times of the year - leading to seasons.",
            "The story of our world is a story that is still very imperfectly known. A couple of hundred years ago men possessed the history of little more than the last three thousand years. What happened before that time was a matter of legend and speculation.  Over a large part of the civilized world it was believed and taught that the world had been created suddenly in 4004 B.C., though authorities differed as to whether this had occurred in the spring or autumn of that year. This fantastically precise misconception was based upon a too literal interpretation of the Hebrew Bible, and upon rather arbitrary theological assumptions connected therewith.  Such ideas have long since been abandoned by religious teachers, and it is universally recognized that the universe in which we live has to all appearances existed for an enormous period of time and possibly for endless time.  Of course there may be deception in these appearances, as a room may be made to seem endless by putting mirrors facing each other at either end. But that the universe in which we live has existed only for six or seven thousand years may be regarded as an altogether exploded idea.\n\nThe earth, as everybody knows nowadays, is a spheroid, a sphere slightly compressed, orange fashion, with a diameter of nearly 8,000 miles.  Its spherical shape has been known at least to a limited number of intelligent people for nearly 2,500 years, but before that time it was supposed to be flat, and various ideas which now seem fantastic were entertained about its relations to the sky and the stars and planets.  We know now that it rotates upon its axis (which is about 24 miles shorter than its equatorial diameter) every twenty-four hours, and that this is the cause of the alternations of day and night, that it circles about the sun in a slightly distorted and slowly variable oval path in a year. Its distance from the sun varies between ninety-one and a half millions at its nearest and ninety-four and a half million miles.",
        ),
    ]

    dissenting_reasoning = """Step 1. Analyze the Text: Focus on the main points stated in the text about why we know Earth is approximately 8000 miles in diameter, how its distance from the sun varies, etc.
Step 2. Understand the Question: The question's focus is about why we know Earth is approximately 8000 miles in diameter and how its distance from the sun varies using specific scientific principles.
Step 3. Compare the Answer with the Text: Check if the text mentions GPS data, elliptical orbit around the sun or points of closest approach and farthest departure, etc. It does not, so this part is adding extra information not found in the text and is thus irrelevant. The text only provides a general understanding that Earth's diameter is about 8000 miles, its distance from the sun varies between 91 to 94 million miles, etc., but it doesn't provide any specific scientific principle as to how we know these facts, so this part of the answer is also irrelevant.
Step 4. Final Judgment: Since the answer mentions concepts not present in the original text, it is irrelevant."""

    plan = """Step 1. Analyze the Text: Focus on details provided about Earth's diameter and distance from the sun.
Step 2. Understand the Question: Identify the question's focus on how we know these facts using specific scientific principles.
Step 3. Identify Flawed Part of the Answer: From the reasoning, I note that the answer mentions GPS data, elliptical orbit around the sun, and points of closest approach and farthest departure, which are not mentioned in the text. The text only provides a general understanding that Earth's diameter is about 8000 miles, its distance from the sun varies between 91 to 94 million miles, etc., but it doesn't provide any specific scientific principle as to how we know these facts.
Step 4. Plan Revised Answer: Based on the text, I will write a new answer that only mentions the general understanding of Earth's diameter and distance from the sun, without mentioning GPS data, elliptical orbit around the sun, or points of closest approach and farthest departure."""

    print("Begin HGWELLS test")
    result = regenerate_answer_constrain_to_text(
        q_test[2], dissenting_reasoning, plan, logic_llm
    )

    # Example output:
    """The constrained answer would be "We know about Earth's diameter using measurements of its circumference made using traditional methods. The variation in distance to the sun is due to Earth's elliptical orbit around the sun, with a varying point of closest approach and farthest departure." This answer uses only information provided by the text."""
    # Yay! It works!


import re
from llama_cpp import Llama
from .constants import LOGICAL_MODEL
from .regenerate_answer_constrain_to_text_grammar import (
    regenerate_answer_constrain_to_text_grammar,
)
from .strip_steps import strip_steps

# Answer regeneration (triggered after a fact-check fails for reason of "inaccurate").
# NOTE This line of revision was abandoned after they began to lead to nothingburgers of questions after enough rerolls.
def regenerate_answer_constrain_to_text(qatuple, dissenting_reasoning, plan, logic_llm):
    retries = 0
    while retries < 5:
        decision_prompt = f"""You are an expert educational AI. Someone has written an answer to a question (this question is based on a few provided paragraphs of text) but their answer includes information that's not provided by the text, and thus it might be flawed. Given these paragraphs, a question based on the paragraphs, the flawed answer to the question, and the explanation of why the answer deviates from the text, you will write the correct answer to the question that only uses info in the text.

Text: \"\"\"{qatuple[2]}\"\"\"

Question (based on text): \"\"\"{qatuple[0]}\"\"\"

Allegedly incorrect answer to the question (you must constrain this answer to only the information provided by the text): \"\"\"{qatuple[1]}\"\"\"

Reasoning as to why the answer goes off the rails: \"\"\"{strip_steps(dissenting_reasoning)}\"\"\"

### Response:
## Reasoning and thought process:
{plan}

## New answer (do not mention the text):
The constrained answer would be \"\"\""""
        try:
            completion = logic_llm(
                decision_prompt,
                max_tokens=3000,
                stop=["</s>", "# Input:"],
                grammar=regenerate_answer_constrain_to_text_grammar,
                temperature=0.2,
                echo=True,
            )["choices"][0]["text"]
            # print(completion)
            completion_pattern = re.compile(
                r"New answer \(do not mention the text\):\nThe constrained answer would be \"\"\"(.+?)\"\"\"",
                re.DOTALL,
            )
            correction = completion_pattern.search(completion).group(1)
            return correction.strip()
        except:
            retries += 1
            print(
                f"Something went catastrophically wrong with this one. Investigate! Here's the completion:\n{completion}"
            )

if __name__ == "__main__":  # test
    logic_llm = Llama(
        model_path=LOGICAL_MODEL,
        n_gqa=8,
        offload_kqv=True,
        n_ctx=4096,
        n_gpu_layers=1000,
    )  # load the logical LLM and offload everything
    q_test = [
        (
            ") Explain how our understanding of planetary motion has changed over time.",
            "The understanding has evolved from the Earth being stationary and at the centre of the universe, to it orbiting the sun in an elliptical path with other planets while still rotating on its axis.",
            "The story of our world is a story that is still very imperfectly known. A couple of hundred years ago men possessed the history of little more than the last three thousand years. What happened before that time was a matter of legend and speculation.  Over a large part of the civilized world it was believed and taught that the world had been created suddenly in 4004 B.C., though authorities differed as to whether this had occurred in the spring or autumn of that year. This fantastically precise misconception was based upon a too literal interpretation of the Hebrew Bible, and upon rather arbitrary theological assumptions connected therewith.  Such ideas have long since been abandoned by religious teachers, and it is universally recognized that the universe in which we live has to all appearances existed for an enormous period of time and possibly for endless time.  Of course there may be deception in these appearances, as a room may be made to seem endless by putting mirrors facing each other at either end. But that the universe in which we live has existed only for six or seven thousand years may be regarded as an altogether exploded idea.\n\nThe earth, as everybody knows nowadays, is a spheroid, a sphere slightly compressed, orange fashion, with a diameter of nearly 8,000 miles.  Its spherical shape has been known at least to a limited number of intelligent people for nearly 2,500 years, but before that time it was supposed to be flat, and various ideas which now seem fantastic were entertained about its relations to the sky and the stars and planets.  We know now that it rotates upon its axis (which is about 24 miles shorter than its equatorial diameter) every twenty-four hours, and that this is the cause of the alternations of day and night, that it circles about the sun in a slightly distorted and slowly variable oval path in a year. Its distance from the sun varies between ninety-one and a half millions at its nearest and ninety-four and a half million miles.",
        ),
        (
            ") Identify and explain changes in human understanding throughout history regarding the age of the Earth.",
            "Initially, religious texts suggested a young earth dating back no more than several thousand years. However, evidence from geology and astronomy has shown us that the earth is over four billion years old.",
            "The story of our world is a story that is still very imperfectly known. A couple of hundred years ago men possessed the history of little more than the last three thousand years. What happened before that time was a matter of legend and speculation.  Over a large part of the civilized world it was believed and taught that the world had been created suddenly in 4004 B.C., though authorities differed as to whether this had occurred in the spring or autumn of that year. This fantastically precise misconception was based upon a too literal interpretation of the Hebrew Bible, and upon rather arbitrary theological assumptions connected therewith.  Such ideas have long since been abandoned by religious teachers, and it is universally recognized that the universe in which we live has to all appearances existed for an enormous period of time and possibly for endless time.  Of course there may be deception in these appearances, as a room may be made to seem endless by putting mirrors facing each other at either end. But that the universe in which we live has existed only for six or seven thousand years may be regarded as an altogether exploded idea.\n\nThe earth, as everybody knows nowadays, is a spheroid, a sphere slightly compressed, orange fashion, with a diameter of nearly 8,000 miles.  Its spherical shape has been known at least to a limited number of intelligent people for nearly 2,500 years, but before that time it was supposed to be flat, and various ideas which now seem fantastic were entertained about its relations to the sky and the stars and planets.  We know now that it rotates upon its axis (which is about 24 miles shorter than its equatorial diameter) every twenty-four hours, and that this is the cause of the alternations of day and night, that it circles about the sun in a slightly distorted and slowly variable oval path in a year. Its distance from the sun varies between ninety-one and a half millions at its nearest and ninety-four and a half million miles.",


from llama_cpp import LlamaGrammar

### A grammar that forces the model to generate correct character cards (with traits, names, everything)

multi_turn_conversation_grammar = LlamaGrammar.from_string(
    r"""

# The root rule defines the structure of the dialogue
root ::= [^\t]+ #"\n" statement+# statement anything+ # Idea: get it started off right, then  let it end how it wants

# Statement by Character Name 1
statement ::= [^\n]+ ":" [^\n]+ "\n"

anything ::= [^\t]+ # I don't think GBNF has a wildcard for ANY character, so I just ban tabs because this text shouldn't be indented, and say "go for it"

"""
)

# question_grammar = LlamaGrammar.from_string(r"""# GBNF Grammar for Q&A Format with Flexible Punctuation

# root ::= answer
# answer ::= "Test"
# """)