import torch
import torch.nn as nn
from argparse import Namespace
from copy import deepcopy

import numpy as np

from xuance.common import get_configs, Union, Optional
from xuance.environment import make_envs, DummyVecMultiAgentEnv, SubprocVecMultiAgentEnv
from xuance.torch.agents import OffPolicyMARLAgents, BaseCallback
from xuance.torch.learners import LearnerMAS, REGISTRY_Learners
from xuance.torch.utils import NormalizeFunctions, ActivationFunctions
from xuance.torch.policies import REGISTRY_Policy
from xuance.torch import Module


# Step 1: Create a multi-agent policy.
class MyMARLPolicy(nn.Module):
    """
    An example of self-defined multi-agent policy for Independent DQN learning.

    Args:
        action_space: The action space of the environment.
        n_agents: The number of agents.
        representation: A neural network module responsible for extracting meaningful features.
        hidden_dim: Specifies the number of units in each hidden layer.
        device: The calculating device.
        use_parameter_sharing: Whether to share parameters across agents.
        model_keys: The keys for the models (agent names).
    """

    def __init__(self, action_space, n_agents, representation, hidden_dim, device,
                 use_parameter_sharing=True, model_keys=None, **kwargs):
        super(MyMARLPolicy, self).__init__()
        self.action_space = action_space
        self.n_agents = n_agents
        self.use_parameter_sharing = use_parameter_sharing
        self.device = device
        self.model_keys = model_keys or [f"agent_{i}" for i in range(n_agents)]

        # Build representations and Q-networks for each agent
        if use_parameter_sharing:
            # All agents share the same parameters
            self.representation = representation
            # Get feature dimension from the first model key
            self.feature_dim = self.representation[self.model_keys[0]].output_shapes['state'][0]
            
            # Single shared Q-network
            self.q_net = nn.Sequential(
                nn.Linear(self.feature_dim, hidden_dim),
                nn.ReLU(),
                nn.Linear(hidden_dim, action_space[self.model_keys[0]].n),
            ).to(device)
            self.target_q_net = deepcopy(self.q_net)
        else:
            # Each agent has its own parameters
            self.representations = representation  # representation is already a ModuleDict
            self.q_nets = nn.ModuleDict()
            self.target_q_nets = nn.ModuleDict()
            
            for key in self.model_keys:
                feature_dim = self.representations[key].output_shapes['state'][0]
                
                self.q_nets[key] = nn.Sequential(
                    nn.Linear(feature_dim, hidden_dim),
                    nn.ReLU(),
                    nn.Linear(hidden_dim, action_space[key].n),
                ).to(device)
                self.target_q_nets[key] = deepcopy(self.q_nets[key])

    def forward(self, observation, agent_ids=None, avail_actions=None, **kwargs):
        """
        Forward pass of the policy network.
        
        Args:
            observation: Dict of observations for each agent.
            agent_ids: Agent identifiers (optional).
            avail_actions: Available actions mask (optional).
            **kwargs: Additional keyword arguments.
            
        Returns:
            Tuple containing:
                - outputs: Feature representations for each agent
                - argmax_actions: Greedy actions for each agent  
                - q_values: Q-values for each agent
        """
        batch_size = list(observation.values())[0].shape[0]
        outputs = {}
        argmax_actions = {}
        q_values = {}
        
        if self.use_parameter_sharing:
            # Use shared parameters for all agents
            model_key = self.model_keys[0]  # Get the shared model key
            for agent_key in observation.keys():
                obs_out = self.representation[model_key](observation[agent_key])
                q_out = self.q_net(obs_out['state'])
                argmax_action = q_out.argmax(dim=-1)
                
                outputs[agent_key] = obs_out
                argmax_actions[agent_key] = argmax_action
                q_values[agent_key] = q_out
        else:
            # Use separate parameters for each agent
            for key in self.model_keys:
                obs_out = self.representations[key](observation[key])
                q_out = self.q_nets[key](obs_out['state'])
                argmax_action = q_out.argmax(dim=-1)
                
                outputs[key] = obs_out
                argmax_actions[key] = argmax_action
                q_values[key] = q_out
                
        return outputs, argmax_actions, q_values

    def target(self, observation, agent_ids=None, **kwargs):
        """
        Forward pass using target networks.
        
        Args:
            observation: Dict of observations for each agent.
            agent_ids: Agent identifiers (optional).
            **kwargs: Additional keyword arguments.
            
        Returns:
            Tuple containing:
                - outputs: Feature representations for each agent
                - argmax_actions: Target greedy actions for each agent
                - q_targets: Target Q-values for each agent
        """
        batch_size = list(observation.values())[0].shape[0]
        outputs = {}
        argmax_actions = {}
        q_targets = {}
        
        if self.use_parameter_sharing:
            model_key = self.model_keys[0]  # Get the shared model key
            for agent_key in observation.keys():
                obs_out = self.representation[model_key](observation[agent_key])
                q_target = self.target_q_net(obs_out['state'])
                argmax_action = q_target.argmax(dim=-1)
                
                outputs[agent_key] = obs_out
                argmax_actions[agent_key] = argmax_action.detach()
                q_targets[agent_key] = q_target.detach()
        else:
            for key in self.model_keys:
                obs_out = self.representations[key](observation[key])
                q_target = self.target_q_nets[key](obs_out['state'])
                argmax_action = q_target.argmax(dim=-1)
                
                outputs[key] = obs_out
                argmax_actions[key] = argmax_action.detach()
                q_targets[key] = q_target.detach()
                
        return outputs, argmax_actions, q_targets

    def copy_target(self):
        """Reset the parameters of target Q network as the Q network."""
        if self.use_parameter_sharing:
            for ep, tp in zip(self.q_net.parameters(), self.target_q_net.parameters()):
                tp.data.copy_(ep)
        else:
            for key in self.model_keys:
                for ep, tp in zip(self.q_nets[key].parameters(), self.target_q_nets[key].parameters()):
                    tp.data.copy_(ep)


# Step 2: Create the multi-agent learner.
class MyMARLLearner(LearnerMAS):
    """
    Custom multi-agent learner implementing Independent DQN learning.
    
    This learner extends the base LearnerMAS class to provide custom
    implementation for multi-agent Q-learning with independent agents.
    """
    
    def __init__(self, config, model_keys, agent_keys, policy, callback):
        super(MyMARLLearner, self).__init__(config, model_keys, agent_keys, policy, callback)
        # Build the optimizer.
        self.optimizer = torch.optim.Adam(self.policy.parameters(), self.config.learning_rate, eps=1e-5)
        self.loss = nn.MSELoss()  # Build a loss function
        self.sync_frequency = config.sync_frequency  # The period to synchronize the target network

    def update(self, sample):
        """
        Update the policy networks using a batch of training samples.
        
        Args:
            sample: Dictionary containing training batch data with keys:
                - obs: Current observations for all agents
                - actions: Actions taken by all agents  
                - obs_next: Next observations for all agents
                - rewards: Rewards received by all agents
                - terminals: Terminal flags for all agents
                
        Returns:
            Dict containing training information and losses.
        """
        info = {}
        self.iterations += 1
        
        # Get a batch of training samples for all agents
        # Use the actual keys from the sample data
        actual_agent_keys = list(sample['obs'].keys())
        obs_batch = {key: torch.as_tensor(sample['obs'][key], device=self.device) for key in actual_agent_keys}
        act_batch = {key: torch.as_tensor(sample['actions'][key], device=self.device) for key in actual_agent_keys}
        next_batch = {key: torch.as_tensor(sample['obs_next'][key], device=self.device) for key in actual_agent_keys}
        rew_batch = {key: torch.as_tensor(sample['rewards'][key], device=self.device) for key in actual_agent_keys}
        ter_batch = {key: torch.as_tensor(sample['terminals'][key], dtype=torch.float, device=self.device) for key in actual_agent_keys}

        # Forward passes for all agents
        _, _, q_eval = self.policy(obs_batch)
        _, _, q_next = self.policy.target(next_batch)
        
        # Compute losses for all agents
        total_loss = 0
        agent_losses = {}
        
        for key in actual_agent_keys:
            # Now each agent has its own Q values in the output
            q_next_action = q_next[key].max(dim=-1).values
            q_eval_action = q_eval[key].gather(-1, act_batch[key].long().unsqueeze(-1)).reshape(-1)
            target_value = rew_batch[key] + (1 - ter_batch[key]) * self.gamma * q_next_action
            
            # Compute loss for this agent
            agent_loss = self.loss(q_eval_action, target_value.detach())
            agent_losses[key] = agent_loss.item()
            total_loss += agent_loss

        # Backward and optimizing steps
        self.optimizer.zero_grad()
        total_loss.backward()
        self.optimizer.step()

        # Synchronize the target network
        if self.iterations % self.sync_frequency == 0:
            self.policy.copy_target()

        # Set the variables you need to observe
        info.update({
            'total_loss': total_loss.item(),
            'iterations': self.iterations,
        })
        
        # Add individual agent losses
        for key in actual_agent_keys:
            info[f'loss_{key}'] = agent_losses[key]
            info[f'q_eval_{key}'] = q_eval[key].mean().item()

        return info


# Step 3: Create the multi-agent agent.
class MyMARLAgents(OffPolicyMARLAgents):
    """Multi-agent version of the custom DQN implementation."""
    
    def __init__(self, config: Namespace, 
                 envs: Union[DummyVecMultiAgentEnv, SubprocVecMultiAgentEnv], 
                 callback: Optional[BaseCallback] = None):
        super(MyMARLAgents, self).__init__(config, envs, callback)
        
        # Initialize epsilon-greedy parameters
        self.start_greedy, self.end_greedy = config.start_greedy, config.end_greedy
        self.delta_egreedy = (self.start_greedy - self.end_greedy) / config.decay_step_greedy
        self.e_greedy = self.start_greedy

        self.policy = self._build_policy()  # Build the policy module
        self.memory = self._build_memory()  # Build the replay buffer
        REGISTRY_Learners['MyMARLLearner'] = MyMARLLearner  # Registry your pre-defined learner
        self.learner = self._build_learner(self.config, self.model_keys, self.agent_keys, self.policy, self.callback)  # Build the learner

    def _build_policy(self) -> Module:
        """
        Build multi-agent policy.
        
        Constructs the custom multi-agent policy with appropriate representation
        networks and Q-networks based on configuration settings.
        
        Returns:
            Module: The constructed multi-agent policy.
        """
        normalize_fn = NormalizeFunctions[self.config.normalize] if hasattr(self.config, "normalize") else None
        initializer = torch.nn.init.orthogonal_
        activation = ActivationFunctions[self.config.activation]
        device = self.device

        # Build representation
        representation = self._build_representation(self.config.representation, self.observation_space, self.config)
        
        # Build custom multi-agent policy
        policy = MyMARLPolicy(
            action_space=self.action_space,
            n_agents=self.n_agents,
            representation=representation,
            hidden_dim=64,  # You can make this configurable
            device=device,
            use_parameter_sharing=self.use_parameter_sharing,
            model_keys=self.model_keys
        )
        return policy


if __name__ == '__main__':
    config = get_configs(file_dir="new_marl.yaml")  # Get the config settings from .yaml file
    config = Namespace(**config)  # Convert the config from dict to argparse
    envs = make_envs(config)  # Make vectorized multi-agent environments
    agents = MyMARLAgents(config, envs)  # Instantiate your pre-build multi-agent class

    if not config.test_mode:  # Training mode.
        train_steps = config.running_steps // config.parallels
        agents.train(train_steps)  # Train your agents
        agents.save_model("final_train_model.pth")  # After training, save the model
    else:  # Testing mode.
        def env_fn():
            config.parallels = config.test_episode
            return make_envs(config)
        
        agents.load_model(path=agents.model_dir_load)  # Load pre-trained model
        scores = agents.test(env_fn, config.test_episode)  # Test your agents
        print(f"Mean Score: {np.array(scores).mean()}, Std: {np.array(scores).std()}")

    agents.finish()  # Finish the agents
    envs.close()  # Close the environments
