
import re
from PIL import Image
import base64
import os

from openai import OpenAI
client = OpenAI(
  api_key=os.getenv("OPENAI_API_KEY")
)

def encode_image(image_path):
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode("utf-8")
    
def construct_messages(prompt_parts):
    messages = []
    current_text = ""
    
    for part in prompt_parts:
        if isinstance(part, str):
            if current_text:
                current_text += "\n" + part
            else:
                current_text = part
        elif isinstance(part, dict) and "image_path" in part:
            if current_text:
                messages.append({"type": "text", "text": current_text})
                current_text = ""
            base64_image = encode_image(part["image_path"])
            messages.append({"type": "image_url", "image_url": {"url": f"data:image/png;base64,{base64_image}"}})
    
    if current_text:
        messages.append({"type": "text", "text": current_text})
    
    return messages

def get_model_response(prompt_parts, model, temp=0.0, max_tokens=2000, preconstructed_message = None):
    if model.lower() == "gpt4o":
        model = "gpt-4o"
    elif model.lower() == "gpt4o-mini":
        model = "gpt-4o-mini"
    else:
        raise ValueError("Unsupported model type. Choose 'GPT4o' or.")
    
    if preconstructed_message: messages = preconstructed_message
    else: messages = [{"role": "user", "content": construct_messages(prompt_parts)}]
    
    response = client.chat.completions.create(
        model=model,
        messages=messages,
        temperature=temp,
        max_tokens=max_tokens,
        top_p=1,
        frequency_penalty=0,
        presence_penalty=0
    )
    response_text = response.choices[0].message.content
    return response_text

def extract_content(response, format, remove_new_lines=True):
    def remove_newlines(text):
        return re.sub(r'\n', '', text)
    response_mod = response
    if remove_new_lines:
        response_mod = remove_newlines(response)
    if format == "text":
        match = re.search(r"```answer_text(.*?)```", response_mod, re.DOTALL)
        if not match:
            match = re.search(r"```answer_text(.*?)", response_mod, re.DOTALL)
        if not match:
            match = re.search(r"```(.*?)```", response_mod, re.DOTALL)
        if not match:
            match = re.search(r"```(.*?)", response_mod, re.DOTALL)
        if match:
            return match.group(1).strip()
        else:
            return ''
    elif format == "code":
        match = re.search(r"```code(.*?)```", response_mod, re.DOTALL)
        if not match:
            match = re.search(r"```python(.*?)```", response_mod, re.DOTALL)
        if match:
            return match.group(1).strip()
        else:
            return ''
    elif format == "action_reasoning":
        match = re.search(r'```action_reasoning(.*?)```', response_mod, re.DOTALL)
        if match:
            return match.group(1).strip()
        else:
            return ''
    elif format == "reasoning":
        match = re.search(r'```reasoning(.*?)```', response_mod, re.DOTALL)
        if match:
            return match.group(1).strip()
        else:
            return ''
    elif format == "current_state":
        match = re.search(r'```current_state(.*?)```', response_mod, re.DOTALL)
        if match:
            return match.group(1).strip()
        else:
            return ''
    elif format == "action_taken":
        match = re.search(r'```action_taken(.*?)```', response_mod, re.DOTALL)
        if match:
            return match.group(1).strip()
        else:
            return ''
    elif format == "yes_no":
        match = re.search(r'```yes_no(.*?)```', response_mod, re.DOTALL)
        if match:
            content = match.group(1).strip().lower()
            if "yes" in content and "no" not in content:
                return True
            elif "no" in content and "yes" not in content:
                return False
            else:
                return None
        else:
            return False
    elif format == "number":
        match = re.search(r"```number(.*?)```", response_mod, re.DOTALL)
        if match:
            try:
                return float(match.group(1).strip())
            except ValueError:
                return 0
        else:
            return 0
    elif format == "ranking":
        match = re.search(r"```ranking(.*?)```", response_mod, re.DOTALL)
        if match:
            ranking_content = match.group(1).strip()
            return [id.strip() for id in ranking_content.split(',') if id.strip()]
        else:
            return []
    elif format == "error":
        match = re.search(r'```error(.*?)```', response_mod, re.DOTALL)
        if match:
            return match.group(1).strip()
        else:
            return ''
    else:
        return None

def encode_initial_state(state_file_path, problem_file_path, model, description):
    with open(state_file_path, 'r') as file:
        initial_state = file.read()
    with open(problem_file_path, 'r') as file:
        initial_format = file.read()
    prompt = f"Given this problem and its natural language description:\n{initial_format}\n{description}\nWrite the natural language description of the initial state for this problem:\n\n{initial_state}"
    encoded_state = extract_content(get_model_response([prompt], model), "text")
    return encoded_state


def get_nl_description(file_path, model):
    with open(file_path, 'r') as file:
        initial_format = file.read()
    prompt = f"Write a natural language description of the task for this problem:\n\n{initial_format}"
    response = get_model_response([prompt], model)
    description = extract_content(response, "text")
    return description

def get_feasible_actions(problem_file_path, model, description):
    with open(problem_file_path, 'r') as file:
        initial_format = file.read()
    prompt = f"Given this problem and its natural language description:\n{initial_format}\n{description}\nList the feasible actions, where each action is represented with a unique short phrase/word, for this problem in Python syntax with no Python code in the response."
    response = get_model_response([prompt], model)
    actions = extract_content(response, "text")
    return actions.split('\n')

def next_action(problem_description, possible_actions, initial_state, current_state, goal_state, model, temp, chosen_actions, previous_attempt=None, error_message=None):
    # Base prompt describing the problem, the current state, and the goal
    prompt_parts = [
        "Consider this problem description:",
        problem_description,
        "This is the initial state of the problem:",
        initial_state.state_nl,
        "Our goal state is:",
        goal_state.state_nl,
        "And this is a visualization of the goal state:",
        {"image_path": goal_state.diagram_picture},
        "This is a list of possible actions for this problem. The arguments to these actions must be strings, in the order given:",
        possible_actions,
        f"After taking the actions listed below we have arrived at our current state which is at depth {current_state.depth}:",
        str(current_state.action_path),
        "This is a description of our current state:",
        current_state.state_nl,
        "This is a visualization of the current state:",
        {"image_path": current_state.diagram_picture},
        "This is the diagram encoding of the current state (Diagram encoding encodes a geometrical layout of the scene. It is a collection of statements describing the relative/absolute position, relative/absolute size, status (includes color if relevant), and a text identifier for each object in the problem):",
        current_state.diagram_encoding,
        #"This is the reasoning for why we took the last few actions in the path:",
        #str(last_few_action_reasoning),
        "Your task is to choose the next best action that is allowed from the current state, does not violate any constraints, and brings us the closest to the goal state.",
        "Think step by step about what should be the next actions. Then provide a very short explanation for why the action you chose is the best action and what is your long term plan (in one sentence). The reasoning should be consise and must be provided in the following format:",
        "```action_reasoning",
        "<reasoning>",
        "```",
        "Next, write the chosen action in the format below. The action name must match one of the names given in the list above, and all arguments should be included in the correct order:",
        "```action_taken",
        "<action chosen, no other information>",
        "```",
        "Finally, describe the new state resulting from taking this action, in the format below. This should contain a consise description of the updated status. For each object in the state include a short statement and describe the status of each object. Be very clear. ",
        "```current_state",
        "<status of all objects in the current state>",
        "```",
    ]

    # If this is NOT the very first child from the current state (i.e., there are already child_states),
    # then we inform the model about the previously chosen actions for this state.
    if len(current_state.child_states) > 0:
        prompt_parts.append(
            "Suggest a valid action from the current state that differes meaningfully from the actions below:"
        )
        prompt_parts.append({str(chosen_actions)})

    # If we have a previous attempt and an error message, mention it
    if previous_attempt and error_message:
        prompt_parts.extend([
            "This is the previous attempt:",
            previous_attempt,
            "It failed with the following error message, avoid this error:",
            error_message
        ])
        
    # Call out to the model
    response = get_model_response(prompt_parts, model, temp=temp)

    # Extract the relevant blocks from the response
    action_reasoning = extract_content(response, "action_reasoning")
    current_state_nl = extract_content(response, "current_state")
    action_taken = extract_content(response, "action_taken")

    return action_reasoning, current_state_nl, action_taken


def check_action_validity(problem_description, current_state, action_taken, new_state_nl, new_state_reasoning, new_state, goal_state, possible_actions, initial_state, model):
   
    """hard_neg_dir = f"{domain_name}<PATH_REMOVED>"
    
    if os.path.exists(hard_neg_dir) and os.path.getsize(hard_neg_dir) > 0:
        with open(hard_neg_dir, "r") as f:
            hard_negatives = f.read() 
    else:
     (fOur goal is to verify wether the selected action is feasible given the state (i.e. to verify that all of the preconditions of the selected action are satisfied and choosing this action does not violate any rules)
        Here are some hypothetical state & action pairs, where the action selected is invalid from the given state:
        {hard_negatives} if hard_negatives else ""),
        """
        
    prompt_parts = [
        "Consider the following problem description:",
        problem_description,
        "This is the initial_state of the problem:",
        initial_state.state_nl,
        "This is a visualization of the initial state of the problem:",
        {"image_path": initial_state.diagram_picture},
        "This is the goal_state of the problem:",
        goal_state.state_nl,
        "This is a visualization of the goal state of the problem:",
        {"image_path": goal_state.diagram_picture},
        "We are currently in the following state (current_state):",
        current_state.state_nl,
        "This is a visualization of the current_state:",
        {"image_path": current_state.diagram_picture},
        "This is the diagram encoding of the current state (Diagram encoding encodes a geometrical layout of the scene. It is a collection of statements describing the relative/absolute position, relative/absolute size, status (includes color if relevant), and a text identifier for each object in the problem):",
        current_state.diagram_encoding,
        "Below is the list of the actions taken from the initial state to reach this state:",
        str(current_state.action_path),
        "From this current_state we have the following suggested action to arrive at a new state:",
        action_taken,
        "It is suggested that after taking this action we would arrive at the new_state described below:",
        new_state_nl,
        "This is a suggested visualization of the new state is given below:",
        {"image_path": new_state.diagram_picture},
        "This is the diagram encoding of the new state:",
        new_state.diagram_encoding,
        "The following is the reasoning behind this choice of action:",
        new_state_reasoning,
        "And below I have attached all actions possible in this problem domain:",
        str(possible_actions),
        """Given the constraints outlined in the problem description, answer and think step by step about the questions below:
        Is the action valid and can it be applied to the objects involved (e.g., are the objects involved clear (if applicable)? Is the correct action selected for that object based on the current_state and preconditions of the action? etc)? (To answer this you have to first write down every precondition of the action chosen specific to the affected variables. Then for each precondition explain if it is met or violated in the currect_state's visualization, description, and diagram encoding. Think carefully and step by step.
        Does the action bring us closer to the goal state and align with the goal state's constraints? (To answer this you must write the affects of the action taken and explain wether or not any of the affects violate any of the goal state's constraints)
        Does the new state visualization accurately reflect the changes caused by the action? (to answer yes, the number of the objects in the new_state's visualization must be euqal to the number of objects in the visualization of the current_state and initial state) (Also the status of the objects should be accurately visualizes, with no objects floating if that doesn't make sense, or no extra grid cells for example)
        Does the text description and diagram encoding of the new state accurately update the status of the affected objects? 
        Does the action's syntax and arguments match the expected format and types (refer to the possible actions list)?
        Is the reasoning behind the action selection logical and aligned with achieving the goal state? 
        Answer all of these questions, and explain why. Think step by step. At the end, provide a final answer in the following format: answer "yes" if all responses to the questions above are "yes"; otherwise, answer "no".""",
        "```yes_no",
        "<your final yes or no answer>",
        "```",
        "After writing the code block above, if your final answer was 'no', provide a very short phrase describing the error. Be as specific as possible (i.e. summerize your reasoning above in 1 sentence). Return your answer in the following format:",
        "```error", 
        "<your error description>",
        "```"
    ]
    response = get_model_response(prompt_parts, model)
    validity = extract_content(response, "yes_no")
    error_description = extract_content(response, "error")
    prompt_parts.append("""Make sure to follow these steps when returning your answer:
                        write ```yes_no
                        then write your anwer yes or no
                        then write ```
                        If you asnwered no above you must return the reason for a no answer int he follwoing steps: 
                        first write ```error
                        then write the error, ie why you answered no
                        lastly write ```""")
    if validity == None:
        response = get_model_response(prompt_parts, model)
        validity = extract_content(response, "yes_no")
        error_description = extract_content(response, "error")
    
    validity = validity if validity else False
    return validity, error_description

def get_cost(problem_description, action_path, state_nl, model):
    prompt = [
        f"""
        Given the following problem description:
        {problem_description}
        We have taken the actions below:
        {action_path}
        Now we are at the following state:
        {state_nl}
        Provide the cost of reaching this state as a numerical value, in the format provided below. The cost represents the metric we are trying to minimize in order to reach the goal state, which is typically quantified by the number of actions taken unless an optimization goal is specified otherwise in the problem statement.
        ```number
        <cost, as a single number>
        ```
        """
    ]
    response = get_model_response(prompt, model)
    cost = extract_content(response, "number")
    return float(cost)

def generate_diagram_encoding(problem_description, initial_state, child_state, model, temp, previous_attempt=None, error_message=None):
    prompt_parts = [
        "Consider the following problem description:",
        problem_description,
        "This is the initial state of the problem:",
        initial_state.state_nl,
        "Below is the diagram encoding of the initial state. Diagram encoding encodes a geometrical layout of the scene. It is a collection of statements describing the relative/absolute position, relative/absolute size, status (includes color if relevant), and a text identifier for each object in the problem. ",
        initial_state.diagram_encoding,
        "We have taken the actions below from the initial state:",
        str(child_state.parent.action_path),
        "And we have arrived at the following state of the problem (the current_state):",
        child_state.parent.state_nl,
        "This is a visualization of the current_state:",
        {"image_path": child_state.parent.diagram_picture},
        "The diagram encoding of the current state is:",
        child_state.parent.diagram_encoding,
        "After taking this new action:",
        child_state.action_taken,
        "We have reached the following new state (the new_state):",
        child_state.state_nl,
        "Generate the diagram encoding for the new state, where the status of the objects affected by the action are correctly updated and the status of all other objects remain the same as in the current_state. In the status you must include if the object is clear or not.",
        "For this you have to identify the few objects effected by the action and update theri status and/or position in the encoding",
        "Note that there should be at least one statement (a statement inlcudes a text identifier, position, size, etc) for each object/location in the scene for the encoding to be valid.",
        "IMPORTANT: You have to make sure the number of statements is equal to the number of statement in the diagram encoding of the current_state and the diargam encoding of the initial state. Return your answer in the following format",
        "```answer_text",
        "<final answer no additional text>",
        "```"
    ]
    if previous_attempt and error_message:
        prompt_parts.extend([
            "This is the previous attempt:",
            previous_attempt,
            "This attempt failed to pass a unit test with the following error. Avoid this err:",
            error_message
        ])
    response = get_model_response(prompt_parts, model, temp)
    diagram_encoding = extract_content(response, "text")
    return diagram_encoding

def test_diagram_encoding(problem_description, initial_state, child_state, model):
    prompt_parts = [
        "Consider the following problem description:",
        problem_description,
        "This is the initial state of the problem:",
        initial_state.state_nl,
        "Below is the diagram encoding of the initial state. Diagram encoding encodes a geometrical layout of the scene. It is a collection of statements describing the relative/absolute position, relative/absolute size, status (includes color if relevant), and a text identifier for each object in the problem. ",
        initial_state.diagram_encoding,
        "We have taken the actions below form the initial_state:",
        str(child_state.parent.action_path),
        "We have arrived to the following state (the current_state):",
        child_state.parent.state_nl,
        "This is a visualization of the current_state:",
        {"image_path": child_state.parent.diagram_picture},
        "The diagram encoding of the current state is:",
        child_state.parent.diagram_encoding,
        "After taking this new action:",
        child_state.action_taken,
        "We have reached the following new child_state:",
        child_state.state_nl,
        "A potential diagram encoding for the child_state is:",
        child_state.diagram_encoding,
        """Your task is to check whether this diagram encoding is correct:
        
        For this you have to check the following: Explain your reasnoing and iterate through each statement in the encoding when checking any of the constraints below.
        - Are the number of statement (i.e. the number of objects encoded) in the diagram encoding of the child state equal to the number of statements in the diagram encoding of the current_state and the inital_state?
        - Consider the objects affected by the action taken, is the status of these objects correctly updated in the suggested diagram encoding?
        - Consider the objects not effected by the action, are they included in the suggested diagram encoding for the child_state with the same status as before (as in the current state's diagram encoding)?
        - Is the number of stataments in the current_state diagram encoding and the child_state diagram encoding equal? (i.e. Are all objects included in the new diagram?)
        
        Think step by step. If you said no to nay of the questions above, or if any of the statements is wrong with respect to the current state, action taken, or problem constraints, return no; say yes otherwise. Return your final answer in the following format:""",
        "```yes_no",
        "<your final answer, yes or no>",
        "```",
        "After writing the code block above, if your final answer was 'no' , provide a very short phrase describing the error. Be as specific as possible (i.e. summerize your reasoning above in 1 sentence). Return your answer in the following format:",
        "```error",
        "<your error description>",
        "```"
    ]
    response = get_model_response(prompt_parts, model)
    validity = extract_content(response, "yes_no")
    error_description = extract_content(response, "error")
    prompt_parts.append("""Make sure to follow these steps when returning your answer:
                        write ```yes_no
                        then write your anwer yes or no
                        then write ```
                        If you asnwered no above you must return the reason for a no answer int he follwoing steps: 
                        first write ```error
                        then write the error, ie why you answered no
                        lastly write ```""")
    if validity == None:
        response = get_model_response(prompt_parts, model)
        validity = extract_content(response, "yes_no")
        error_description = extract_content(response, "error")
    
    validity = validity if validity else False
    return validity, error_description

def rank_states(states, problem_description, goal_state, model):
    # Prepare prompt
    prompt_parts = []
    prompt_parts.append("Consider the following problem description:")
    prompt_parts.append(problem_description)
    prompt_parts.append("And the goal state is:")
    prompt_parts.append(goal_state.state_nl)
    prompt_parts.append("We have generated the following states at the current depth:")
    for state in states:
        prompt_parts.append(f"State ID: {state.id}")
        prompt_parts.append(f"State Description: {state.state_nl}")
        prompt_parts.append(f"This is the action path taken to reach this state: {state.action_path}")
        if state.diagram_picture:
            prompt_parts.append("This is a visualization of the state:")
            prompt_parts.append({"image_path": state.diagram_picture})
    prompt_parts.append("Your task is to first iterate through each state and determine how many of the constraints of the goal state is satisfied in the current state. If no constraints are satisfied for all states given then calculate how close we are relatively to satifying a constraint. Then you must rank the states from best to worst based on the heuristic you calculated, i.e, the states where a higher number of goal constraints are met should be ranked as better.")
    #prompt_parts.append("Important: First rank states with unique object configurations, from best (closest to the goal) to worst. Place states with identical object configurations but different IDs at the end of the sorted list.")
    prompt_parts.append("Provide the final ranking in the following format:")
    prompt_parts.append("```ranking")
    prompt_parts.append("<State IDs in order from best to worst, separated by commas>")
    prompt_parts.append("```")
    prompt_parts.append("Explain your reasoning and think step by step above the code block. For each state, mention exactly which of the goal constraints are satisfied.")
    response = get_model_response(prompt_parts, model)
    reasoning = response  # store the full response for logging
    ranking = extract_content(response, "ranking")
    if not ranking:
        # Return states in the same order as received
        return states, reasoning
    # Convert IDs to integers and map to states
    ranked_ids = [int(id.strip()) for id in ranking if id.strip().isdigit()]
    id_to_state = {state.id: state for state in states}
    ranked_states = [id_to_state[id] for id in ranked_ids if id in id_to_state]
    return ranked_states, reasoning

def generate_diagram_code(domain_name, problem_description, child_state, initial_state, model, temp, save_path, previous_attempt=None, error_message=None):
    example_ini_diagram_reasoning_path = f"{domain_name}<PATH_REMOVED>"
    
    if os.path.exists(example_ini_diagram_reasoning_path) and os.path.getsize(example_ini_diagram_reasoning_path) > 0:
        with open(example_ini_diagram_reasoning_path, "r") as f:
            example_ini_diagram_reasoning = f.read()
    else:
        example_ini_diagram_reasoning = None
        
    prompt_parts = [
        "Here the task is to write matplotlib code to visualize a state. Consider the following problem description:",
        problem_description,
        "This is the initial state of the problem",
        initial_state.state_nl,
        "And this is a visualization of the initial_state:",
        {"image_path": initial_state.diagram_picture},
        "We got the visualization above for the initial_state using the following matplotlib code:",
        initial_state.diagram_code,
        "We also have some reasoning for how the diagram should be drawn, ie what is the meaning behind each shape, color, different sizes and locations. But you should note this reasoning is for a random state but the meanings and the reasoning behind the diagram remains the same:" if example_ini_diagram_reasoning else "",
        example_ini_diagram_reasoning if example_ini_diagram_reasoning else "",
        f"Now after taking the following series of actions {str(child_state.parent.action_path)} we are in the following state (current_state):",
        child_state.parent.state_nl,
        "This is a visualization of the current state:",
        {"image_path": child_state.parent.diagram_picture},
        "And this is the matplotlib code we used to draw the diagram for the current_state:",
        child_state.parent.diagram_code,
        "After taking the following action:",
        child_state.action_taken,
        "We have reached the following new child state (child_state):",
        child_state.state_nl,
        "This is a diagram encoding of the child_state. Diagram encoding encodes how the scene should be geometrically drawn. It is a collection of statements describing the relative/absolute position, relative/absolute size, status (includes color if relevant), and a text identifier for each object in the child_state:",
        child_state.diagram_encoding,
        """Your taks is to generate the matplotlib code that can draw a diagram for the child_state where each object is depicted as described in the diagram encoding given above.
        You have to make sure to do the following:
        - Ensure all objects encoded in the diagram encoding are present in the visualization, with no missing objects or overlapping ones. Also, any object present in the initial state or the current_state diagram should be present in the child_state's diagram (with an updated state). No 2 objects can overlap
        - Make sure that the textual label and status of each object are included inside its corresponding shape.
        - Ensure all text in the diagram is readable, clear, and free of major overlaps or illegible portions.
        - Ensure the diagram is physically plausible, ensuring no elements defy physical expectations (e.g., no floating objects unless appropriate).
        - Visualize the status of objects based on the legend of the given diagrams and the diagram reasoning provided.""",
        "Think step by step about how each object should be coded",
        f"Important: The code must save the figure at the following location: {save_path}",
        "The final answer must be enclosed in",
        "```code",
        "<final python code, no additional text>",
        "```"
    ]
    if previous_attempt and error_message:
        prompt_parts.extend([
            "This is the previous attempt:",
            previous_attempt,
            "This attempt failed to execute with the following error:",
            error_message
        ])
    response = get_model_response(prompt_parts, model, temp=temp)
    diagram_code = extract_content(response, "code", remove_new_lines=False)
    return diagram_code

def test_diagram(problem_description, child_state, domain_name, model):
    example_ini_diagram_reasoning_path = f"{domain_name}<PATH_REMOVED>"
    
    if os.path.exists(example_ini_diagram_reasoning_path) and os.path.getsize(example_ini_diagram_reasoning_path) > 0:
        with open(example_ini_diagram_reasoning_path, "r") as f:
            example_ini_diagram_reasoning = f.read()
    else:
        example_ini_diagram_reasoning = None
        
    prompt_parts = [
        "Consider the following problem description:",
        problem_description,
        "We are currently in the following state (current_state):",
        child_state.parent.state_nl,
        "This is a visualization of the current_state:",
        {"image_path": child_state.parent.diagram_picture},
        f"Below is some reasoning explaining the meaning behind each shape, location, color, and other visual cues in the diagram (the diagram reasoning). {example_ini_diagram_reasoning}" if example_ini_diagram_reasoning else "",
        "The diagram encoding (a geometrical encoding of the scene which includes a collection of statements each describing the relative/absolute position, relative/absolute size, status (includes color if relevant), and a text identifier for each object in the scene) of the current state is:",
        child_state.parent.diagram_encoding,
        "After taking the following action:",
        child_state.action_taken,
        "We have reached the following new child state (child_state):",
        child_state.state_nl,
        "This is the diagram encoding of the child_state is:",
        child_state.diagram_encoding,
        "The task is to determine whether the suggested diagram below accurately visaulizes the child_state:",
        {"image_path": child_state.diagram_picture},
        """I want you to analyze the new diagram and answer these questions with an explanation.
        - Are all of the objects in the diagram encoding of the child_state and objects in the diagram of the currect_state present in the diagram? (to answer yes you have to make sure no object is missing from the visaulization and no 2 objects overlap.)
        - Is the status of objects accurately visualized based on the legend and the diagram reasoning? Is the status updated correctly according to the action taken? 
        - Is the textual label and the status of each object included inside its shape? 
        - Is all of the text in the diagram readable, clear, and with no major overlaps or illegible portions?
        - Does the diagram appear physically plausible (e.g., no floating objects if that wouldn't make sense, etc.)? """,
        "Think step by step about each question",
        "Say yes, if your answer to all of the questions above is yes and the diagram is accurate, correct, and clear, no otherwise, in the following format:",
        "```yes_no",
        "<your final answer, yes or no>",
        "```",
        "After writing the code block above, if your final answer was 'no', provide a very short phrase describing the error. Be as specific as possible (i.e. summerize your reasoning above in 1 sentence). Return your answer in the following format:",
        "```error",
        "<your error description>",
        "```"
    ]
    response = get_model_response(prompt_parts, model)
    validity = extract_content(response, "yes_no")
    error_description = extract_content(response, "error")
    prompt_parts.append("""Make sure to follow these steps when returning your answer:
                        write ```yes_no
                        then write your anwer yes or no
                        then write ```
                        If you asnwered no above you must return the reason for a no answer int he follwoing steps: 
                        first write ```error
                        then write the error, ie why you answered no
                        lastly write ```""")
    if validity == None:
        response = get_model_response(prompt_parts, model)
        validity = extract_content(response, "yes_no")
        error_description = extract_content(response, "error")
    
    validity = validity if validity else False
    return validity, error_description

def test_diagram_ini_g(problem_description, initial_state, goal_state, domain_name, model, goal=False):
    example_ini_diagram_reasoning_path = f"{domain_name}<PATH_REMOVED>"
    
    if os.path.exists(example_ini_diagram_reasoning_path) and os.path.getsize(example_ini_diagram_reasoning_path) > 0:
        with open(example_ini_diagram_reasoning_path, "r") as f:
            example_ini_diagram_reasoning = f.read()
    else:
        example_ini_diagram_reasoning = None
    
    state = goal_state if goal else initial_state
    state_type = "goal" if goal else "initial"
    
    prompt_parts = [
        "Consider the following problem description:",
        problem_description,
        "The following is a description of the initial state of the problem:",
        initial_state.state_nl,
        "And the following is the goal state we are trying to arrive at:",
        goal_state.state_nl,
        #f"The following is diagram encoding of the {state_type}. Diagram encoding consists of statements that detemine the geometric layout of the objects inthe sene, inlcuding the name, status, relative position and relative silze of each object.",
        #state.diagram_encoding,        
        f"The task is to determine whether the suggested diagram below accurately visaulizes the {state_type} state:",
        {"image_path": state.diagram_picture},
        
        f"Below if some reasoning explaining the meaning behind each shape, location, color, and other visual cues in the diagram (the diagram reasoning). {example_ini_diagram_reasoning}" if example_ini_diagram_reasoning else "",

        """I want you to analyze the new diagram and answer these questions with an explanation.
        - Are all of the objects in the state state description present in the diagram? For this you have make sure no object is missing from the visaulization and no 2 objects overlap.
        - Is the status of objects accurately visualized based on the legend and the diagram reasoning? 
        - Is the textual label and the status of each object included inside its shape? Is it clear which status/label belongs to which object?
        - Is the visualization of different objects of the same type consistent? (e.g. if there are 10 blcoks in the diagram, are all blcoks visualized using rectangles of the same size and shape?) Also check for consistancy in the meaning behind shapes and colors with respect to the current_state diagram (e.g. if green encodes clear objects in current_state's diagram, it should also be the color of the clear objects int the child_state diagram)
        - Is all of the text in the diagram readable, clear, and with no major overlaps or illegible portions?
        - Does the diagram appear physically plausible (e.g., no floating objects if that wouldn't make sense, etc.)? """,
        "Think step by step about each question",
        "Say yes, if your answer to all of the questions above is yes and the diagram is accurate, correct, and clear, no otherwise, in the following format:",
        "```yes_no",
        "<your final answer, yes or no>",
        "```",
        "After writing the code block above, if your final answer was 'no' , provide a very short phrase describing the error. Be as specific as possible (i.e. summerize your reasoning above in 1 sentence). Return your answer in the following format:",
        "```error",
        "<your error description>",
        "```",
    ]
    response = get_model_response(prompt_parts, model)
    validity = extract_content(response, "yes_no")
    error_description = extract_content(response, "error")
    prompt_parts.append("""Make sure to follow these steps when returning your answer:
                        write ```yes_no
                        then write your anwer yes or no
                        then write ```
                        If you asnwered no above you must return the reason for a no answer int he follwoing steps: 
                        first write ```error
                        then write the error, ie why you answered no
                        lastly write ```""")
    if validity == None:
        response = get_model_response(prompt_parts, model)
        validity = extract_content(response, "yes_no")
        error_description = extract_content(response, "error")
    
    validity = validity if validity else False
    return validity, error_description

def generate_diagram_encoding_ini_g(problem_description, initial_state, goal_state, model, temp, domain_name, goal=False, previous_attempt=None, error_message=None):
    example_encoding_path = f"{domain_name}<PATH_REMOVED>"
    example_state_path = f"{domain_name}<PATH_REMOVED>"
    
    if os.path.exists(example_encoding_path) and os.path.getsize(example_encoding_path) > 0:
        with open(example_encoding_path, "r") as f:
            example_encoding = f.read() 
            
        with open(example_state_path, "r") as f:
            hypothetical_state = f.read()
    else:
        example_encoding = None
        hypothetical_state = None
    
    initial_state_description = f"""
    This is the description of the initial state of this problem:
    {initial_state.state_nl}
    The diagram encoding we previously acquired for the initial state is:
    {initial_state.diagram_encoding}
    Now I want to write the diagram encoding of the goal state of this instance:
    {goal_state.state_nl} 
    Note: If there is no information provided for an object in the goal state description but the object is present in the initial state, we should locate the object at the bottom of the page with status no constraints.
    """ if goal else f"""{initial_state.state_nl}
    You should encode the status of all objects described in the initial state in the diagram encoding."""
    
   
    example = f"""
    Consider this example of writing a collection of descriptions of the relative/absolute position, relative/absolute size, status (includes color if relevant), and a text identifier for each object in the scene, called a diagram encoding.
    problem description:
    A set of robots use different colors to paint patterns in floor tiles. The robots can move around the floor tiles in four directions (up, down, left, and right). Robots paint with one color at a time but can change their spray guns to any available color. In this example the available colors are white and black. Robots can only paint the tile that is in front (up) and behind (down) them, and once a tile has been painted, no robot can stand on it.
    An incomplete diagram encoding of the following scene is:
    We are currently at state 0 with robot1 in tile 3_1 with a white paint gun, robot2 in tile 2_2 with a white paint gun, and none of the tiles are colored
    Diagram encoding (partially provided, must encode all of the tiles):
    (text/identifier: tile_0-1, shape: rectangle, size: 4x3, position:lower left corner, status: not painted)
    (text/identifier: tile_0-2, shape: rectangle, size: 4x3/same as tile_0-1, position:to right of tile_0-1 & to left of tile_0-3 (same x), status: not painted)
    … for all tiles
    (text/identifier: robot1, shape: circle, size: smaller than tile_3-1, position:inside tile_3-1, status: white paint gun)
    (text/identifier: robot2, shape: circle, size: smaller than tile_2-2, position:inside tile_2-2, status: white paint gun)
    (text/identifier: color_choice_white, shape: none, size: text size smaller than tile_0-1, position:lower left corner & below tile_0-1, status: white color)
    (text/identifier: color_choice_black, shape: none, size: text size smaller than tile_0-1, position:lower left corner & to the right of color_choice_white & same x as color_choice_white, status: black color)
    Now consider the following problem description:
    {problem_description}
    """ if not example_encoding else f"""
    In this problem domain:
    {problem_description}
    Consider this example of writing a collection of descriptions of the relative/absolute position, relative/absolute size, status (includes color if relevant), and a text identifier for each object in the scene, called a diagram encoding.
    The diagram encoding below is written for a random hypothetical instantiation of the problem domain.
    The hypothetical state:
    {hypothetical_state}
    The diagram encoding for the hypothetical state:
    {example_encoding}
    """
    
    prompt_parts = [example] + [f"""
    I want to write a similar diagram encoding specified to the following instance of the domain above:
    {initial_state_description}
    Generate the diagram encoding for this state, describing the relative/absolute position, relative/absolute size, status (includes color if relevant), and a text identifier for each object in the scene.
    There should be at least one statement (a statement inlcudes a text identifier, position, size, etc) for each object/location in the scene.
    The final answer must be enclosed in
    ```answer_text
    <final answer no additional text>
    ```
    """]
    if previous_attempt and error_message:
        prompt_parts.extend([
            "Additional info: This is the previous attempt:",
            previous_attempt,
            "This attempt failed to pass a unit test with the following error. Avoid this err:",
            error_message
        ])
    response = get_model_response(prompt_parts, model, temp)
    diagram_encoding = extract_content(response, "text")
    return diagram_encoding

def test_diagram_encoding_ini_g(problem_description, initial_state, goal_state, model, goal=False):
    state = goal_state if goal else initial_state
    state_type = "goal" if goal else "initial"
    additional_instruction = """
    Note: If there is no information provided for an object in the goal state description but the object is present in the initial state, one must include the object at the bottom of the page with status no constraints."""
    prompt_parts = [f"""
    Consider the following problem description:
    {problem_description}
    The following is a description of the {state_type} state:
    {state.state_nl}
    The diagram encoding (which is made of a description of the relative/absolute position, relative size, status (includes color if relevant), and a text identifier for each object in the scene) of the {state_type} state is:
    {state.diagram_encoding}
    Your task is to check whether this diagram encoding is correct. For every item in the diagram encoding, write a description about wether the object description matches the state description and if the status of the object is correct according to the natural language description of the state. Think step by step and iterate thought evey object. The final answer must be a yes or a no. If any of the descriptions is wrong, say no, yes otherwise, in the following format.
    ```yes_no
    <final answer, yes or no>
    ```
    {additional_instruction}
    After writing the code block above, if your final answer was 'no' , provide a very short phrase describing the error:
    ```error
    <error description>
    ```
    """]
    response = get_model_response(prompt_parts, model)
    validity = extract_content(response, "yes_no")
    error_description = extract_content(response, "error")
    prompt_parts.append("""Make sure to follow these steps when returning your answer:
                        write ```yes_no
                        then write your anwer yes or no
                        then write ```
                        If you asnwered no above you must return the reason for a no answer int he follwoing steps: 
                        first write ```error
                        then write the error, ie why you answered no
                        lastly write ```""")
    if validity == None:
        response = get_model_response(prompt_parts, model)
        validity = extract_content(response, "yes_no")
        error_description = extract_content(response, "error")
    
    validity = validity if validity else False
    return validity, error_description

def generate_diagram_code_ini_g(domain_name, problem_description, initial_state, goal_state, model, temp, save_path, goal=False, previous_attempt=None, error_message=None):
    example_ini_code_path = f"{domain_name}<PATH_REMOVED>"
    example_ini_image_path = f"{domain_name}<PATH_REMOVED>"
    example_ini_state_path = f"{domain_name}<PATH_REMOVED>"
    example_ini_diagram_reasoning_path = f"{domain_name}<PATH_REMOVED>"
    
    example_goal_code_path = f"{domain_name}<PATH_REMOVED>"
    example_goal_image_path = f"{domain_name}<PATH_REMOVED>"
    example_goal_state_path = f"{domain_name}<PATH_REMOVED>"
    
    if os.path.exists(example_ini_code_path) and os.path.getsize(example_ini_code_path) > 0:
        with open(example_ini_code_path, "r") as f:
            example_ini_code = f.read() 
            
        with open(example_ini_state_path, "r") as f:
            hypothetical_ini_state = f.read()
            
        with open(example_ini_diagram_reasoning_path, "r") as f:
            example_ini_diagram_reasoning = f.read()
    else:
        example_ini_code = None
        hypothetical_ini_state = None
        example_ini_diagram_reasoning = None
            
    if os.path.exists(example_goal_code_path) and os.path.getsize(example_goal_code_path) > 0:
        with open(example_goal_code_path, "r") as f:
            example_goal_code = f.read() 
            
        with open(example_goal_state_path, "r") as f:
            hypothetical_goal_state = f.read()
    else:
        example_goal_code = None
        hypothetical_goal_state = None


    state = goal_state if goal else initial_state
    state_type = "goal" if goal else "initial"
    
    prompt_parts = [f"""
        Consider the following problem description:
        {problem_description}
        
        The goal is to write a matplotlib script to visualize a state in the problem domain above."""]
    
    if not goal: 
        prompt_parts.extend(
        [f"""Here's an example implementation for the hypothetical state below:
        {hypothetical_ini_state}
        
        The matplotlib implementation:
        {example_ini_code}
        
        The reasoning behind the implementation for this arbitrary instance and an explanation of the meaning behind each shape, color, different sizes and shape is given below:
        {example_ini_diagram_reasoning}
        
        This code produced the following diagram for the random hypothetical state above:""",
        {"image_path": example_ini_image_path} if example_ini_code else ""] +
        [f"""Now I want you to implement a similar diagram but specific to the objects of the state below:
        {state.state_nl}
        """])
    
    else:  
        if example_ini_code and example_goal_code:
            prompt_parts.extend([f"""
            Here's the implmentations for an initial and goal state pair for a hypothetical insatnce of this problem:
            Hypothetical initial state:
            {hypothetical_ini_state}
            
            Hypothetical goal state:
            {hypothetical_goal_state}
            
            The target matplotlib code for the hypothetical initial state:
            {example_ini_code}
            
            The reasoning behind the implementation for this arbitrary instance and an explanation of the meaning behind each shape, color, different sizes and shape is given below:
            {example_ini_diagram_reasoning}
            
            The code generated the diagram below:""",
            {"image_path": example_ini_image_path},
            
            f"""The target matplotlib code for the hypothetical goal state:
            {example_goal_code}
            
            which generated the diagram below:""",
            {"image_path": example_goal_image_path},
            
            f"""
            Now I want you to generate the matplotlib code to draw the diagram specified to the goal state of the instance below (Note if there are multiple possible states for a given object, include that information as text at the bottom of the diagram.):"""]) 
        
        prompt_parts.extend([f"""The following is a description of the initial state of the instance:
        {initial_state.state_nl}
        This is a visualization of the initial state we previously acquired:
        """,
        {"image_path": initial_state.diagram_picture},
        f"""The follwoing is the goal state. I want you to implement diagram code specific to the objects of this goal state below:
        {state.state_nl}
        Note: If there is no information provided for an object in the goal state description but the object is present in the initial state: If you can make a plausible default position/status exists for that object, draw the object at the default position with the default status. Otherwise, write a textual description at the bottom of the page about the object with status no constraint.
        """
        ])
    
    prompt_parts.extend([
        f"The following is diagram encoding of the {state_type} state which you should implment a matplotlib script for. Diagrma encoding is a collection of descriptions of the relative/absolute position, relative/absolute size, status (includes color if relevant), and a text identifier for each object the state:",
        state.diagram_encoding,
        "Generate the matplotlib code that draws a diagram for the {state_type} state where each object is a shape with the size and position as described in the diagram encoding, and inside the rectangle is the text identifier and status of the object. Try to visualize the status if possible (for example using a color). The text identifiers should be easily readable and must not protrude the bounds of the shape of the object. Ensure the contrast between the text and the background of the text is as high as possible.",
        f"The code must save the figure at the following location: {save_path}",
        "The final answer must be enclosed in",
        "```code",
        "<final python code, no additional text>",
        "```"
    ])
    if previous_attempt and error_message:
        prompt_parts.extend([
            "This is the previous attempt:",
            previous_attempt,
            "This attempt failed to execute with the following error:",
            error_message
        ])
    response = get_model_response(prompt_parts, model, temp=temp)
    diagram_code = extract_content(response, "code", remove_new_lines=False)
    return diagram_code

def is_unique_action(problem_description, current_state, action_taken, new_state_nl, model):
    prompt_parts = [f"""
    Consider the following problem description:
    {problem_description}
    We are currently at the following state:
    {current_state.state_nl}
    The goal is to determine whether the action below is unique and leads to a new child state compared to the other child states we have explored previously.
    The new action chosen is:
    {action_taken}
    New state description is:
    {new_state_nl}
    The following is a list of actions previously explored:
    """]
    for child in current_state.child_states:
        prompt_parts.append(f"Action: {child.action_taken}, Child State: {child.state_nl}")
    prompt_parts.extend([
        "Is the new action and resulting child state unique compared to the child states previously explored? Answer with 'yes', if the new action is not in the list of previously explored actions, 'no' otherwise. Explain your answer, and report the final answer in the following format:",
        "```yes_no",
        "<final answer, yes or no>",
        "```"
    ])
    response = get_model_response(prompt_parts, model)
    is_unique = extract_content(response, "yes_no")
    prompt_parts.append("""Make sure to follow these steps when returning your answer:
                        write ```yes_no
                        then write your anwer yes or no
                        finally write ````""")
    if is_unique == None:
        response = get_model_response(prompt_parts, model)
        is_unique = extract_content(response, "yes_no")
    
    is_unique = is_unique if is_unique else False
    return is_unique

def check_action_path(problem_description, initial_state, current_state, goal_state, actions, possible_actions, model):
    prompt_parts = [
        f"Consider the following problem description:\n{problem_description}",
        f"This is a list of valid actions for this problem:\n{str(possible_actions)}",
        "This is the initial state of the problem:",
        initial_state.state_nl,
        "Below is a visualization of the initial state:",
        {"image_path": initial_state.diagram_picture},
        "This is the goal state of the problem:",
        goal_state.state_nl,
        "Below is a visualization of the goal state:",
        {"image_path": goal_state.diagram_picture},
        "We believe we are currently at the following state of the problem:",
        current_state.state_nl,
        "The current state is visualized below:",
        {"image_path": current_state.diagram_picture},
        "This is the diagram encoding of the current state (Diagram encoding encodes a geometrical layout of the scene. It is a collection of statements describing the relative/absolute position, relative/absolute size, status (includes color if relevant), and a text identifier for each object in the problem):",
        current_state.diagram_encoding,
        "Here are the sequence actions we took from the initial state to arrive at this state:",
        str(actions),
        "The parent state of the current_state is:",
        current_state.parent.state_nl,
        "Here is a visualization (diagram) of the parent state:",
        {"image_path": current_state.parent.diagram_picture},
        "And the diagram encoding of the parent state is:",
        current_state.parent.diagram_encoding,
        "Your task is to determine the following:",
        "Determine whether choosing the last action makes sense with respect to the domain constraints and actions taken previously.",
        "Determine if taking this action violates is allowed by considering each precondition of the action and checking to see if it is satisfied in the parent diagram and diagram encoding. Also you have to iterate through current state's diagram encoding to make sure the affects of the action are correctly updated. For this you hvae to think carefully and step by step. You have to explain your reasoning in depth and explain wether each precondition is satisfied or not."
        "Is taking the last action logical for achieving the goal state? Is it getting us closer to the goal state? If there is no way for us to reach the goal state fromt he current state, then the action is invalid.",
        "Are the natrual language and diagram encoding description of the state and the diagram encoding of the state is accurate given the sequence of actions taken from the initial state?",
        "Is the visualization of the currect state accurate? Is the status of all objects correctly visualized with respect to the sequence of actions taken? Important: for the diagram to be valid all objects in the initial state must be present in the current state's visualization",#Note the visualizations might have errors.",
    ]
    prompt_parts.extend([
        "At the end, answer with 'yes' if the last action and the current state's natural langauge description and visualization are all valid. Answer 'no' otherwise. Report the final yes-no answer in the following format:",
        "```yes_no",
        "<final answer, yes or no>",
        "```",
        "Then provide a short explanation why the action or the current state is valid or invalid (be very specific if the action is invalid, e.g. say which precondition is not met and for which action), in the following format:",
        "```answer_text",
        "<your reasoning>",
        "```",
        "Explain your reasoning and think step by step above the code blocks.",

        
        #"If you answered 'no' to the question above, identify the action that violated the constraints by its position, counting from 0 starting from the last action in the path.",
        #"```number",
        #"<final answer, the index of the invalid action>",
        #"```"
    ])
    response = get_model_response(prompt_parts, model)
    is_valid = extract_content(response, "yes_no")
    reasoning = extract_content(response, "text")
    #invalid_action = None
    #if is_valid == False:
        #invalid_action = int(extract_content(response, "number"))
   
    if is_valid == None:
        prompt_parts.append("""Make sure to follow these steps when returning your answer:
                        write ```yes_no
                        then write your anwer yes or no
                        then write ```
                        Similarly return a summary of your reasoning by first writing 
                        ```answer_text
                        then write your reaosning
                        and finally write ```""")
        response = get_model_response(prompt_parts, model)
        is_valid = extract_content(response, "yes_no")
        reasoning = extract_content(response, "text")
        #if not is_valid:
            #invalid_action = int(extract_content(response, "number"))
    
        is_valid = is_valid if is_valid != None else False
    #print(response)
    return is_valid, reasoning#, invalid_action 

def test_physical_plausibility(problem_description, initial_state, model):
    prompt_parts = [
        "Consider the following problem description:",
        problem_description,
        "The following is a description of the initial state of the problem:",
        initial_state.state_nl,
        "This is a suggested visualization of the initial state:",
        {"image_path": initial_state.diagram_picture},
        """I want you to analyze the visualization above and answer these questions with an explanation.
        Does this diagram make sense physically?
        Is the relative distance of the objects accurate (for example say no if an object is randomly hovering in the air)?
        Is the relative placement of objects accurate according to the problem statement and the description of the initial state?
        Is any laws of physics or common sense broken by the demonstration?
        Is there any object missing from the scene?
        Provide a detailed explanation, think step by step, and conclude with a yes or no answer in the following format:""",
        "```yes_no",
        "<final answer, yes or no>",
        "```",
        "After writing the code block above, if your final answer was 'no' , provide a very short phrase describing the error. Be as specific as possible (i.e. summerize your reasoning above in 1 sentence). Return your answer in the following format:",
        "```error",
        "<error description>",
        "```"
    ]
    response = get_model_response(prompt_parts, model)
    validity = extract_content(response, "yes_no")
    error_description = extract_content(response, "error")
    prompt_parts.append("""Make sure to follow these steps when returning your answer:
                        write ```yes_no
                        then write your anwer yes or no
                        then write ```
                        If you asnwered no above you must return the reason for a no answer int he follwoing steps: 
                        first write ```error
                        then write the error, ie why you answered no
                        lastly write ```""")
    if validity == None:
        response = get_model_response(prompt_parts, model)
        validity = extract_content(response, "yes_no")
        error_description = extract_content(response, "error")
    
    validity = validity if validity else False
    return validity, error_description



def check_goal_state(problem_description, current_state, goal_state, initial_state, model):
    prompt_parts = [
        "Given the following problem description:",
        problem_description,
        "where this is the initial state of the problem:",
        initial_state.state_nl,
        "And below is a visualization of the initial state:",
        {"image_path": initial_state.diagram_picture},
        "After taking the follwoing actions:",
        str(current_state.action_path),
        "We are currently in the following state:",
        current_state.state_nl,
        "This is a visualization of the current state:",
        {"image_path": current_state.diagram_picture},
        "The goal state is:",
        goal_state.state_nl,
        "This is a visualization of the goal state:",
        {"image_path": goal_state.diagram_picture},
        "Does the current state (only evluate the corrent state not any given example plans) satisfy the constraints of the goal state (ie is the current state our goal state as defined by the constraints)? Ensure the status of objects is as described in the goal state. Check every the status of every object and every constraint in the current state against the goal state description and diagram"
        "Pay close attention to the action path. Does the action path transform the initial state into the final state? The final answer should be reported in the following format as a yes (if all of the constraints are satisfied and the action path correctly trasnforms the initial state into the goal state) or no",
        "```yes_no",
        "<final answer, only yes or no>",
        "```",
        "Think step by step and explain your reasoning in depth above this code block."
    ]
    
    response = get_model_response(prompt_parts, model)
    validity = extract_content(response, "yes_no")
    prompt_parts.append("""Make sure to follow these steps when returning your answer:
                        write ```yes_no
                        then write your anwer yes or no
                        finally write ```""")
    if validity == None:
        response = get_model_response(prompt_parts, model)
        validity = extract_content(response, "yes_no")
    return validity if validity else False