import pickle
from LTL_implementer import OBS_VARIABLES_LIST, ACTION_VARIABLES_LIST, SOFT_LTL_RULES_SAYCAN
import asyncio
from utils import get_response_4o, get_response_4o_mini, get_response_o3


def get_active_keys(one_hot_dict):
    return [key for key, value in one_hot_dict.items() if value == 1]


def format_trajectory_for_critic(trajectory):
    lines = []
    for i in range(len(trajectory)):
        step = trajectory[i]
        next_step = trajectory[i + 1] if i + 1 < len(trajectory) else None

        obs = step["observation"]
        obs_next = next_step["observation"] if next_step else None

        action_keys = [k for k, v in step["action"].items() if v]
        result = "success" if not (next_step and next_step.get("previous_error")) else f"Error: {next_step['previous_error']}"

        # Show only the APs that changed between this step and the next
        if obs_next:
            changed_aps = [k for k in obs if obs.get(k) != obs_next.get(k)]
        else:
            changed_aps = []

        lines.append(f"Step {i + 1}:\n"
                     f"- Observations: {', '.join([k for k in obs if obs[k]]) or 'None'}\n"
                     f"- Action: {', '.join(action_keys) if action_keys else 'action_idle'}\n"
                     f"- Changed APs: {', '.join(changed_aps) if changed_aps else 'None'}\n"
                     f"- Result: {result}\n")
    
    return "\n".join(lines)





async def analyze_trace(filename):
    with open(filename, 'rb') as f:
        trace = pickle.load(f)

    print(format_trajectory_for_critic(trace))


    intro_prompt = f"""
    You are an expert critic observing the trajectory of a Minecraft agent. The goal of the agent is to mine a diamond.

    You are given:
    1) A list of atomic observation and action variables
    2) A list of failures that occurred in trajectories
    3) Existing LTL rules implemented

    You will be given a series of steps taken by the agent, including observations, actions, and success/failure of the action with an error message.
    Your task is to analyze the trajectory and provide LTL laws that constrain the agent's actions in order to boost efficiency and performance. 
    The laws should be in the form of LTL formulas, and you should provide a brief explanation of each law. The boolean variables used in the laws are defined as follows:

    Observation Variables: 
    {", ".join(OBS_VARIABLES_LIST)}

    The observations in the atomic proposition space are described as follows:

    - `obs_has_x` corresponds to having the item `x` in the inventory of the agent. 
    - `obs_near_crafting_table` or `obs_near_furnace` define if the agent is within an interacting distance of a crafting table or furnace.
    - `obs_has_x_equipped` corresponds to an object `x` (e.g., an iron pickaxe) actively equipped.
    - Only one item can be equipped at any point in time. 
    - You may propose additional observation variables if needed to express useful rules.

    Action Variables:
    {", ".join(ACTION_VARIABLES_LIST)}

    The actions the agent can perform are limited to a few types:

    - `action_mine_x`: mines the item `x`. Certain blocks require certain tools:
        - stone: wood pickaxe or better
        - iron: stone pickaxe or better
        - diamond: iron pickaxe or better
    - `action_craft_x`: crafts an item `x`, if prerequisites and (if needed) a crafting table are present.
    - `action_smelt_iron`: smelts raw iron into ingots using fuel and a furnace.
    - `action_equip_x`: equips a tool for mining. Only one tool can be equipped at a time.
    - `action_explore`: used to find resources not currently visible.
    - `action_place_x`: places an item like a crafting table or furnace to enable usage.
    """
    main_prompt = f"""
    You are a symbolic critic observing the trajectory of a Minecraft agent. The agent is inefficient, often repeats work, and occasionally causes errors like trying to mine without the right tool or crafting without the ingredients.

    GOAL OF AGENT: MINE A DIAMOND  
    YOUR GOAL: PROPOSE LTL LAWS THAT PROMOTE THE AGENT'S PROGRESS TOWARD THIS GOAL AND PREVENT INEFFICIENCIES.

    THINK OF THE BASIC TASK GRAPH REQUIRED TO MINE A DIAMOND, and PROPOSE LTL LAWS TO GUIDE THE AGENT ALONG THAT GRAPH.

    Before forcing any action, think of checking if the subgoals to do that action are met.


    ### Your task:
    1. Decompose the task of mining a diamond into symbolic subgoals: acquiring wood, crafting tools, smelting, equipping tools, etc.
    2. For each subgoal transition (e.g., "has stone => craft stone_pickaxe"), propose an **LTL law** that enables or encourages this step.
    3. Also identify any **errors** or **inefficiencies** in the trajectory. For each one, propose an LTL law to prevent that mistake in the future.
    4. Focus on writing LTL laws in the form `G(condition => X(action))`. Use observation variables for `condition`, and action variables for `action`.
    5. Avoid overly specific or redundant laws. Try to generalize from the plan, not just from individual steps.
    6. You may also propose **new boolean observation variables** if needed to express useful constraints (e.g., `obs_has_stone_pickaxe`, `obs_seen_diamond_block`).
    7. Your final set of LTL laws should include:
    - ≥3 rules that **encourage efficient, goal-aligned behavior**
    - ≥1 rule that **discourages observed inefficient behavior**

    ### Existing Inputs:
    - Existing LTL laws: {SOFT_LTL_RULES_SAYCAN}
    - Action and observation variables:
        - Actions: {", ".join(ACTION_VARIABLES_LIST)}
        - Observations: {", ".join(OBS_VARIABLES_LIST)}


    ### Agent Trajectory:
    {format_trajectory_for_critic(trace)}

    
    ### Output Format:
    Return the following in your response:

    ---

    ### Reasoning:
    1. **Plan Decomposition**: Write out the full high-level plan to mine a diamond as a sequence of symbolic subgoals (e.g., get wood → make planks → craft tools → smelt → equip → mine).
    2. **Plan  Conversion**:  List the sequence of obs props and action props that correspond to this plan
    2. **Positive Constraints**: For at least four transitions in this plan, propose a rule of the form `G(preconditions => X(useful_action))` that helps the agent complete the task efficiently.
    3. **Negative Constraints**: Identify any mistakes in the trajectory (e.g., crafting without ingredients, mining without tools), and write `G(bad_condition => X(¬bad_action))` rules to prevent them.
    4. **Coverage**: Ensure your rules cover multiple stages of the plan (not just early or late stages).
    5. **Reusability**: The rules should generalize and not rely on specific step numbers.


    ---
    Laws:
    - `efficiency_laws`: a list of LTL rules that **promote efficient behaviors**, each with a brief explanation.
    - `inefficiency_laws`: a list of LTL rules that **prevent mistakes**, each with a brief explanation.
    - Mention which part of the plan each law corresponds to.
    """



    print("="   * 40)
    response = await get_response_o3(intro_prompt, main_prompt)
    print(response)



async def handle_conflicts(filename):

    with open(filename, 'rb') as f:
        trace = pickle.load(f)



    intro_prompt = f"""
    You are an expert critic observing the trajectory of a Minecraft agent. The goal of the agent is to mine a diamond.

    Sometimes, the laws you impose are too constraining and prevent all possible actions. Your job is to break these deadlocks

    You are given:
    1) A list of atomic observation and action variables

    You will be given a particular timestep where the LTL laws led to no feasible action, and your task is to resolve the conflict by either modifying or deleting one of the LTL laws.
    
    The laws should be in the form of LTL formulas, and you should provide a brief explanation of each law. The boolean variables used in the laws are defined as follows:

    Observation Variables: 
    {", ".join(OBS_VARIABLES_LIST)}

    The observations in the atomic proposition space are described as follows:

    - `obs_has_x` corresponds to having the item `x` in the inventory of the agent. 
    - `obs_near_crafting_table` or `obs_near_furnace` define if the agent is within an interacting distance of a crafting table or furnace.
    - `obs_has_x_equipped` corresponds to an object `x` (e.g., an iron pickaxe) actively equipped.
    - Only one item can be equipped at any point in time. 
    - You may propose additional observation variables if needed to express useful rules.

    Action Variables:
    {", ".join(ACTION_VARIABLES_LIST)}

    The actions the agent can perform are limited to a few types:

    - `action_mine_x`: mines the item `x`. Certain blocks require certain tools:
        - stone: wood pickaxe or better
        - iron: stone pickaxe or better
        - diamond: iron pickaxe or better
    - `action_craft_x`: crafts an item `x`, if prerequisites and (if needed) a crafting table are present.
    - `action_smelt_iron`: smelts raw iron into ingots using fuel and a furnace.
    - `action_equip_x`: equips a tool for mining. Only one tool can be equipped at a time.
    - `action_explore`: used to find resources not currently visible.
    - `action_place_x`: places an item like a crafting table or furnace to enable usage.
    """

    for t in trace:
        main_prompt = f"""
        You are a symbolic critic observing the trajectory of a Minecraft agent. The agent is inefficient, often repeats work, and occasionally causes errors like trying to mine without the right tool or crafting without the ingredients.

        GOAL OF AGENT: MINE A DIAMOND
        YOUR GOAL: You are given a timestep where the LTL laws led to no feasible action, and your task is to resolve the conflict by either modifying or deleting one of the LTL laws.

        Given the set of observations, no actions are allowed in the given instance. Modify one or both of the laws to break this deadlock.

        ### Your task:
        1. First, reason about which rule is less useful given the constraints
        2. Modify that law
        3. Make sure there is atleast one feasible action in the given state after the modification of laws.

        ### Rules you should output:
        - Each rule should follow the form: `G(condition => X(action))`
        - Use **observation variables** (e.g., obs_has_x, obs_near_x, obs_equipped_x) in the `condition`.
        - Use **action variables** (e.g., action_mine_x, action_craft_x) in the `action`.
        - Conditions should reflect **states that actually occurred** in the trajectory, so the rule can generalize and not be overly specific or invalid. The corresponding action should either approve or disapprove of the corresponding action in the trajectory. 
        - Make sure the rules do not block **all** possible actions. The agent always needs at least one valid option.
        - Make sure to state which timesteps your rule is based on.
        - If necessary, propose **new observation variables** that could help express useful rules.



        ### Format:
        - Output the LTL laws as a list of strings, each in proper syntax
        - Then provide a brief explanation of each rule and how it prevents error or improves efficiency
        - Output as many laws as you deem necessary, forcing efficient actions and disallowing inefficient ones

        ### Example LTL Law:
        G(obs_has_raw_iron ∧ obs_near_furnace => X(action_smelt_iron))
        Explanation: Smelting iron early helps the agent craft a better pickaxe sooner.

        ### Inputs:

        - Existing LTL laws 
            Rule 2: G(obs_has_2x_plank & !obs_has_2x_stick -> X(action_craft_stick))
            Number of actions allowed: 1
            Filtered actions:
            ['action_craft_stick']
            Observation Propositions (obs_props):
            obs_has_plank: True
            obs_has_2x_plank: True
            obs_has_3x_plank: True
            obs_has_wood_pickaxe: True
            obs_has_stone_pickaxe: True
            obs_has_fuel: True
            obs_near_crafting_table: True
            obs_iron_in_chunk: True
            obs_coal_in_chunk: True
            obs_wood_pickaxe_equipped: True
            ==================================================
            Rule 7: G(obs_wood_pickaxe_equipped & !obs_has_3x_cobble -> X(action_mine_stone))
            Number of actions allowed: 1
            Filtered actions:
            ['action_mine_stone']
            Observation Propositions (obs_props):
            obs_has_plank: True
            obs_has_2x_plank: True
            obs_has_3x_plank: True
            obs_has_wood_pickaxe: True
            obs_has_stone_pickaxe: True
            obs_has_fuel: True
            obs_near_crafting_table: True
            obs_iron_in_chunk: True
            obs_coal_in_chunk: True
            obs_wood_pickaxe_equipped: True

        These rules are too constraining, modify them. 


        - Feasible actions per step are available (so do not block everything)
        - Action and observation variables:
            - Actions: {", ".join(ACTION_VARIABLES_LIST)}
            - Observations: {", ".join(OBS_VARIABLES_LIST)}


        Answer with the following three things

        1. Identify the action the agent should take given this state
        2. Identify which rules are blocking that action from happening
        3. Modify those rules.
        

        Modify one of the rules, or delete one of them, so that the agent can take a feasible action at this timestep.
        """

        #print(t['obs_props'])
        print("="   * 40)

        response = await get_response_4o(intro_prompt, main_prompt)
        print(response)

        x = input()

import argparse


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Process some variables.")
    parser.add_argument('--file', type=str, required=True, help='Name of the user')
    parser.add_argument('--mode', type=int, required=True, help='0 = analyze trace, 1 = handle conflicts')

    args = parser.parse_args()
    FILENAME = args.file

    # Run the async function to analyze the trace
    if( args.mode == 0):
        asyncio.run(analyze_trace(FILENAME))
    else:
        asyncio.run(handle_conflicts(FILENAME))


# # Load the first trace
# with open('saved_trajectory_2.pkl', 'rb') as f:
#     trace1 = pickle.load(f)



# # Print the result
# print("Combined Trace:")
# for i, step in enumerate(trace1):
#     print(f"\n=== Step {i + 1} ===")
#     print("Observation:")
#     print(get_active_keys(step["observation"]))
#     print("Action:")
#     print(get_active_keys(step["action"]))
#     print("Code:")
#     print(step["code"].strip())
#     if(get_active_keys(step["action"]) == []):
#         # Accept input and manually enter the action prop
#         print()
#         action_input = input("Enter the action (comma-separated): ")
#     print(f"Previous Error: {step['previous_error']}")

