import os
import subprocess
import json
import requests
API_KEY = os.environ.get("ARK_API_KEY")
url = "https://ark.cn-beijing.volces.com/api/v3/chat/completions"
API_KEY = os.environ.get("ARK_API_KEY") 
API_URL = "https://ark.cn-beijing.volces.com/api/v3/chat/completions"
room={"Kitchen","Living Room","Bedroom","Bathroom"}
room_id={"1","2","3","4"}
action_list={
    'OpenObject',
    'CloseObject',
    'PickupObject',
    'PutObject',
    'CleanObject',
    'HeatObject',
    'CoolObject',
    'ToggleObjectOn',
    'ToggleObjectOff',
    'SliceObject'}
RECEPTACLES = {
        'BathtubBasin',
        'Bowl',
        'Cup',
        'Drawer',
        'Mug',
        'Plate',
        'Shelf',
        'SinkBasin',
        'Box',
        'Cabinet',
        'CoffeeMachine',
        'CounterTop',
        'Fridge',
        'GarbageCan',
        'HandTowelHolder',
        'Microwave',
        'Pot',
        'StoveBurner',
        'SideTable',
        'ToiletPaperHanger',
        'TowelHolder',
        'Safe',
        'ArmChair',
        'Toilet',
        'Dresser',
        'LaundryHamper',
        'Toaster',
    }
Openable={
    'Box',
    'Cabinet',
    'Drawer',
    'Fridge',
    'Microwave',
    'Safe',
}
similar_category_map = {
    'Cup': ['Mug'],
    'Mug': ['Cup'],
    'Spoon': ['Ladle'],
    'Ladle': ['Spoon'],
    'SaltShaker': ['PepperShaker'],
    'PepperShaker': ['SaltShaker'],
    'ToiletPaperRoll': ['PaperTowelRoll'],
    'PaperTowelRoll': ['ToiletPaperRoll'],
    'Towel': ['HandTowel'],
    'HandTowel': ['Towel'],
    'Glassbottle': ['WineBottle'],
    'WineBottle': ['Glassbottle'],
    'Pencil': ['Pen'],
    'Pen': ['Pencil'],
    'CoffeeTable': ['SideTable', 'Desk'],
    'SideTable': ['CoffeeTable', 'DiningTable'],
    'Desk': ['DiningTable', 'CoffeeTable'],
    'DiningTable': ['Desk', 'SideTable']
}
objects={
'AlarmClock', 'Apple', 'AppleSliced', 'BaseballBat', 'BasketBall', 'Book', 'Bowl', 'Box', 'Bread', 'BreadSliced', 'ButterKnife', 'CD', 'Candle', 'CellPhone', 'Cloth', 'CreditCard', 'Cup', 'DeskLamp', 'DishSponge', 'Egg', 'Faucet', 'FloorLamp', 'Fork', 'Glassbottle', 'HandTowel', 'HousePlant', 'Kettle', 'KeyChain', 'Knife', 'Ladle', 'Laptop', 'LaundryHamperLid', 'Lettuce', 'LettuceSliced', 'LightSwitch', 'Mug', 'Newspaper', 'Pan', 'PaperTowel', 'PaperTowelRoll', 'Pen', 'Pencil', 'PepperShaker', 'Pillow', 'Plate', 'Plunger', 'Pot', 'Potato', 'PotatoSliced', 'RemoteControl', 'SaltShaker', 'ScrubBrush', 'ShowerDoor', 'SoapBar', 'SoapBottle', 'Spatula', 'Spoon', 'SprayBottle', 'Statue', 'StoveKnob', 'TeddyBear', 'Television', 'TennisRacket', 'TissueBox', 'ToiletPaper', 'ToiletPaperRoll', 'Tomato', 'TomatoSliced', 'Towel', 'Vase', 'Watch', 'WateringCan', 'WineBottle'
}
map_save_large_objects = [
    'ArmChair', 'BathtubBasin', 'Bed', 'Cabinet', 'Cart', 'CoffeeMachine', 'CoffeeTable', 'CounterTop',
    'Desk', 'DiningTable', 'Drawer', 'Dresser', 'Fridge', 'GarbageCan', 'Microwave', 'Ottoman', 'Safe',
    'Shelf', 'SideTable', 'SinkBasin', 'Sofa', 'StoveBurner', 'TVStand', 'Toilet'
]
# 将 receptacles 中不在 objects 里的元素加进去
merged_objects = sorted(objects.union(set(map_save_large_objects)))
rooms = {
    "Kitchen": {
        "ALL": [
            "Apple", "AppleSliced", "Bowl", "Bread", "BreadSliced", "ButterKnife",
            "Cabinet",  "CoffeeMachine", "CounterTop",  "Egg","Fork", "Fridge", "Kettle", "Knife", "Lettuce", "LettuceSliced",
            "LightSwitch", "Microwave", "Pan", "PepperShaker", "Plate", "Pot",
            "Potato", "PotatoSliced", "SaltShaker", "Sink", "Spatula", "Spoon",
            "StoveBurner", "StoveKnob", "Toaster", "Tomato", "TomatoSliced", "WineBottle"
        ],
        "some": [
            "Box","Chair", "Candle", "Cup", "Drawer","DiningTable","GarbageCan", "Glassbottle", "HousePlant",
            "KeyChain", "Laptop", "Mug", "Newspaper", "Painting", "PaperTowel",
            "PaperTowelRoll", "Pen", "Pencil", "Safe", "ScrubBrush", "SoapBottle",
            "SprayBottle", "Statue", "Towel", "Vase", "Watch", "WateringCan", "Window"
        ]
    },
    "Living Room": {
        "ALL": [
            "ArmChair", "Chair", "FloorLamp", "LightSwitch", "RemoteControl",
            "SideTable", "Sofa", "Television"
        ],
        "some": [
            "BaseballBat", "BasketBall", "Box", "Candle", "Cup", "Desk", "DeskLamp",
            "GarbageCan", "Glassbottle", "HousePlant", "KeyChain", "Laptop", "Mug",
            "Newspaper", "Ottoman", "Painting", "PaperTowel", "PaperTowelRoll", "Pen",
            "Pencil", "Pillow", "Safe", "Statue", "TeddyBear", "TennisRacket",
            "TissueBox", "Vase", "Watch", "WateringCan", "Window", "WineBottle"
        ]
    },
    "Bedroom": {
        "ALL": [
            "AlarmClock", "Bed", "Desk", "DeskLamp", "Dresser", "LightSwitch",
            "Laptop", "PaperTowel", "PaperTowelRoll", "Pen", "Pencil", "Pillow", "TeddyBear"
        ],
        "some": [
            "ArmChair", "BaseballBat", "BasketBall", "Box", "Candle", "Chair",
            "FloorLamp", "GarbageCan", "Glassbottle", "HousePlant", "KeyChain",
            "LaundryHamper", "LaundryHamperLid", "Mug", "Newspaper", "Ottoman",
            "Painting", "Poster", "RemoteControl", "Safe", "Statue", "Television",
            "TennisRacket", "TissueBox", "Towel", "Vase", "Watch", "WateringCan", "Window"
        ]
    },
    "Bathroom": {
        "ALL": [
            "BathtubBasin", "Cabinet", "Drawer", "Faucet", "HandTowel", "HandTowelHolder",
            "LightSwitch", "Plunger", "ShowerCurtain", "ShowerDoor", "ShowerGlass",
            "ShowerHead", "Sink", "SoapBar", "SoapBottle", "Toilet", "ToiletPaper",
            "ToiletPaperHanger", "Towel", "TowelHolder"
        ],
        "some": [
            "Box", "Candle", "CounterTop", "Cup", "GarbageCan", "Glassbottle",
            "LaundryHamper", "LaundryHamperLid", "Mug", "Newspaper", "Painting",
            "PaperTowel", "PaperTowelRoll", "Pen", "Pencil", "ScrubBrush", "SprayBottle",
            "Statue", "TissueBox", "Vase", "Watch", "Window"
        ]
    }
}
'''def call_ark_from_py36(user_input):
    payload = {
        "api_key": "6d67df38-2dde-4e4e-951a-22e1cd117ef1",  # Or use os.environ.get("ARK_API_KEY")
        "model": "deepseek-r1-250120",
        "messages": [
            {"role": "user", "content": user_input}
        ]
    }

    python38_path = os.path.expanduser("/scratch/lijh/.conda/envs/janus/bin/python")
    ark_runner_script = "ark_example.py"

    process = subprocess.Popen(
        [python38_path, ark_runner_script],
        stdin=subprocess.PIPE,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE
    )

    out, err = process.communicate(input=json.dumps(payload).encode("utf-8"))

    if err and err.strip():
        print("❌ Subprocess Error:", err.decode())
    else:
        result = json.loads(out.decode("utf-8"))
        print("✅ Ark Reply:", result)
        return result['result'].strip()'''
def call_ark_from_py36(user_input):
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {API_KEY}"
    }
    data = {
        "model": "deepseek-r1-250120",
        "messages": [
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": user_input}
        ],
        "temperature": 0.3
    }
    response = requests.post(API_URL, headers=headers, data=json.dumps(data))
    if response.status_code == 200:
        result = response.json()
        return result['choices'][0]['message']['content'].strip()
    else:
        print("调用失败：", response.status_code)
        print(response.text)
        return ""
def call_ark_from_py37(user_input):
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {API_KEY}"
    }
    data = {
        "model": "deepseek-r1-distill-qwen-32b-250120",
        "messages": [
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": user_input}
        ],
        "temperature": 0.3
    }
    response = requests.post(API_URL, headers=headers, data=json.dumps(data))
    if response.status_code == 200:
        result = response.json()
        return result['choices'][0]['message']['content'].strip()
    else:
        print("调用失败：", response.status_code)
        print(response.text)
        return ""
def rename(goal_long,goal_short,list_of_actions_s):
    template = """
You are an intelligent agent operating in a home-like virtual environment.


Here is what you know:

- smallObjects: {Objects1}
- largeObjects: {Objects2}

- Your current task is: {Task}
- Task description: {TaskDescription}
- Current initial plan: {InitialPlan}


---

### Objective

Your job is to **verify and correct the initial plan** based on the task description and the known objects.

---

### What You Must Do

1. **Check for Object Prediction Errors in the Initial Plan**
   - Carefully read the **task description**.
   - For every object mentioned in the **initial plan**, verify:
     - Is it consistent with the description?
     - Does it appear in the known object lists ({Objects1} or {Objects2})?
     - Is it a likely mistaken reference (e.g., confusing "cup" with "mug")?

2. **Fix Incorrect Object Names**
   - If any object in the plan seems incorrect or mismatched with the task description:
     - **Replace it** using the correct name from {Objects1} or {Objects2}.
     - The new object must appear in the known object list.
   - Only correct names based on strong evidence in the task description.



---

### Output Format

You must return a **single line** in the exact format:

`('target','action'),('target','action'),...`

This line must represent the **corrected and complete plan**.  
- Modify object names if they were wrong.  
  
- If the plan was already correct, return it as is.  
- **All object names must come from the  object lists.**
You **must not** change the number of steps in the plan.
- You **must not** change the action verbs.
- You **must not** change the order of actions.
- You may **only** replace object names based on the task description and known objects.
your object must be chosen from smallObjects and largeObjects 

---

### Example Reasoning Flow

- Initial plan says: `('Cup','PickObject')`
- Task description says: "Pick up the mug from the table"
- Known objects: `['Mug']`, no 'Cup' in the list → Replace 'Cup' with 'Mug'

---

Begin by correcting or confirming the initial plan based on task description and object lists.
"""


    prompt = template.format(
        Objects1=objects,
        Objects2=map_save_large_objects,

        Task=goal_short,
        Objects=merged_objects,
        TaskDescription=goal_long,

        InitialPlan=list_of_actions_s,
    )

    result=call_ark_from_py37(prompt)
    clean_plan= extract_plan_from_reply(result)
    return  clean_plan


import re

def rename_object(goal_long,goal_short,list_of_actions_pointer_s,list_of_actions_s,detected_list,step1):
    template = """
    You are an intelligent agent operating in a home-like virtual environment.

    This task will fail after 1000 steps if not completed.

    Here is what you know:

    - smallObjects: {Objects1}
    - largeObjects: {Objects2}
    - Your current task is: {Task}
    - Task description: {TaskDescription}
    - Current initial plan: {InitialPlan}
    - You have been stuck at: {Stuck} for {step1} steps
    - The detected list is: {detected_list}
    - similar_category_map: {similar_category_map}
    ---

    ### Your Task

    Your job is to check whether the **object name** in the **current action**, while possibly correct according to the `task` and `task description`, might have been **visually misidentified in the environment due to segmentation or perception errors**.

To handle this, you should:
- Use `similar_category_map` to find possible alternative object names that could have been confused with the current one.
- Compare them in light of:
  • The `task description`  
  • The known object lists (`smallObjects`, `largeObjects`)  
  • The `detected_list`

    **Important constraints:**
    1. You are **not allowed** to change the action verb.
    2. You are **not allowed** to insert, delete, or reorder any steps in the plan.
    3. You are **only allowed** to correct the object name **if** all of the following conditions are met:
       - The object name exists in the `similar_category_map`
       - A better alternative can be found **within the provided alternatives in `similar_category_map`**
       - The alternative object:
         • Matches the task description better, and  
         • **Appears in the `detected_list`**
         • Is present in either `smallObjects` or `largeObjects`

    4. If the object name is already correct, or no better candidate exists that satisfies all criteria, **keep it unchanged**.
    5. Return the revised plan in **this format**:
       ('object1', 'action1'), ('object2', 'action2'), ..., ('objectN', 'actionN')
       - The number of steps must be identical to the initial plan.
       - The action verbs and step order must remain unchanged.

    You must be extremely conservative when making changes — **if there is any doubt**, prefer to leave the object name as is.
    """
    prompt = template.format(
        Objects1=objects,
        Objects2=map_save_large_objects,
        Openable=Openable,
        Task=goal_short,
        Objects=merged_objects,
        TaskDescription=goal_long,
        step1=step1,

        InitialPlan=list_of_actions_s,
        Stuck=list_of_actions_s[list_of_actions_pointer_s],
        detected_list=detected_list,
        similar_category_map=similar_category_map,

    )

    result = call_ark_from_py36(prompt)
    clean_plan = extract_plan_from_reply(result)
    return clean_plan


def extract_plan_from_reply(reply):
    """
    从模型的回复中提取 Location 和 Plan。
    返回 (location, plan_list)，例如：
    ('Microwave', [('Bowl', 'PickupObject'), ('Microwave', 'OpenObject')])
    """
    # 提取 location（位于开头，直到第一个逗号）
    '''location_match = re.match(r"\s*([^,]+)\s*,", reply)
    location = location_match.group(1).strip() if location_match else "None"
'''
    # 提取 plan 中的所有 ('xxx','yyy') 结构
    plan_matches = re.findall(r"\('([^']+)'\s*,\s*'([^']+)'\)", reply)

    plan = [(obj.strip(), act.strip()) for obj, act in plan_matches]

    return plan
if __name__ == '__main__':
    # 用法举例
    initial_plan = [('Bowl', 'PickupObject'), ('Cabinet', 'OpenObject')]
    new_plan = [
        ('Microwave', 'OpenObject'),  # 1
        ('Bowl', 'PickupObject'),  # 2
        ('Drawer', 'OpenObject'),  # 3 ← 满足：新增 + 下接 Pickup
        ('Apple', 'PickupObject')  # 4
    ]


    print(get_new_open_action_positions(initial_plan, new_plan))