from llama_cpp import LlamaGrammar

question_grammar = LlamaGrammar.from_string(
    r"""
root ::= (question-one answer "\n")

# Define the question structure with a number followed by content and ending punctuation
# question ::= number ".) " [^\n]+ [?.!] "\n" # maybe blacklist ?!. along with newlines

# Define the answer structure
answer ::= "Answer: " [^\n]+ "\n"

# Define a number (in this case, limiting to any three-digit number for simplicity)
number ::= [1-9] [0-9]? [0-9]?

# Define content as a sequence of characters excluding the word "paragraph" and using not_paragraph to build up the content
# content ::= (not-paragraph "paragraph")* #not_paragraph

question-one ::= "1.) " [^\n]+ [?.!] "\n" # maybe blacklist ?!. along with newlines
# ws ::= [ \t\n]*


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)}\"\"\"


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]}\"\"\"


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


import re
from .question_relevant_grammar import question_relevant_grammar
from llama_cpp import Llama
from .constants import LOGICAL_MODEL

# POSSIBLE TODO:
# Add an "Off the rails insane forever generation" check, like on this output:
# I haven't seen this happen since I started using 70bs though.
"""
4.) What did ancient people believe about the shape and movement of our world, before scientific discoveries indicated otherwise? (Hint: The earth is a spheroid, not flat. The universe existed before the year 4004 BCE. The earth rotates on its axis and orbits around the sun.) You are allowed to dramatically change/revamp/rewrite this question's content. The goal here is to create a different question that still requires information from the paragraphs, but doesn't require knowledge of flawed assumptions connected with biblical interpretations or an understanding of events before the creation of the universe, which goes beyond the generally accepted scientific evidence provided in the text.  As you see fit. (Don't just rephrase the old question.) The goal here is to create a different question that still requires information from the paragraphs, but doesn't require knowledge of flawed assumptions connected with biblical interpretations or an understanding of events before the creation of the universe, which goes beyond the generally accepted scientific evidence provided in the text.  As you see fit. (Don't just rephrase the old question.) So you should create a new question that only requires information from the paragraphs to solve. It is not necessary to refer to the text at all. Also note: An example about how NOT to ask questions is if the text states fact X, but does not explain how X was established, do not ask a question "How do we know X". But instead you might consider asking how X relates to other facts in the paragraph, or how these facts together lead to a progression of ideas, "Explain how X, Y, and Z are related" for instance. Here's what we know: 1. The earth is a spheroid, not flat (paragraph 3). 2. The universe existed before the year 4004 BCE (paragraph 5). 3. The earth rotates on its axis and orbits around the sun (paragraph 5).
Answer: We don't have an answer.
"""

# Answer vetting
def check_question(qatuple, logic_llm):
    retries = 0
    while retries <= 4:
        decision_prompt = f"""You are an expert educational AI. Given a paragraph or two from a larger text, and a question based on the paragraphs, you will make a determination as to whether the question tests ONLY information in the paragraphs. Essentially: you will check if the question is answerable, given the information in the paragraphs. Your task includes first analyzing the text, thinking through whether or not the question reflects aspects of the paragraphs provided.


import re
from .check_qatuple_context_grammar import check_qatuple_context_grammar
from llama_cpp import Llama
from .constants import LOGICAL_MODEL

# For the reword step (ONLY USE IF JUDGEMENT IS REWORD, OTHERWISE WE JUST IGNORE THE LAST BIT OF THE GEN)
def extract_question_answer(response):
    # Define the regex pattern to match the question and answer
    pattern = r"### Question Rewording \(using text details as reference\):\nQuestion: (.+?)\nAnswer: (.+?)\n"

    # Search for the pattern in the response
    match = re.search(pattern, response)

    # Extract and return the question and answer if a match is found
    if match:
        question = match.group(1).strip()
        answer = match.group(2).strip()
        return question, answer
    else:
        return None, None


from llama_cpp import LlamaGrammar

make_regenerate_answer_plan_grammar = LlamaGrammar.from_string(
    r"""
root ::= analyze-step understand-step identify-step plan-step

analyze-step ::= "Step " [0-9]?[0-9] ". " "Analyze the Text:" [^\n]+ "\n"

understand-step ::= "Step " [0-9]?[0-9] ". " "Understand the Question:" [^\n]+ "\n"

identify-step ::= "Step " [0-9]?[0-9] ". " "Identify the Incorrect Part of the Answer:" [^\n]+ "\n"

plan-step ::= "Step " [0-9]?[0-9] ". " "Plan a Corrected Answer:" [^\n]+ "\n"
"""
)

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

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


from llama_cpp import LlamaGrammar

regenerate_answer_constrain_to_text_grammar = LlamaGrammar.from_string(
    r"""

root ::= reasoning

reasoning ::= [^\n]+ "\"\"\""
"""
)


import re
from .questions_grammar import questions_grammar
from llama_cpp import Llama
from .constants import LOGICAL_MODEL

# Answer generation code, begin.
# structure: define a series of helpers, then define the control flow, exception handling, retries etc. in a for loop that iterates over the processed sequences of paragraphs at the end in another cell

# Since some paragraphs can be much shorter
# First off, question generation.

# Each local LLM function essentially has 3 phases: prompt, regex to extract response, and reaction to that response.
# However I'm not going to build an abstraction for that because I need fine control.

# If any function fails to make things, it won't throw, it'll just return None.

# Strengths of open source AI: hella cheap, very customizable, you can call it as much as you want
# Downside: you need very good regexes to catch its outputs

def generate_short_questions(text, existing_question_tuples, logic_llm):


from llama_cpp import LlamaGrammar

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

scenario_grammar = LlamaGrammar.from_string(
    r"""

root ::= reasoning-start

reasoning-start ::= [^\n\t]+ "."

"""
)


from llama_cpp import LlamaGrammar

judge_paragraph_grammar = LlamaGrammar.from_string(
    r"""

root ::= identify-content-step evaluate-relevance-step assess-contexts-and-formats-step assess-possibility-step determine-suitability-step check-contextual-completeness-step final-step "\n"

identify-content-step ::= "Step " [0-9]?[0-9] ". " "Identify Paragraph Content: " [^\n]+ "\n"

evaluate-relevance-step ::= "Step " [0-9]?[0-9] ". " "Evaluate Educational Relevance: " [^\n]+ "\n"

assess-contexts-and-formats-step ::= "Step " [0-9]?[0-9] ". " "Assess Specific Contexts and Formats:" "\n" context-format-bullets

assess-possibility-step ::= "Step " [0-9]?[0-9] ". " "Assess the Possibility of Formulating Questions: " [^\n]+ "\n"

determine-suitability-step ::= "Step " [0-9]?[0-9] ". " "Determine Suitability for Educational Purposes: " [^\n]+ "\n"

check-contextual-completeness-step ::= "Step " [0-9]?[0-9] ". " "Check for Contextual Completeness: " [^\n]+ "\n"

final-step ::= "Step " [0-9]?[0-9] ". " "Final Judgment: " ("Unsuitable" | "Suitable" | "suitable" | "unsuitable") "\n"


from llama_cpp import LlamaGrammar

ensure_answer_consistent_grammar = LlamaGrammar.from_string(
    r"""

root ::= understand-question-step compare-question-step understand-answer-step compare-step final-step "\n"

# step ::= "Step " [0-9]?[0-9] ". " ( "Analyze" | "Understand" | "Compare" | "Skip" | "Final" ) [^\n]+ "\n"

understand-question-step ::= "Step " [0-9]?[0-9] ". " "Understand the provided question:" [^\n]+ "\n"

compare-question-step ::= "Step " [0-9]?[0-9] ". " "Compare the conversation's question: " [^\n]+ "\n"

understand-answer-step ::= "Step " [0-9]?[0-9] ". " "Understand the provided answer:" [^\n]+ "\n"

# compare-step ::= "Step " [0-9]?[0-9] ". " "Compare the " ("first" | "second" | "third" | "fourth" | "fifth" | "sixth" | "seventh" | "eighth" | "ninth") " Part of the Answer with the Text: check if the text " [^\n]+ "\n"

compare-step ::= "Step " [0-9]?[0-9] ". " "Compare the conversation's answer:" [^\n]+ "\n"

final-step ::= "Step " [0-9]?[0-9] ". " "Final Judgement: " ("Inconsistent" | "Consistent") "\n"
"""


from itertools import product
import random

def combine_traits(personality_matrix):  # GPT-generated
    # Using itertools.product to generate all possible combinations
    combinations = product(*personality_matrix)

    # Joining each combination into a single string
    combined_traits = [
        "\n".join(combination).strip().replace("\n\n", "\n")
        for combination in combinations
    ]

    return combined_traits

def special_instructions(n=1, non_axis_traits=False, non_axis_traits_only=False):
    """
    documentation todo
    """


import re

def has_sequential_chars(string1, string2, n):
    """
    Check if any n sequential characters from string1 appear in string2.

    Args:
    string1 (str): The first string to check.
    string2 (str): The second string in which to look for sequences.
    n (int): The length of the sequence to check.

    Returns:
    bool: True if any n sequential characters from string1 are found in string2, False otherwise.
    """

    # Check if n is larger than the length of string1.
    if n > len(string1):
        return False, ""

    # Iterate over string1 and check for each n-length substring in string2


from llama_cpp import LlamaGrammar

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

question_plan_grammar = LlamaGrammar.from_string(
    r"""

# root ::= reasoning-start
# At least 3 steps
root ::= identify-step generate-step brainstorm-step relationships-step if-then-step make-suitable-step

# no-questions-after-here ::= "\nI will not ask any questions about the following information: " [^\n]+ "."

identify-step ::= "Step " [0-9]?[0-9] ". " "Identify Key Topics:" [^\n]+ "\n"

# generate-step ::= "Step " [0-9]?[0-9] ". " "Determine Information-Rich Areas:" [^\n]+ "\n"

brainstorm-step ::= "Step " [0-9]?[0-9] ". " "Brainstorm and Develop Questions Testing Recall:" [^\n]+ "\n"

relationships-step ::= "Step " [0-9]?[0-9] ". " "Devise Questions" [^\n]+ "\n"


from llama_cpp import LlamaGrammar

proofread_output_grammar = LlamaGrammar.from_string(
    r"""

root ::= analyze-step step+ "\n\nBegin Edit: " [^\n]+

step ::= "Step " [0-9]?[0-9] ". " ( "Analyze" | "Understand" | "Compare" | "Skip" | "Notice" | "Note" | "There is" | "Error" | "I found" | "End" | "There are" ) [^\n]+ "\n"

analyze-step ::= "Step " [0-9]?[0-9] ". " "Analyze" [^\n]+ "\n"
"""
)
# question_grammar = LlamaGrammar.from_string(r"""# GBNF Grammar for Q&A Format with Flexible Punctuation

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