#from modules.agents import REGISTRY as agent_REGISTRY
from modules.agents.FeUdal_agent import Feudal_ManagerAgent, Feudal_WorkerAgent
from components.action_selectors import REGISTRY as action_REGISTRY
import torch


# This multi-agent controller shares parameters between agents
class FeUdalMAC:
    def __init__(self, scheme, groups, args):
        self.n_agents = args.n_agents
        self.args = args
        self.input_shape = self._get_input_shape(scheme)
        self.manager_input_shape = self._get_manager_input_shape(scheme)
        self._build_agents(self.manager_input_shape, self.input_shape)
        self.agent_output_type = args.agent_output_type

        self.action_selector = action_REGISTRY[args.action_selector](args)

        self.manager_hidden_states = None
        self.manager_cell_states = None
        self.worker_hidden_states = None
        self.worker_cell_states = None
        
        self.single_past_goals = []
        self.batch_past_goals = []

    def select_actions(self, ep_batch, t_ep, t_env, bs=slice(None), test_mode=False):
        # Only select actions for the selected batch elements in bs
        avail_actions = ep_batch["avail_actions"][:, t_ep]
        agent_outputs, goal_outs, value_outs = self.forward(ep_batch, t_ep, test_mode=test_mode)
        chosen_actions = self.action_selector.select_action(agent_outputs[bs], avail_actions[bs], t_env, test_mode=test_mode)
        return chosen_actions

    def forward(self, ep_batch, t, test_mode=False):
        manager_inputs = self._build_manager_inputs(ep_batch, t)
        worker_inputs = self._build_inputs(ep_batch, t)
        
        # 更新為處理 LSTM 的 hidden states 和 cell states
        manager_goal, manager_value, (new_manager_hidden, new_manager_cell) = self.manager_agent(
            manager_inputs, 
            (self.manager_hidden_states, self.manager_cell_states)
        )
        self.manager_hidden_states = new_manager_hidden
        self.manager_cell_states = new_manager_cell

        worker_agent_outs, (new_worker_hidden, new_worker_cell), self.single_past_goals, self.batch_past_goals = self.worker_agent(
            worker_inputs, 
            (self.worker_hidden_states, self.worker_cell_states),
            self.single_past_goals,
            self.batch_past_goals,
            manager_goal
        )
        self.worker_hidden_states = new_worker_hidden
        self.worker_cell_states = new_worker_cell

        return worker_agent_outs.view(ep_batch.batch_size, self.n_agents, -1), manager_goal.view(ep_batch.batch_size, self.n_agents, -1), manager_value.view(ep_batch.batch_size, self.n_agents, -1)

    def init_hidden(self, batch_size):
        # 初始化 LSTM 的 hidden states 和 cell states
        manager_hidden, manager_cell = self.manager_agent.init_hidden()
        worker_hidden, worker_cell = self.worker_agent.init_hidden()
        
        self.manager_hidden_states = manager_hidden.unsqueeze(0).expand(batch_size, self.n_agents, -1)
        self.manager_cell_states = manager_cell.unsqueeze(0).expand(batch_size, self.n_agents, -1)
        self.worker_hidden_states = worker_hidden.unsqueeze(0).expand(batch_size, self.n_agents, -1)
        self.worker_cell_states = worker_cell.unsqueeze(0).expand(batch_size, self.n_agents, -1)

    def parameters(self):
        """返回所有可訓練參數"""
        # 將生成器轉換為列表後再合併
        manager_params = list(self.manager_agent.parameters())
        worker_params = list(self.worker_agent.parameters())
        return manager_params + worker_params

    def manager_parameters(self):
        return self.manager_agent.parameters()

    def worker_parameters(self):
        return self.worker_agent.parameters()

    def load_state(self, other_mac):
        self.manager_agent.load_state_dict(other_mac.manager_agent.state_dict())
        self.worker_agent.load_state_dict(other_mac.worker_agent.state_dict())

    def cuda(self):
        self.manager_agent.cuda()
        self.worker_agent.cuda()

    def cpu(self):
        self.manager_agent.cpu()
        self.worker_agent.cpu()

    def get_device(self):
        """返回模型所在的設備（CPU或GPU）"""
        # 直接檢查第一個參數的設備
        if hasattr(self, 'manager_agent'):
            return next(self.manager_agent.parameters()).device
        elif hasattr(self, 'worker_agent'):
            return next(self.worker_agent.parameters()).device
        return torch.device("cpu")  # 默認返回 CPU

    def save_models(self, path):
        torch.save(self.manager_agent.state_dict(), "{}/manager_agent.torch".format(path))
        torch.save(self.worker_agent.state_dict(), "{}/worker_agent.torch".format(path))
        

    def load_models(self, path):
        self.manager_agent.load_state_dict(torch.load("{}/manager_agent.torch".format(path), map_location=lambda storage, loc: storage))
        self.worker_agent.load_state_dict(torch.load("{}/worker_agent.torch".format(path), map_location=lambda storage, loc: storage))

    def _build_agents(self, manager_input_shape, worker_input_shape):
        self.manager_agent = Feudal_ManagerAgent(manager_input_shape, self.args)
        self.worker_agent = Feudal_WorkerAgent(worker_input_shape, self.args)

    def _build_inputs(self, batch, t):
        # Assumes homogenous agents with flat observations.
        # Other MACs might want to e.g. delegate building inputs to each agent
        bs = batch.batch_size
        inputs = []
        inputs.append(batch["obs"][:, t])  # b1av
        if self.args.obs_last_action:
            if t == 0:
                inputs.append(torch.zeros_like(batch["actions_onehot"][:, t]))
            else:
                inputs.append(batch["actions_onehot"][:, t-1])
        if self.args.obs_agent_id:
            inputs.append(torch.eye(self.n_agents, device=batch.device).unsqueeze(0).expand(bs, -1, -1))

        inputs = torch.cat([x.reshape(bs*self.n_agents, -1) for x in inputs], dim=1)
        return inputs

    def _get_input_shape(self, scheme):
        input_shape = scheme["obs"]["vshape"]
        if self.args.obs_last_action:
            input_shape += scheme["actions_onehot"]["vshape"][0]
        if self.args.obs_agent_id:
            input_shape += self.n_agents

        return input_shape

    def _build_manager_inputs(self, batch, t):
        # 假設 manager 需要不同的輸入結構
        bs = batch.batch_size
        manager_inputs = []
        manager_inputs.append(batch["obs"][:, t])  # 假設有一個 manager 特定的觀察
        # if self.args.obs_last_action:
        #     if t == 0:
        #         manager_inputs.append(torch.zeros_like(batch["actions_onehot"][:, t]))
        #     else:
        #         manager_inputs.append(batch["actions_onehot"][:, t-1])
        # # 不需要代理ID，移除以下程式碼
        # if self.args.obs_agent_id:
        #     manager_inputs.append(torch.eye(self.n_agents, device=batch.device).unsqueeze(0).expand(bs, -1, -1))

        manager_inputs = torch.cat([x.reshape(bs*self.n_agents, -1) for x in manager_inputs], dim=1)
        return manager_inputs
    
    def _get_manager_input_shape(self, scheme):
        input_shape = scheme["obs"]["vshape"]
        # if self.args.manager_obs_last_action:
        #     input_shape += scheme["manager_actions_onehot"]["vshape"][0]
        # if self.args.manager_obs_agent_id:
        #     input_shape += self.n_agents
        return input_shape
