from my_agent.agent import Agent
from utils.llm import rules, prepend_history, query_llm
from levels.utils import convert_to_prompt, compute_dependency
import random
import copy
import re

class AuctionAgent(Agent):
    """
    An agent that participates in auctions for tasks.
    """
    def __init__(self, env, model, agent_id, total_num_agents, with_feedback, with_notes, look_ahead_steps):
        super().__init__(env, model, agent_id, total_num_agents, with_feedback, with_notes, look_ahead_steps)
        
        # Currency and bidding attributes
        self.currency = 100  # Initial currency amount
        self.assigned_tasks = []  # Tasks assigned to this agent
        self.completed_tasks = []  # Tasks completed by this agent
        self.failed_tasks = []  # Tasks failed by this agent
        
        # Performance tracking
        self.success_rate = 0.8  # Initial success rate estimate
        self.action_history = []  # History of actions taken
        self.action_success_history = []  # History of action successes
        
        # Location and state tracking
        self.current_location = None
        self.holding_item = None
        
    def update_state(self, obs):
        """Update the agent's internal state based on observation"""
        for agent in obs.agents:
            if agent['id'] == self.agent_id:
                self.current_location = agent['location']
                self.holding_item = agent['hold']
                break

    def step(self, obs, step, verbose=False, agent_tasks=None):
        """
        Executes a single step for the agent based on the given observation and assigned tasks.

        Args:
            obs (dict): The observation data for the current step.
            step (int): The current step number.
            verbose (bool, optional): If True, prints detailed debug information. Defaults to False.
            agent_tasks (list, optional): List of tasks assigned to this agent through auction. Defaults to None.

        Returns:
            list: A list of parsed actions generated by the agent.
        """
        # Update history with feedback and suggestions
        if self.with_feedback and step != 0:
            self.feedback = '-execution error messages:\n  --  ' + str(self.env.feedback) + '\n'

        if self.with_notes and step != 0:
            self.suggestions = '-execution suggestions:\n  --  ' + str(self.env.suggestions) + '\n'
            if 'agent ids cannot be the same' in self.feedback:
                self.suggestions += f'  --  You can only control and plan the actions for agent{self.agent_id}. \n'
        
        # Add auction-assigned tasks to the prompt with detailed instructions
        auction_tasks_prompt = ""
        if agent_tasks:
            # Prioritize tasks
            prioritized_tasks = self.prioritize_tasks(agent_tasks, obs)
            
            auction_tasks_prompt = f"-assigned tasks (in priority order):\n"
            
            # Group tasks by parent task for better organization
            tasks_by_parent = {}
            for task_id in prioritized_tasks:
                task_parts = task_id.split('_')
                parent_task = task_parts[0].rstrip('0123456789')  # Remove task ID number
                action_type = task_parts[1]
                
                # Extract item and location
                remaining_parts = task_parts[2:]
                if len(remaining_parts) >= 2:
                    item = remaining_parts[0]
                    location = remaining_parts[1]
                elif len(remaining_parts) == 1:
                    item = remaining_parts[0]
                    location = ""
                else:
                    item = ""
                    location = ""
                
                if parent_task not in tasks_by_parent:
                    tasks_by_parent[parent_task] = []
                
                tasks_by_parent[parent_task].append({
                    'id': task_id,
                    'action_type': action_type,
                    'item': item,
                    'location': location
                })
            
            # Add tasks organized by parent task
            for parent_task, subtasks in tasks_by_parent.items():
                auction_tasks_prompt += f"  -- For {parent_task}, you need to:\n"
                for subtask in subtasks:
                    action_desc = f"    * {subtask['action_type']} "
                    if subtask['item']:
                        action_desc += f"{subtask['item']} "
                    if subtask['location']:
                        action_desc += f"at {subtask['location']}"
                    auction_tasks_prompt += action_desc + "\n"
            
            # Add guidance on how to execute the tasks
            auction_tasks_prompt += "\nTo execute these tasks, you need to generate appropriate actions like:\n"
            auction_tasks_prompt += "  - goto_agent{}_location (to move to a specific location)\n".format(self.agent_id)
            auction_tasks_prompt += "  - get_agent{}_item_location (to pick up an item from a location)\n".format(self.agent_id)
            auction_tasks_prompt += "  - put_agent{}_location (to put an item you're holding into a location)\n".format(self.agent_id)
            auction_tasks_prompt += "  - activate_agent{}_location (to activate a tool at a location)\n".format(self.agent_id)
            auction_tasks_prompt += "\nFocus on completing one subtask at a time in a logical sequence.\n"
            auction_tasks_prompt += "Consider your current location and what you're holding when deciding what to do next.\n"
            
            # Add current state information
            auction_tasks_prompt += f"\nYour current state:\n"
            auction_tasks_prompt += f"  - Location: {self.current_location}\n"
            auction_tasks_prompt += f"  - Holding: {self.holding_item if self.holding_item else 'nothing'}\n"
            auction_tasks_prompt += f"  - Currency: {self.currency}\n"
        
        # Create the full prompt
        prompt = self.feedback + self.suggestions + convert_to_prompt(obs) + auction_tasks_prompt + '-action:\n'

        print(f"Agent {self.agent_id} prompt:\n{prompt}")

        # Cap message length
        if len(self.history) < self.look_ahead_steps + self.initial_history_length:
            self.history = prepend_history(self.history, prompt, verbose=verbose)
        else:
            self.history = (self.history[:self.initial_history_length] + 
                          self.history[-(self.look_ahead_steps-1):])
            self.history = prepend_history(self.history, prompt, verbose=verbose)

        # Generate action
        action = query_llm(self.history, model=self.model)
        
        try:
            parsed_actions = self.extract_actions(action)
            if parsed_actions:
                parsed_actions = [parsed_actions[0]]
        except Exception as e:
            print(f"Error parsing actions for agent {self.agent_id}: {e}")
            parsed_actions = []

        if parsed_actions:
            self.update_history(parsed_actions, role='assistant', verbose=verbose)
            # Add to action history
            self.action_history.append(parsed_actions[0])
            return [parsed_actions[0]]
        
        return parsed_actions

    def prioritize_tasks(self, agent_tasks, obs):
        """
        Prioritize the assigned tasks based on various factors.
        
        Args:
            agent_tasks: List of task IDs assigned to this agent
            obs: The current observation
            
        Returns:
            list: A prioritized list of task IDs
        """
        task_priorities = []
        
        for task_id in agent_tasks:
            task_parts = task_id.split('_')
            parent_task = task_parts[0].rstrip('0123456789')  # Remove task ID number
            action_type = task_parts[1]
            
            # Extract item and location
            remaining_parts = task_parts[2:]
            if len(remaining_parts) >= 2:
                item = remaining_parts[0]
                location = remaining_parts[1]
            elif len(remaining_parts) == 1:
                item = remaining_parts[0]
                location = ""
            else:
                item = ""
                location = ""
            
            # Calculate priority based on various factors
            priority = 0
            
            # Factor 1: Task urgency (based on lifetime)
            for task_name, lifetime in zip(obs.current_tasks_name, obs.current_tasks_lifetime):
                if parent_task == task_name:
                    # Higher priority for tasks with shorter lifetime
                    priority -= lifetime * 10
            
            # Factor 2: Current location and item
            if location == self.current_location:
                # Higher priority for tasks at current location
                priority += 50
            
            if action_type == "put" and item == self.holding_item:
                # Higher priority for putting an item we're already holding
                priority += 100
            
            if action_type == "get" and not self.holding_item:
                # Higher priority for getting items when we're not holding anything
                priority += 75
            
            # Factor 3: Task dependencies
            if action_type == "activate":
                # Activation usually comes after putting items
                priority -= 25
            
            if action_type == "put" and location == "servingtable":
                # Highest priority for completing a dish
                priority += 200
            
            task_priorities.append((task_id, priority))
        
        # Sort tasks by priority (highest first)
        sorted_tasks = [task_id for task_id, _ in sorted(task_priorities, key=lambda x: x[1], reverse=True)]
        
        return sorted_tasks

    def bid_subtask(self, subtask, agent_state, task_info):
        """Bid for a subtask"""
        # Create a prompt for the LLM to bid on the subtask
        subtask_id = subtask["id"]
        task_parts = subtask_id.split('_')
        task_name = task_parts[0].rstrip('0123456789')  # Remove task ID number
        action_type = task_parts[1]
        item = task_parts[2] if len(task_parts) > 2 else ""
        location = task_parts[3] if len(task_parts) > 3 else ""
        
        # Get information about other agents
        other_agents_info = ""

        for agent in self.env.agents:
            agent_name = agent.name
            # extract the number from the agent name
            agent_id = int(agent_name.split('_')[-1])
            if agent_id != self.agent_id:
                other_agents_info += f"Agent {agent_name} is at {agent.location} and {'holding ' + agent.holding.name if agent.holding else 'not holding any item'}.\n"
        
        # Get information about task dependencies
        task_dependencies = ""
        if action_type == "get":
            task_dependencies = f"To get {item}, you need to be at {location}.\n"
        elif action_type == "put":
            task_dependencies = f"To put {item}, you need to be holding {item} and be at {location}.\n"
        elif action_type == "activate":
            task_dependencies = f"To activate {location}, you need to be at {location} and not holding any item.\n"
        
        # Create a bidding prompt
        bid_prompt = f"""
You are agent{self.agent_id}, currently at location {self.current_location} and {'holding ' + self.holding_item if self.holding_item else 'not holding any item'}.
You have {self.currency} currency units available.

You need to bid on the following subtask:
- Subtask ID: {subtask_id}
- Action: {action_type} {item} at {location}
- Parent Task: {task_name}
- Reward if completed: {subtask["reward"]}
- Penalty if failed: {subtask["penalty"]}
- Task lifetime remaining: {task_info['lifetime']}

Environment information:
{other_agents_info}

Task dependencies:
{task_dependencies}

Your success rate on previous tasks is {self.success_rate:.2f}.
Your current location is {self.current_location}.

Please provide a bid amount based on:
1. Your current location relative to where the subtask needs to be performed (closer is better)
2. Whether you're already holding an item needed for this subtask
3. Your available currency (don't bid more than you have)
4. Your estimated ability to complete the task in time
5. The potential reward and penalty
6. The task's remaining lifetime (more urgent tasks may deserve higher bids)

Respond with only a number representing your bid amount (between 0 and your available currency).
"""
        
        # Query the LLM for a bid
        history = [("user", bid_prompt)]
        bid_response = query_llm(history, model=self.model, max_tokens=20)
        
        # Print the full prompt and response for debugging
        # print(f"Agent {self.agent_id} bid prompt:\n{bid_prompt}")
        print(f"Agent {self.agent_id} bid response: {bid_response}")
        
        # Extract the bid amount from the response
        try:
            # Use regex to find the first number in the response
            bid_match = re.search(r'\d+(\.\d+)?', bid_response)
            if bid_match:
                bid_amount = float(bid_match.group(0))
                # Ensure bid is within available currency
                bid_amount = min(bid_amount, self.currency)
                bid_amount = max(bid_amount, 0)
                return bid_amount
            else:
                # Default bid if no number found
                return 0
        except Exception as e:
            print(f"Error parsing bid from agent {self.agent_id}: {e}")
            return 0

class AuctionSystem:
    """
    Manages the auction process for task allocation among agents.
    """
    def __init__(self, env, model, agents, costs):
        self.env = env
        self.model = model
        self.agents = agents  # List of AuctionAgent instances
        self.costs = costs  # Cost matrix
        
        # Auction state
        self.task_subtasks = {}  # Maps tasks to their subtasks
        self.task_lifetime = {}  # Maps tasks to their lifetime
        self.task_assignments = {}  # Maps subtasks to assigned agents
        self.task_rewards = {}  # Maps tasks to their rewards
        self.task_penalties = {}  # Maps tasks to their penalties
        self.tasks_queue = []
        self.last_task_id = 0
        
        # Performance tracking
        self.completed_tasks = []
        self.failed_tasks = []
        
    def decompose_task(self, task_name, task_id, obs):
        """
        Decompose a task into subtasks.
        
        Args:
            task_name: The name of the task
            obs: The current observation
            
        Returns:
            list: A list of subtasks
        """
        # Get the task definition
        task_def = self.env.task_manager.task_def['task_list'][task_name]
        
        # Compute the dependency graph for the task
        all_required_items, tools, reward_mappings, mappings = compute_dependency(task_name, return_mapping=True)
        
        # Create subtasks based on the dependency graph
        subtasks = []
        
        # Add subtasks for getting base ingredients
        for item in all_required_items:
             # Add subtask for getting the item from storage
                subtask_id = f"{task_name}{task_id}_get_{item}_storage"
                subtasks.append(subtask_id)

        # Add subtasks for processing at each location
        for loc_type, items_list in mappings.items():
            for items in items_list:
                # Add subtasks for putting items at the location
                for item in items:
                    subtask_id = f"{task_name}{task_id}_put_{item}_{loc_type}"
                    subtasks.append(subtask_id)
                
                # Add a subtask for activating the location
                subtask_id = f"{task_name}{task_id}_activate_{loc_type}"
                subtasks.append(subtask_id)
                
                # Add a subtask for getting the processed item
                result_item = "_".join(sorted(items)) + f"_{loc_type}"
                if result_item in reward_mappings:
                    processed_item = reward_mappings[result_item]
                    subtask_id = f"{task_name}{task_id}_get_{processed_item}_{loc_type}"
                    subtasks.append(subtask_id)
        
        # Add final subtask for putting the completed dish on the serving table
        subtask_id = f"{task_name}{task_id}_put_{task_name}_servingtable"
        subtasks.append(subtask_id)
        
        subtasks2return = []
        for subtask in subtasks:
            subtasks2return.append(
                {
                    "id": subtask,
                    "reward": 10,
                    "penalty": -5,
                    "assigned_agent": None,
                    "completed": False,
                }
            )
        
        return subtasks2return
    
    def run_auction(self, obs, step, verbose=False):
        """
        Run the auction process for the current step.
        
        Args:
            obs: The current observation
            step: The current step number
            verbose: Whether to print verbose output
            
        Returns:
            list: A list of actions for all agents
        """
        # Update agent states
        for agent in self.agents:
            agent.update_state(obs)

        # Update task queue
        self.update_task_queue(obs)
        print("Current Tasks: ", obs.current_tasks_name)
        print("Current Tasks Lifetime: ", obs.current_tasks_lifetime)
        print("Tasks Queue: ", self.tasks_queue)
        
        # Update subtask completion status
        self.update_subtask_completion(obs)
        
        # Print task and subtask status for debugging
        self.print_task_status()

        # Auction subtasks
        agent_assignments = self.auction_subtasks(obs)
        print("Agent Assignments: ", agent_assignments)
        
        # Get actions from the agent assignments
        actions = []
        for agent in self.agents:
            # Get the assigned tasks for this agent
            assigned_tasks = agent_assignments.get(agent.agent_id, [])
            
            # Print agent status
            print(f"Agent {agent.agent_id} status:")
            print(f"  Location: {agent.current_location}")
            print(f"  Holding: {agent.holding_item}")
            print(f"  Currency: {agent.currency}")
            print(f"  Success rate: {agent.success_rate:.2f}")
            print(f"  Assigned tasks: {assigned_tasks}")
            
            # Get the action from the agent
            agent_action = agent.step(obs, step, verbose=verbose, agent_tasks=assigned_tasks)
            
            # Add the action to the list
            if agent_action:
                actions.extend(agent_action)
                print(f"Agent {agent.agent_id} action: {agent_action}")
            else:
                print(f"Agent {agent.agent_id} has no action")
        return
        return actions
        
    def print_task_status(self):
        """
        Print the current status of all tasks and subtasks.
        """
        print("\n=== TASK STATUS ===")
        for task in self.tasks_queue:
            print(f"Task: {task['task_name']} (ID: {task['task_id']}, Lifetime: {task['lifetime']})")
            
            # Count completed and total subtasks
            completed = sum(1 for subtask in task['subtasks'] if subtask['completed'])
            total = len(task['subtasks'])
            print(f"  Progress: {completed}/{total} subtasks completed")
            
            # Print details of uncompleted subtasks
            print("  Uncompleted subtasks:")
            for subtask in task['subtasks']:
                if not subtask['completed']:
                    print(f"    - {subtask['id']} (Assigned to: {subtask['assigned_agent']})")
        print("==================\n")

    def update_subtask_completion(self, obs):
        """
        Update the completion status of subtasks based on the observation.
        
        Args:
            obs: The current observation
        """
        # Check for completed subtasks based on agent actions and world state
        for task in self.tasks_queue:
            for subtask in task['subtasks']:
                if subtask['completed']:
                    continue
                    
                subtask_id = subtask['id']
                task_parts = subtask_id.split('_')
                action_type = task_parts[1]
                item = task_parts[2] if len(task_parts) > 2 else ""
                location = task_parts[3] if len(task_parts) > 3 else ""
                
                # Check if the subtask is completed based on the action type
                if action_type == "get":
                    # Check if any agent is holding the item
                    for agent_info in obs.agents:
                        if agent_info['hold'] == item:
                            subtask['completed'] = True
                            print(f"Subtask {subtask_id} completed: Agent {agent_info['id']} got {item}")
                            
                            # If this subtask was assigned to an agent, update their success rate
                            if subtask['assigned_agent'] is not None:
                                for agent in self.agents:
                                    if agent.agent_id == subtask['assigned_agent']:
                                        agent.action_success_history.append(1)
                                        if len(agent.action_success_history) > 10:
                                            agent.action_success_history.pop(0)
                                        agent.success_rate = sum(agent.action_success_history) / len(agent.action_success_history)
                                        break
                            break
                            
                elif action_type == "put":
                    # Check if the item is in the location
                    for loc_info in obs.locations:
                        if loc_info['id'] == location:
                            if item in loc_info['contains']:
                                subtask['completed'] = True
                                print(f"Subtask {subtask_id} completed: {item} put in {location}")
                                
                                # If this subtask was assigned to an agent, update their success rate
                                if subtask['assigned_agent'] is not None:
                                    for agent in self.agents:
                                        if agent.agent_id == subtask['assigned_agent']:
                                            agent.action_success_history.append(1)
                                            if len(agent.action_success_history) > 10:
                                                agent.action_success_history.pop(0)
                                            agent.success_rate = sum(agent.action_success_history) / len(agent.action_success_history)
                                            break
                            break
                            
                elif action_type == "activate":
                    # For activate subtasks, we need to check if the result of the activation is present
                    # This is more complex and depends on the specific task
                    # For now, we'll use a heuristic: check if any agent has recently been at the location
                    # and if the assigned agent is no longer holding any items they were before
                    
                    if subtask['assigned_agent'] is not None:
                        for agent in self.agents:
                            if agent.agent_id == subtask['assigned_agent']:
                                if agent.current_location == location and not agent.holding_item:
                                    # Mark as completed if agent is at the location and not holding anything
                                    subtask['completed'] = True
                                    print(f"Subtask {subtask_id} completed: {location} activated by agent {agent.agent_id}")
                                    
                                    # Update agent success rate
                                    agent.action_success_history.append(1)
                                    if len(agent.action_success_history) > 10:
                                        agent.action_success_history.pop(0)
                                    agent.success_rate = sum(agent.action_success_history) / len(agent.action_success_history)
                                    break

    def auction_subtasks(self, obs):
        """
        Run the auction for each of the uncompleted subtask in the task queue.
        
        Args:
            obs: The current observation
            
        Returns:
            dict: A dictionary mapping agents to their assigned subtasks
        """
        agent_assignments = {agent.agent_id: [] for agent in self.agents}
        
        # Collect all uncompleted subtasks from all tasks in the queue
        uncompleted_subtasks = []
        for task in self.tasks_queue:
            for subtask in task['subtasks']:
                if not subtask['completed']:
                    # Store the previously assigned agent before resetting
                    previously_assigned = subtask['assigned_agent']
                    # Reset the assigned agent for re-auctioning
                    subtask['assigned_agent'] = None
                    uncompleted_subtasks.append((subtask, task, previously_assigned))
        
        if not uncompleted_subtasks:
            return agent_assignments
            
        print(f"Running auction for {len(uncompleted_subtasks)} uncompleted subtasks")
        
        # For each uncompleted subtask, collect bids from all agents
        for subtask, task_info, previously_assigned in uncompleted_subtasks:
            print(f"Auctioning subtask: {subtask['id']}")
            
            # Collect bids from all agents
            bids = []
            for agent in self.agents:
                # Get agent state
                agent_state = {
                    'location': agent.current_location,
                    'holding': agent.holding_item,
                    'currency': agent.currency
                }
                
                # Get bid from agent
                bid_amount = agent.bid_subtask(subtask, agent_state, task_info)
                
                # Give a bonus to the previously assigned agent to provide continuity
                if previously_assigned is not None and agent.agent_id == previously_assigned:
                    # Add a 20% bonus to the bid of the previously assigned agent
                    bid_amount *= 1.2
                    print(f"Agent {agent.agent_id} gets a 20% bonus for previously working on this subtask")
                
                bids.append((agent.agent_id, bid_amount))
                print(f"Agent {agent.agent_id} bid {bid_amount} for subtask {subtask['id']}")
            
            # Find the highest bidder
            if bids:
                highest_bidder = max(bids, key=lambda x: x[1])
                agent_id, bid_amount = highest_bidder
                
                # If the highest bid is greater than 0, assign the subtask to the agent
                if bid_amount > 0:
                    # Assign the subtask to the agent
                    subtask['assigned_agent'] = agent_id
                    agent_assignments[agent_id].append(subtask['id'])
                    
                    # Deduct the bid amount from the agent's currency
                    # If this was the previously assigned agent, give them a discount
                    actual_bid = bid_amount
                    if previously_assigned is not None and agent_id == previously_assigned:
                        actual_bid = bid_amount / 1.2  # Remove the bonus for payment calculation
                    
                    for agent in self.agents:
                        if agent.agent_id == agent_id:
                            agent.currency -= actual_bid
                            # Only add to assigned_tasks if not already there
                            if subtask['id'] not in agent.assigned_tasks:
                                agent.assigned_tasks.append(subtask['id'])
                            print(f"Subtask {subtask['id']} assigned to agent {agent_id} for {actual_bid}")
                            break
        
        return agent_assignments

    def update_task_queue(self, obs):
        """
        Update the task queue based on the observation.
        """
        for task in self.tasks_queue:
            task['lifetime'] -= 1
            task['reward'] = 50 - task['lifetime']

        # remove task in the queue whose time is out
        # if just failed, remove the task in the queue
        if obs.just_failed:
            for task in self.tasks_queue:
                if task['lifetime'] <= 0:
                    self.tasks_queue.remove(task)
                    # Penalize the assigned agents
                    self.penalize_agents_for_failed_task(task)


        # if just succeeded, remove the task in the queue with same task name and largest lifetime
        if obs.task_just_success:
            # Group tasks by name and find the one with highest lifetime for each successful task
            for success_task in obs.task_just_success:
                matching_tasks = [t for t in self.tasks_queue if t['task_name'] == success_task]
                if matching_tasks:
                    task_to_remove = max(matching_tasks, key=lambda x: x['lifetime'])
                    self.tasks_queue.remove(task_to_remove)
                    # Reward the assigned agents
                    self.reward_agents_for_completed_task(task_to_remove)


        if len(self.tasks_queue) != len(obs.current_tasks_name):
            last_task = obs.current_tasks_name[-1]
            last_task_lifetime = obs.current_tasks_lifetime[-1]
            new_task = {
                'task_name': last_task,
                'task_id': self.last_task_id,
                'lifetime': last_task_lifetime,
                'subtasks': self.decompose_task(last_task, self.last_task_id, obs),
                'reward': 50 - last_task_lifetime, 
                'penalty': -20
            }
            self.tasks_queue.append(new_task)
            self.last_task_id += 1

    def reward_agents_for_completed_task(self, task):
        """
        Reward agents who contributed to completing a task.
        
        Args:
            task: The completed task
        """
        # Find all agents who were assigned subtasks for this task
        assigned_agents = set()
        for subtask in task['subtasks']:
            if subtask['assigned_agent'] is not None:
                assigned_agents.add(subtask['assigned_agent'])
        
        # Calculate reward per agent
        if assigned_agents:
            reward_per_agent = task['reward'] / len(assigned_agents)
            
            # Distribute rewards to agents
            for agent in self.agents:
                if agent.agent_id in assigned_agents:
                    agent.currency += reward_per_agent
                    agent.completed_tasks.append(task['task_name'])
                    
                    # Update success rate
                    agent.action_success_history.append(1)
                    if len(agent.action_success_history) > 10:
                        agent.action_success_history.pop(0)
                    agent.success_rate = sum(agent.action_success_history) / len(agent.action_success_history)
                    
                    print(f"Agent {agent.agent_id} rewarded {reward_per_agent} for completing task {task['task_name']}")
    
    def penalize_agents_for_failed_task(self, task):
        """
        Penalize agents who failed to complete a task.
        
        Args:
            task: The failed task
        """
        # Find all agents who were assigned subtasks for this task
        assigned_agents = set()
        for subtask in task['subtasks']:
            if subtask['assigned_agent'] is not None and not subtask['completed']:
                assigned_agents.add(subtask['assigned_agent'])
        
        # Calculate penalty per agent
        if assigned_agents:
            penalty_per_agent = task['penalty'] / len(assigned_agents)
            
            # Distribute penalties to agents
            for agent in self.agents:
                if agent.agent_id in assigned_agents:
                    agent.currency += penalty_per_agent  # Penalty is negative
                    agent.failed_tasks.append(task['task_name'])
                    
                    # Update success rate
                    agent.action_success_history.append(0)
                    if len(agent.action_success_history) > 10:
                        agent.action_success_history.pop(0)
                    agent.success_rate = sum(agent.action_success_history) / len(agent.action_success_history)
                    
                    print(f"Agent {agent.agent_id} penalized {penalty_per_agent} for failing task {task['task_name']}")

    

    

    
