

def code_gen_prompt(domain_pddl, observation, goal, spec, skill_code, available_skills, object_list_position, feedback, prev_code, exploration_knowledge, frozen_code_part = ""): 
    prompt_template = """
    You are an AI assistant tasked with generating executable Python code that controls a robot in a simulated environment. The goal is to ensure the robot can achieve the specified objective by executing a sequence of actions (plan). Your task is to complete the executable code using the provided inputs and by implementing any missing skills.

    ### **Input Details:**  
    1. **Domain PDDL** (Planning Domain Definition Language):  
       The PDDL domain describes the available actions and predicates in the environment. It includes information about actions (such as `move`, `pick`, `drop`) and their parameters, preconditions, and effects.  
       - Example structure:
         ```
         (:action pick
            :parameters (?obj ?room ?gripper)
            :precondition (and (ball ?obj) (room ?room) (free ?gripper) (at ?obj ?room) (at-robby ?room))
            :effect (and (carry ?obj ?gripper) (not (at ?obj ?room)) (not (free ?gripper)))
         )
         ```  
       Provided domain PDDL:
       ```
       {domain_pddl}
       ```

    2. **Observation (Initial State Description)**:  
       The observation represents the initial state of the environment in PDDL format. This typically includes the locations of objects, the robot’s initial location, and other relevant information. The observation does not include the goal state but provides context for the planning process.  
       - Example structure:
         ```
         (define (problem initial-problem)
            (:objects ball1 ball2 rooma roomb)
            (:init (at-robby rooma) (at ball1 rooma) (at ball2 rooma))
         )
         ```  
       Provided observation:
       ```
       {observation}
       ```

    3. **Goal (Natural Language Description)**:  
       The goal specifies what the robot must accomplish in the environment, described in plain language. It could describe moving objects to specific locations, interacting with objects, or completing complex tasks.  
       - Example: "Move all balls from Room A to Room B and return to the starting point."  
       Provided goal:
       ```
       {goal}
       ```

    4. **Specification (Code Generation Guidelines)**:  
       The specification outlines the exact requirements and constraints for generating the executable Python code. The generated code must strictly follow the specification guidelines to ensure correctness and consistency.  
       Provided specification:
       ```
       {spec}
       ```

    5. **Skill Code (Python Implementations of Actions)**:  
       The skill code consists of Python implementations of some of the available actions (such as `move`, `pick`). Each skill function defines how the robot performs a particular action. 
       - Example:
         ```python
         def move(env, from_room, to_room):
             # Implementation of move action
             pass

         def pick(env, obj, room, gripper):
             # Implementation of pick action
             pass
         ```  
       Provided skill code:
       ```
       {skill_code}
       ```

    6. **Executable Code Skeleton**:  
       A partial Python function that sets up the environment and outlines the overall control flow for executing the oracle plan. The skeleton may include some predefined logic (like initializing the environment and defining typed literals) but requires completion, handling control flow to execute the plan.  
       - Skeleton structure:
         ```python
# run_skeleton_task.py (Fully Generic Skeleton)
#
# Description:
#  - This code is a skeleton that shows how to set up an environment, retrieve
#    object positions, and execute a plan without being tied to a specific skill
#    like pick or place.
#  - Use this as a base for creating your own task logic.
#  - The sections that are task-dependent (e.g., manipulation steps) are marked
#    with TODO or placeholders.

import numpy as np
from pyrep.objects.shape import Shape
from pyrep.objects.proximity_sensor import ProximitySensor

from env import setup_environment, shutdown_environment

from skill_code import * (you don't need to redefine primitives like move, pick, place)

from video import init_video_writers, recording_step, recording_get_observation

# This module is assumed to provide object positions or other info about the environment
from object_positions import get_object_positions


def run_skeleton_task():
    '''Generic skeleton for running any task in your simulation.'''
    print("===== Starting Skeleton Task =====")
    
    # === Environment Setup ===
    env, task = setup_environment()
    try:
        # Reset the task to its initial state
        descriptions, obs = task.reset()

        # (Optional) Initialize video writers for capturing your simulation
        init_video_writers(obs)

        # Wrap the task steps for recording (if needed)
        original_step = task.step
        task.step = recording_step(original_step)
        original_get_obs = task.get_observation
        task.get_observation = recording_get_observation(original_get_obs)

        # === Retrieve Object Positions ===
        # Example usage: positions = 'object_1': (0,0,0), 'object_2': (1,1,1), ... (dictionary format)
        positions = get_object_positions()

        # TODO: Replace these placeholder keys with your actual objects
        # object_1_pos = positions['object_1']
        # object_2_pos = positions['object_2']

        # === Example Task Plan (Generic) ===
        #
        # Here’s where you would define the logical steps of your task. For instance:
        #
        # 1) Approach and manipulate object_1
        # 2) Move or interact with object_1 in some way
        # 3) Approach and manipulate object_2
        # 4) Complete the task
        #
        # Since this is a skeleton, we don't include any details on "pick" or "place" 
        # or other skills. Instead, you’d add them here if/when needed.
        #
        # Example (commented out):
        #
        # print("[Task] Manipulating object_1 at:", object_1_pos)
        # obs, reward, done = pick(
        #     env,
        #     task,
        #     target_pos=object_1_pos,
        #     approach_distance=0.15,
        #     max_steps=100,
        #     threshold=0.01,
        #     approach_axis='z',
        #     timeout=10.0
        # )
        #
        # if done:
        #     print("[Task] Task ended after manipulating object_1!")
        #     return
        #
        # print("[Task] Placing object_1 at:", object_2_pos)
        # obs, reward, done = place(
        #     env,
        #     task,
        #     target_pos=object_2_pos,
        #     approach_distance=0.15,
        #     max_steps=100,
        #     threshold=0.01,
        #     approach_axis='z',
        #     timeout=10.0
        # )
        # if done:
        #     print("[Task] Task completed successfully! Reward:", reward)
        # else:
        #     print("[Task] Task not completed yet (done=False).")

        # TODO: Continue adding more plan steps as needed...

    finally:
        # Always ensure the environment is properly shutdown
        shutdown_environment(env)

    print("===== End of Skeleton Task =====")


if __name__ == "__main__":
    run_skeleton_task()
         ``` 

    7. **Available Skill Names (List of Skill Names)**:  
       A list of action names that can potentially be used in the plan. These names correspond to the actions described in the domain PDDL (e.g., `move`, `pick`, `drop`). 
       - Example: `['move', 'pick', 'drop', 'place']`  
       Provided available skill names:
       ```
       {available_skills}
       ```

   8. **Object List**:
      {object_list_position}
      A list of object names in the environment.

   9. **Feedback (Corrections and Issues in Previous Code Generation):**  
       The feedback section contains identified issues in the previous generated code and necessary corrections. If feedback is none, it means there are no issues in the previous code generation.
       
       - Feedback: {feedback}
       - Previous Code: {prev_code}

   10. **Exploration Knowledge**:
      If you think some observations are lacking, you can explore them through exploration knowledge.
      - Exploration Knowledge: {exploration_knowledge}

   ### **Already Validated (Frozen) Code**:
    - Keep everything in [Frozen Code Start] ... [Frozen Code End] exactly the same.
    - Generate or modify the code *after* the Frozen Code region so that it achieves the goal and follows the specification.
    - Do not use code blocks in your final answer. Provide plain text.
   [Frozen Code Start]
    {frozen_code_part}
    [Frozen Code End]

   

Task Requirements:
	1.	Complete the Executable Code:
	•	Fill in the missing portions of the executable code skeleton by incorporating only the provided (predefined) skills and control flow logic.
	•	Ensure that the generated code can execute the oracle plan step-by-step and achieve the described goal.
	2.	Preserve the Skeleton Imports:
	•	Do not modify or remove any of the existing imports in the provided skeleton (skill_code, env, video, etc.).
	•	Use the externally provided functions and modules exactly as they are, without defining any new skill functions.

⸻

Guidelines for Implementation:
	•	Use Only Predefined Skills:
	•	The actions (move, pick, drop, place, etc.) provided in skill_code must be used directly.
	•	Do not create new skill/action functions; simply call the existing ones where needed in the plan logic.
	•	Error Handling:
	•	Gracefully handle exceptions (e.g., type mismatches, missing objects) during action execution.
	•	Naming Conventions:
	•	Maintain consistent naming for variables, objects, and other references throughout the code.
	•	Plan Execution:
	•	The generated executable code should dynamically handle actions in the provided oracle plan and achieve the specified goal state.
	•	Output Format:
	•	Do not use code blocks. Provide the generated code in plain text format.

⸻

Additional Explanations for External Modules:
	•	env (Environment Setup)
	•	Contains functions like setup_environment() and shutdown_environment() which manage the creation and teardown of the RLBench simulation.
	•	You should call these functions at the beginning and end of your main task logic to ensure proper environment handling.
	•	skill_code (Skills / Actions)
	•	Provides the implementations for various skill actions (e.g., move, pick, drop, place, etc.) used by the robot in the simulation.
	•	You must call these pre-existing skill functions directly to perform the necessary steps in your task.
	•	Do not create or redefine skill functions.
	•	video (Recording / Observation)
	•	Offers helper functions (e.g., init_video_writers, recording_step, recording_get_observation) to record or capture the simulation steps.
	•	Typically used to generate debug videos or logs of the robot’s activities.
	•	object_positions (Environment Information)
	•	Contains utilities like get_object_positions() to retrieve positional or state information about objects in the environment.
	•	Use these for tasks that require precise knowledge of object locations.
     
**Important**:
- Do not use code block. Output should be in plain text format.
    """
    
    return prompt_template.format(
         domain_pddl=domain_pddl,
         observation=observation,
         goal=goal,
         skill_code=skill_code,
         spec=spec,
         available_skills=available_skills,
         object_list_position=object_list_position,
         feedback=feedback,
         prev_code=prev_code,
         exploration_knowledge=exploration_knowledge,
         frozen_code_part=frozen_code_part
    )


def realworld_code_gen_prompt(domain_pddl, observation, goal, spec, skill_code, available_skills, object_list_position, feedback, prev_code, exploration_knowledge, frozen_code_part = "", plan_content = None): 
   #  print("realworld_code_gen_prompt")
    prompt_template = """
    You are an AI assistant tasked with generating executable Python code that controls a robot in a simulated environment. The goal is to ensure the robot can achieve the specified objective by executing a sequence of actions (plan). Your task is to complete the executable code using the provided inputs and by implementing any missing skills.

    ### **Input Details:**  
    1. **Domain PDDL** (Planning Domain Definition Language):  
       The PDDL domain describes the available actions and predicates in the environment. It includes information about actions (such as `move`, `pick`, `drop`) and their parameters, preconditions, and effects.  
       - Example structure:
         ```
         (:action pick
            :parameters (?obj ?room ?gripper)
            :precondition (and (ball ?obj) (room ?room) (free ?gripper) (at ?obj ?room) (at-robby ?room))
            :effect (and (carry ?obj ?gripper) (not (at ?obj ?room)) (not (free ?gripper)))
         )
         ```  
       Provided domain PDDL:
       ```
       {domain_pddl}
       ```

    2. **Observation (Initial State Description)**:  
       The observation represents the initial state of the environment in PDDL format. This typically includes the locations of objects, the robot’s initial location, and other relevant information. The observation does not include the goal state but provides context for the planning process.  
       - Example structure:
         ```
         (define (problem initial-problem)
            (:objects ball1 ball2 rooma roomb)
            (:init (at-robby rooma) (at ball1 rooma) (at ball2 rooma))
         )
         ```  
       Provided observation:
       ```
       {observation}
       ```

    3. **Goal (Natural Language Description)**:  
       The goal specifies what the robot must accomplish in the environment, described in plain language. It could describe moving objects to specific locations, interacting with objects, or completing complex tasks.  
       - Example: "Move all balls from Room A to Room B and return to the starting point."  
       Provided goal:
       ```
       {goal}
       ```

    4. **Specification (Code Generation Guidelines)**:  
       The specification outlines the exact requirements and constraints for generating the executable Python code. The generated code must strictly follow the specification guidelines to ensure correctness and consistency.  
       Provided specification:
       ```
       {spec}
       ```

    5. **Skill Code (Python Implementations of Actions)**:  
       The skill code consists of Python implementations of some of the available actions (such as `move`, `pick`). Each skill function defines how the robot performs a particular action. 
       - Example:
         ```python
            0: lambda: self._execute_pick_and_place(pose0, pose1, gripper_force, axis),
            1: lambda: self.primitive_skill.execute_pick(gripper_force=gripper_force, axis=axis),
         ```  
       Provided skill code:
       ```
       {skill_code}
       ```

    6. **Executable Code Skeleton**:  
       
       A partial Python function that sets up the environment and outlines the overall control flow for executing the oracle plan. The skeleton may include some predefined logic (like initializing the environment and defining typed literals) but requires completion, handling control flow to execute the plan.  
      In the example below, the axis parameter determines the approach direction: 0 indicates approach along the x-axis, 1 along the y-axis, and 2 along the z-axis. When grasping objects like drawer handles, the robot should approach along the x-axis, whereas for general objects, it should approach along the z-axis.
      You should redefine the def Long_horizon_task2_oracle_seq(): function within the skeleton. However, you must use only the given objects and the provided skills, and the structure of skill function calls must follow the same format as in the skeleton.
      Skill function must be called in the same way as in the skeleton.(robot.execute_movement(...))
       - Skeleton structure:
         ```python
# run_skeleton_task.py (Fully Generic Skeleton)
#
# Description:
#  - This code is a skeleton that shows how to set up an environment, retrieve
#    object positions, and execute a plan without being tied to a specific skill
#    like pick or place.
#  - Use this as a base for creating your own task logic.
#  - The sections that are task-dependent (e.g., manipulation steps) are marked
#    with TODO or placeholders.

from Object_list import Object
from skill_mapping import RobotController

def Long_horizon_task2_oracle_seq():

    robot = RobotController()

    robot.execute_movement(Object["drawer_top_handle"], mode = 1, axis = 0)
    robot.execute_movement(Object["drawer_top_handle"], mode = 4, distance=0.1)

    robot.execute_movement(Object["trash"], mode = 1, axis = 2)
    robot.execute_movement(mode = 8)
    robot.execute_movement(Object["trash_bin"], mode = 2, axis = 2)
    robot.execute_movement(mode = 8)

    robot.execute_movement(Object["yellow_dice"], mode = 1, axis = 2)
    robot.execute_movement(mode = 8)
    robot.execute_movement(Object["drawer_top_place_left"], mode = 2, axis = 2)
    robot.execute_movement(Object["green_dice"], mode = 1, axis = 2)
    robot.execute_movement(mode = 8)
    robot.execute_movement(Object["drawer_top_place_right"], mode = 2, axis = 2)
    robot.execute_movement(mode = 8)



def main():
    Long_horizon_task2_oracle_seq()
 
if __name__ == "__main__":
    main()

         ``` 

    7. **Available Skill Names (List of Skill Names)**:  
       A list of action names that can potentially be used in the plan. These names correspond to the actions described in the domain PDDL (e.g., `move`, `pick`, `drop`). 
       - Example: `['move', 'pick', 'drop', 'place']`  
       Provided available skill names:
       ```
       {available_skills}
       ```

   8. **Object List**:
      {object_list_position}
      A list of object names in the environment.

   9. **Feedback (Corrections and Issues in Previous Code Generation):**  
       The feedback section contains identified issues in the previous generated code and necessary corrections. If feedback is none, it means there are no issues in the previous code generation.
       
       - Feedback: {feedback}
       - Previous Code: {prev_code}

   10. **Exploration Knowledge**:
      If you think some observations are lacking, you can explore them through exploration knowledge.
      - Exploration Knowledge: {exploration_knowledge}
   
   11. **Plan Content**:
      The plan content is the oracle plan that you should follow. You can use the skills in the plan content to generate the code.
      - Plan Content: {plan_content}

   ### **Already Validated (Frozen) Code**:
    - Keep everything in [Frozen Code Start] ... [Frozen Code End] exactly the same.
    - The [Frozen Code Start] and [Frozen Code End] markers are only for reference and should not be included in the actual code. Only the content between them should be included in your implementation.
    - Generate or modify the code *after* the Frozen Code region so that it achieves the goal and follows the specification.
    - Do not use code blocks in your final answer. Provide plain text.
   [Frozen Code Start]
    {frozen_code_part}
    [Frozen Code End]

   

Task Requirements:
	1.	Complete the Executable Code:
	•	Fill in the missing portions of the executable code skeleton by incorporating only the provided (predefined) skills and control flow logic.
	•	Ensure that the generated code can execute the oracle plan step-by-step and achieve the described goal.
	2.	Preserve the Skeleton Imports:
	•	Do not modify or remove any of the existing imports in the provided skeleton.
	•	Use the externally provided functions and modules exactly as they are, without defining any new skill functions.

⸻

Guidelines for Implementation:
	•	Use Only Predefined Skills:
	•	The actions (move, pick, drop, place, etc.) provided in skill_code must be used directly.
	•	Do not create new skill/action functions; simply call the existing ones where needed in the plan logic.
	•	Error Handling:
	•	Gracefully handle exceptions (e.g., type mismatches, missing objects) during action execution.
	•	Naming Conventions:
	•	Maintain consistent naming for variables, objects, and other references throughout the code.
	•	Plan Execution:
	•	The generated executable code should dynamically handle actions in the provided oracle plan and achieve the specified goal state.
	•	Output Format:
	•	Do not use code blocks. Provide the generated code in plain text format.

     
**Important**:
- Do not use code block. Output should be in plain text format.
    """
    
    return prompt_template.format(
         domain_pddl=domain_pddl,
         observation=observation,
         goal=goal,
         skill_code=skill_code,
         spec=spec,
         available_skills=available_skills,
         object_list_position=object_list_position,
         feedback=feedback,
         prev_code=prev_code,
         exploration_knowledge=exploration_knowledge,
         frozen_code_part=frozen_code_part,
         plan_content=plan_content
    )
