"""
Multi-Agent Reinforcement Learning Module

This module implements multi-agent reinforcement learning algorithms and utilities.
"""

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim

class MADDPGAgent:
    def __init__(self, state_dim, action_dim, hidden_dim=64):
        self.actor = nn.Sequential(
            nn.Linear(state_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, action_dim),
            nn.Tanh()
        )
        self.critic = nn.Sequential(
            nn.Linear(state_dim + action_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, 1)
        )
        self.actor_optimizer = optim.Adam(self.actor.parameters(), lr=1e-3)
        self.critic_optimizer = optim.Adam(self.critic.parameters(), lr=1e-3)

    def act(self, state):
        state = torch.FloatTensor(state)
        action = self.actor(state)
        return action.detach().numpy()

    def update(self, states, actions, rewards, next_states, dones):
        states = torch.FloatTensor(states)
        actions = torch.FloatTensor(actions)
        rewards = torch.FloatTensor(rewards).unsqueeze(1)
        next_states = torch.FloatTensor(next_states)
        dones = torch.FloatTensor(dones).unsqueeze(1)

        # Update critic
        Q_values = self.critic(torch.cat([states, actions], dim=1))
        next_actions = self.actor(next_states)
        next_Q_values = self.critic(torch.cat([next_states, next_actions], dim=1))
        target_Q_values = rewards + (1 - dones) * 0.99 * next_Q_values

        critic_loss = nn.MSELoss()(Q_values, target_Q_values.detach())
        self.critic_optimizer.zero_grad()
        critic_loss.backward()
        self.critic_optimizer.step()

        # Update actor
        actor_loss = -self.critic(torch.cat([states, self.actor(states)], dim=1)).mean()
        self.actor_optimizer.zero_grad()
        actor_loss.backward()
        self.actor_optimizer.step()

class MultiAgentEnvironment:
    def __init__(self, num_agents, state_dim, action_dim):
        self.num_agents = num_agents
        self.state_dim = state_dim
        self.action_dim = action_dim

    def reset(self):
        return [np.random.rand(self.state_dim) for _ in range(self.num_agents)]

    def step(self, actions):
        # Simplified environment dynamics
        next_states = [np.random.rand(self.state_dim) for _ in range(self.num_agents)]
        rewards = [np.random.rand() for _ in range(self.num_agents)]
        dones = [False for _ in range(self.num_agents)]
        return next_states, rewards, dones, {}

def train_maddpg(num_agents, state_dim, action_dim, num_episodes=1000):
    env = MultiAgentEnvironment(num_agents, state_dim, action_dim)
    agents = [MADDPGAgent(state_dim, action_dim) for _ in range(num_agents)]

    for episode in range(num_episodes):
        states = env.reset()
        episode_reward = 0

        while True:
            actions = [agent.act(state) for agent, state in zip(agents, states)]
            next_states, rewards, dones, _ = env.step(actions)
            episode_reward += sum(rewards)

            for i, agent in enumerate(agents):
                agent.update([states[i]], [actions[i]], [rewards[i]], [next_states[i]], [dones[i]])

            states = next_states

            if any(dones):
                break

        if episode % 100 == 0:
            print(f"Episode {episode}, Average Reward: {episode_reward / num_agents:.2f}")

    return agents

def inter_agent_communication(agents, messages):
    """
    Simulates inter-agent communication by exchanging messages between agents.
    """
    for i, agent in enumerate(agents):
        received_messages = [msg for j, msg in enumerate(messages) if j != i]
        # Process received messages (simplified)
        agent.process_messages(received_messages)

def main():
    num_agents = 3
    state_dim = 5
    action_dim = 2

    trained_agents = train_maddpg(num_agents, state_dim, action_dim)

    # Simulate inter-agent communication
    messages = [f"Message from Agent {i}" for i in range(num_agents)]
    inter_agent_communication(trained_agents, messages)

if __name__ == "__main__":
    main()
