import torch.nn as nn
import torch.nn.functional as F
import torch

class Feudal_ManagerAgent(nn.Module):
    def __init__(self, input_shape, args):
        super(Feudal_ManagerAgent, self).__init__()
        self.args = args

        # Manager network
        self.manager_fc1 = nn.Linear(input_shape, args.manager_hidden_dim)
        #self.manager_rnn = nn.GRUCell(args.manager_hidden_dim, args.manager_hidden_dim)
        self.manager_rnn = nn.LSTMCell(args.manager_hidden_dim, args.manager_hidden_dim)
        self.manager_fc2 = nn.Linear(args.manager_hidden_dim, args.state_dim)

        # 目標生成
        self.goal_network = nn.Linear(args.manager_hidden_dim, args.state_dim)
        
        # 狀態值估計 V_t^M
        self.value_network = nn.Linear(args.manager_hidden_dim, 1)

    def init_hidden(self):
        # Initialize hidden states for both manager and worker
        manager_hidden = self.manager_fc1.weight.new(1, self.args.manager_hidden_dim).zero_()
        manager_cell = self.manager_fc1.weight.new(1, self.args.manager_hidden_dim).zero_()
        return (manager_hidden, manager_cell)

    def forward(self, inputs, hidden):
        x = F.relu(self.manager_fc1(inputs))
        h_in, c_in = hidden
        h_in = h_in.reshape(-1, self.args.manager_hidden_dim)
        c_in = c_in.reshape(-1, self.args.manager_hidden_dim)
        h, c = self.manager_rnn(x, (h_in, c_in))
        
        # 生成目標
        goal = self.goal_network(h)
        
        # 估計狀態值
        value = self.value_network(h)
        
        return goal, value, (h, c)
    
class Feudal_WorkerAgent(nn.Module):
    def __init__(self, input_shape, args):
        super(Feudal_WorkerAgent, self).__init__()
        self.args = args
        
        # Worker 網絡
        self.worker_fc1 = nn.Linear(input_shape, args.worker_hidden_dim)
        self.worker_rnn = nn.LSTMCell(args.worker_hidden_dim, args.worker_hidden_dim)
        
        # U_t: Action embedding matrix (1x16)
        self.U_embedding = nn.Linear(args.worker_hidden_dim, args.embedding_dim)
        
        # w_t: 優勢方向/權重 (1x16)
        self.w_network = nn.Linear(args.state_dim, args.embedding_dim)
        
        # 最終的 Q 值輸出
        self.q_network = nn.Linear(args.embedding_dim, args.n_actions)

    def init_hidden(self):
        # Initialize hidden states for both manager and worker
        worker_hidden = self.worker_fc1.weight.new(1, self.args.worker_hidden_dim).zero_()
        worker_cell = self.worker_fc1.weight.new(1, self.args.worker_hidden_dim).zero_()
        return (worker_hidden, worker_cell)
    
    def forward(self, inputs, worker_hidden, single_past_goals, batch_past_goals, goal):
        # Worker RNN
        x = F.relu(self.worker_fc1(inputs))
        h_in, c_in = worker_hidden
        h_in = h_in.reshape(-1, self.args.worker_hidden_dim)
        c_in = c_in.reshape(-1, self.args.worker_hidden_dim)
        h, c = self.worker_rnn(x, (h_in, c_in))
        
        # 生成 U_t (action embedding matrix)
        U_t = self.U_embedding(h)
        
        # 生成 w_t (優勢方向)
        w_t = self.w_network(goal)
        
        # 計算 s_t (weighted state)
        s_t = U_t * w_t  # element-wise multiplication
        
        # 生成 Q 值
        q = self.q_network(s_t)
        
        return q, (h, c), single_past_goals, batch_past_goals
        
        
        
    

