"from . import RLAgent
from common.registry import Registry
import numpy as np
import os
import random
from collections import OrderedDict, deque
import gym

from generator.lane_vehicle import LaneVehicleGenerator
from generator.intersection_phase import IntersectionPhaseGenerator
import torch
from torch import nn
import torch.nn.functional as F
import torch_scatter
import torch.optim as optim
from torch.nn.utils import clip_grad_norm_

from torch_geometric.nn import MessagePassing
from torch_geometric.data import Data, Batch
from torch_geometric.utils import add_self_loops

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using device: {device}")
torch.autograd.set_detect_anomaly(True)

@Registry.register_model('TestLight')
class TestLightAgent(RLAgent):
    def __init__(self, world, rank):
        super().__init__(world, world.intersection_ids[rank])
        self.rank = rank  # Initialize the rank attribute
        self.buffer_size = Registry.mapping['trainer_mapping']['setting'].param['buffer_size']
        self.replay_buffer = deque(maxlen=self.buffer_size)

        self.graph = Registry.mapping['world_mapping']['graph_setting'].graph
        self.world = world
        self.sub_agents = len(self.world.intersections)
        self.edge_idx = torch.tensor(self.graph['sparse_adj'].T, dtype=torch.long).to(device)
        
        self.phase = Registry.mapping['model_mapping']['setting'].param['phase']
        self.one_hot = Registry.mapping['model_mapping']['setting'].param['one_hot']
        self.model_dict = Registry.mapping['model_mapping']['setting'].param

        # Initialize intersection importance with ones for all nodes
        self.intersection_importance = {i: torch.tensor(1.0, device=device) for i in range(self.sub_agents)}

        observation_generators = []
        for inter in self.world.intersections:
            node_id = inter.id if 'GS_' not in inter.id else inter.id[3:]
            node_idx = self.graph['node_id2idx'][node_id]
            tmp_generator = LaneVehicleGenerator(self.world, inter, ['lane_count'], in_only=True, average=None)
            observation_generators.append((node_idx, tmp_generator))
        observation_generators.sort(key=lambda x: x[0])  # Corrected sorting method
        self.ob_generator = observation_generators

        rewarding_generators = []
        for inter in self.world.intersections:
            node_id = inter.id if 'GS_' not in inter.id else inter.id[3:]
            node_idx = self.graph['node_id2idx'][node_id]
            tmp_generator = LaneVehicleGenerator(self.world, inter, ["lane_waiting_count"],
                                                 in_only=True, average='all', negative=True)
            rewarding_generators.append((node_idx, tmp_generator))
        rewarding_generators.sort(key=lambda x: x[0])
        self.reward_generator = rewarding_generators

        queues = []
        for inter in self.world.intersections:
            node_id = inter.id if 'GS_' not in inter.id else inter.id[3:]
            node_idx = self.graph['node_id2idx'][node_id]
            tmp_generator = LaneVehicleGenerator(self.world, inter, ["lane_waiting_count"], 
                                                 in_only=True, negative=False)
            queues.append((node_idx, tmp_generator))
        queues.sort(key=lambda x: x[0])
        self.queue = queues

        delays = []
        for inter in self.world.intersections:
            node_id = inter.id if 'GS_' not in inter.id else inter.id[3:]
            node_idx = self.graph['node_id2idx'][node_id]
            tmp_generator = LaneVehicleGenerator(self.world, inter, ["lane_delay"], 
                                                 in_only=True, average="all", negative=False)
            delays.append((node_idx, tmp_generator))
        delays.sort(key=lambda x: x[0])
        self.delay = delays

        phasing_generators = []
        for inter in self.world.intersections:
            node_id = inter.id if 'GS_' not in inter.id else inter.id[3:]
            node_idx = self.graph['node_id2idx'][node_id]
            tmp_generator = IntersectionPhaseGenerator(self.world, inter, ['phase'],
                                                       targets=['cur_phase'], negative=False)
            phasing_generators.append((node_idx, tmp_generator))
        phasing_generators.sort(key=lambda x: x[0])
        self.phase_generator = phasing_generators

        self.action_space = gym.spaces.Discrete(len(self.world.intersections[0].phases))
        min_ob_length = max([ob[1].ob_length for ob in self.ob_generator])
        self.ob_length = min_ob_length


        self.learning_rate = Registry.mapping['model_mapping']['setting'].param['learning_rate']

        self.target_update_freq = Registry.mapping['trainer_mapping']['setting'].param['target_update_freq']

        self.gamma = Registry.mapping['model_mapping']['setting'].param['gamma']
        self.grad_clip = Registry.mapping['model_mapping']['setting'].param['grad_clip']
        self.epsilon = Registry.mapping['model_mapping']['setting'].param['epsilon']
        self.epsilon_decay = Registry.mapping['model_mapping']['setting'].param['epsilon_decay']
        self.epsilon_min = Registry.mapping['model_mapping']['setting'].param['epsilon_min']

        self.vehicle_max = Registry.mapping['model_mapping']['setting'].param['vehicle_max']
        self.batch_size = Registry.mapping['model_mapping']['setting'].param['batch_size']

        self.episode_count = 0
        self.update_interval = Registry.mapping['trainer_mapping']['setting'].param['update_interval']  # Initial interval for updating importance
        self.decay_factor = Registry.mapping['trainer_mapping']['setting'].param['decay_factor']  # Factor to increase the interval after each update
        self.softmax_temperature = Registry.mapping['trainer_mapping']['setting'].param['softmax_temperature']  # Temperature for softmax
        self.softmax_feather = Registry.mapping['trainer_mapping']['setting'].param['softmax_feather']  # Feather for softmax
    
        self.model = self._build_model().to(device)
        self.target_model = self._build_model().to(device)
        self.optimizer = optim.AdamW(self.model.parameters(), lr=self.learning_rate)
        self.criterion = nn.MSELoss()

    def reset(self):
        # Similar to __init__ but for resetting generators
        observation_generators = []
        for inter in self.world.intersections:
            node_id = inter.id if 'GS_' not in inter.id else inter.id[3:]
            node_idx = self.graph['node_id2idx'][node_id]
            tmp_generator = LaneVehicleGenerator(self.world, inter, ['lane_count'], in_only=True, average=None)
            observation_generators.append((node_idx, tmp_generator))
        observation_generators.sort(key=lambda x: x[0])
        self.ob_generator = observation_generators

        rewarding_generators = []
        for inter in self.world.intersections:
            node_id = inter.id if 'GS' not in inter.id else inter.id[3:]
            node_idx = self.graph['node_id2idx'][node_id]
            tmp_generator = LaneVehicleGenerator(self.world, inter, ["lane_waiting_count"],
                                                 in_only=True, average='all', negative=True)
            rewarding_generators.append((node_idx, tmp_generator))
        rewarding_generators.sort(key=lambda x: x[0])
        self.reward_generator = rewarding_generators

        phasing_generators = []
        for inter in self.world.intersections:
            node_id = inter.id if 'GS_' not in inter.id else inter.id[3:]
            node_idx = self.graph['node_id2idx'][node_id]
            tmp_generator = IntersectionPhaseGenerator(self.world, inter, ['phase'],
                                                       targets=['cur_phase'], negative=False)
            phasing_generators.append((node_idx, tmp_generator))
        phasing_generators.sort(key=lambda x: x[0])
        self.phase_generator = phasing_generators

        queues = []
        for inter in self.world.intersections:
            node_id = inter.id if 'GS_' not in inter.id else inter.id[3:]
            node_idx = self.graph['node_id2idx'][node_id]
            tmp_generator = LaneVehicleGenerator(self.world, inter, ["lane_waiting_count"], 
                                                 in_only=True, negative=False)
            queues.append((node_idx, tmp_generator))
        queues.sort(key=lambda x: x[0])
        self.queue = queues

        delays = []
        for inter in self.world.intersections:
            node_id = inter.id if 'GS_' not in inter.id else inter.id[3:]
            node_idx = self.graph['node_id2idx'][node_id]
            tmp_generator = LaneVehicleGenerator(self.world, inter, ["lane_delay"], 
                                                 in_only=True, average="all", negative=False)
            delays.append((node_idx, tmp_generator))
        delays.sort(key=lambda x: x[0])
        self.delay = delays

    def get_ob(self):
        x_obs = []  # sub_agents * lane_nums,
        for i in range(len(self.ob_generator)):
            ob = self.ob_generator[i][1].generate() / self.vehicle_max
            ob = np.pad(ob, (0, self.ob_length - ob.shape[-1]))
            x_obs.append(ob)
        x_obs = np.array(x_obs, dtype=np.float32)
        return x_obs

    def get_reward(self):
        rewards = []
        for i in range(len(self.reward_generator)):
            rewards.append(self.reward_generator[i][1].generate())
        rewards = np.squeeze(np.array(rewards, dtype=np.float32)) * 12
        return rewards

    def get_phase(self):
        phase = []
        for i in range(len(self.phase_generator)):
            phase.append((self.phase_generator[i][1].generate()))
        phase = (np.concatenate(phase)).astype(np.int8)
        return phase

    def get_queue(self):
        queue = []
        for item in self.queue:
            item = item[1].generate()
            item = np.pad(item, (0, self.ob_length - item.shape[-1]))
            queue.append(item)
        tmp_queue = np.squeeze(np.array(queue, dtype=np.float32))
        queue = np.sum(tmp_queue, axis=1 if len(tmp_queue.shape) == 2 else 0)
        return queue

    def get_delay(self):
        delay = []
        for i in range(len(self.delay)):
            delay.append((self.delay[i][1].generate()))
        delay = np.squeeze(np.array(delay, dtype=np.float32))
        return delay

    def get_action(self, ob, phase, test=False):
        if not test and np.random.rand() <= self.epsilon:
            return self.sample()
        
        observation = torch.tensor(ob, dtype=torch.float32).to(device)
        edge_index = self.edge_idx.to(device)
        
        # Generate masks if necessary
        tgt_mask = self.generate_square_subsequent_mask(observation.size(0)).to(device)
        memory_mask = None  # Define as needed

        q_values, attention_matrices = self.model(x=observation, edge_index=edge_index, train=False, tgt_mask=tgt_mask, memory_mask=memory_mask)
        actions = torch.argmax(q_values, dim=1)
        
        self.update_intersection_importance(attention_matrices)
    
        return actions.cpu().numpy()

    def sample(self):
        return np.random.randint(0, self.action_space.n, self.sub_agents)

    def generate_square_subsequent_mask(self, sz):
        mask = torch.triu(torch.ones(sz, sz) * float('-inf'), diagonal=1)
        return mask

    def update_intersection_importance(self, attention_matrices):
        importance = {}
        last_attention_matrix = attention_matrices[-1]  # Use only the last attention matrix

        # Calculate raw importance scores
        for i in range(self.sub_agents):
            if i < last_attention_matrix.size(0):
                importance[i] = last_attention_matrix[i].mean().detach()
            else:
                # If we don't have attention for this node, use the previous importance or a default value
                importance[i] = self.intersection_importance.get(i, torch.tensor(0.5, device=device))

        # Convert importance dict to tensor
        importance_tensor = torch.tensor([importance[i] for i in range(self.sub_agents)], device=device)

        # Apply softmax with a lower temperature to make the distribution sharper
        temperature = self.softmax_temperature  # Lower value for sharper distribution
        normalized_importance_tensor = F.softmax(importance_tensor / temperature, dim=0)

        # Convert back to dictionary
        normalized_importance = {i: normalized_importance_tensor[i].item() for i in range(self.sub_agents)}

        self.intersection_importance = normalized_importance
        self.model.intersection_importance = normalized_importance
        self.target_model.intersection_importance = normalized_importance
        #PRINT DEBUG STATEMENT
        #print(f"Updated Normalized Intersection Importance: {self.intersection_importance}")
        #print(f"Number of nodes: {self.sub_agents}")
        #print(f"Size of importance dict: {len(self.intersection_importance)}")
        #print(f"Sum of importance scores: {sum(self.intersection_importance.values())}")

    def _build_model(self):
        model = TestLightNet(
            self.ob_length, 
            self.action_space.n, 
            self.intersection_importance, 
            softmax_temperature=self.softmax_temperature,
            softmax_feather=self.softmax_feather,
            update_interval=self.update_interval,
            decay_factor=self.decay_factor,
            **self.model_dict
        )
        return model

    def remember(self, last_obs, last_phase, actions, actions_prob, rewards, obs, cur_phase, done, key):
        self.replay_buffer.append((key, (last_obs, last_phase, actions, rewards, obs, cur_phase)))

    def _batchwise(self, samples):
        batch_list = []
        batch_list_p = []
        actions = []
        rewards = []
        for item in samples:
            dp = item[1]
            state = torch.tensor(dp[0], dtype=torch.float32)
            batch_list.append(Data(x=state, edge_index=self.edge_idx))

            state_p = torch.tensor(dp[4], dtype=torch.float32)
            batch_list_p.append(Data(x=state_p, edge_index=self.edge_idx))
            rewards.append(dp[3])
            actions.append(dp[2])
        batch_t = Batch.from_data_list(batch_list)
        batch_tp = Batch.from_data_list(batch_list_p)
        rewards = torch.tensor(np.array(rewards), dtype=torch.float32)
        actions = torch.tensor(np.array(actions), dtype=torch.long)
        if self.sub_agents > 1:
            rewards = rewards.view(rewards.shape[0] * rewards.shape[1])
            actions = actions.view(actions.shape[0] * actions.shape[1])
        return batch_t, batch_tp, rewards, actions

    def train(self):
        if len(self.replay_buffer) < self.batch_size:
            return 0

        samples = random.sample(self.replay_buffer, self.batch_size)
        b_t, b_tp, rewards, actions = self._batchwise(samples)

        # Ensure all tensors are on the correct device
        b_t = b_t.to(device)
        b_tp = b_tp.to(device)
        rewards = rewards.to(device)
        actions = actions.to(device)

        # Generate masks for the batch
        tgt_mask = self.generate_square_subsequent_mask(b_t.num_graphs).to(device)
        memory_mask = None  # Define as needed

        current_q_values, current_attention = self.model(x=b_t.x, edge_index=b_t.edge_index, train=True, tgt_mask=tgt_mask, memory_mask=memory_mask)
        target_q_values, _ = self.target_model(x=b_tp.x, edge_index=b_tp.edge_index, train=False, tgt_mask=None, memory_mask=None)

        target_q_values = rewards + self.gamma * torch.max(target_q_values, dim=1)[0]

        current_q_values = current_q_values.gather(1, actions.unsqueeze(1)).squeeze(1)

        loss = self.criterion(current_q_values, target_q_values)
        self.optimizer.zero_grad()
        loss.backward()
        clip_grad_norm_(self.model.parameters(), self.grad_clip)
        self.optimizer.step()

        # Update episode count
        self.episode_count += 1

        # Update intersection importance based on the episode count
        if self.episode_count >= self.update_interval:
            self.update_intersection_importance(current_attention)
            self.update_interval = int(self.update_interval * self.decay_factor)  # Increase the interval
            self.episode_count = 0  # Reset episode count

        if self.epsilon > self.epsilon_min:
            self.epsilon *= self.epsilon_decay

        return loss.item()

    def update_target_network(self):
        weights = self.model.state_dict()
        self.target_model.load_state_dict(weights)

    def load_model(self, e):
        model_name = os.path.join(Registry.mapping['logger_mapping']['path'].path,
                                  'model', f'{e}_{self.rank}.pt')        
        self.model = self._build_model()
        self.model.load_state_dict(torch.load(model_name))
        self.target_model = self._build_model()
        self.target_model.load_state_dict(torch.load(model_name))

    def save_model(self, e):
        path = os.path.join(Registry.mapping['logger_mapping']['path'].path, 'model')
        if not os.path.exists(path):
            os.makedirs(path)
        model_name = os.path.join(path, f'{e}_{self.rank}.pt')
        torch.save(self.target_model.state_dict(), model_name)


class TestLightNet(nn.Module):
    def __init__(self, input_dim, output_dim, intersection_importance, **kwargs):
        super(TestLightNet, self).__init__()
        self.model_dict = kwargs
        self.action_space = gym.spaces.Discrete(output_dim)
        self.features = input_dim
        self.intersection_importance = intersection_importance
        self.softmax_temperature = kwargs.get('softmax_temperature', 1.0)
        self.softmax_feather = kwargs.get('softmax_feather', 1.0)
        self.update_interval = kwargs.get('update_interval', 100)
        self.decay_factor = kwargs.get('decay_factor', 1.05)
        
        # Embedding MLP
        self.embedding_MLP = Embedding_MLP(self.features, layers=self.model_dict.get('NODE_EMB_DIM'))

        self.embedding_output_dim = self.model_dict.get('NODE_EMB_DIM')[-1]

        # Lower Transformer (Encoder)
        self.num_lower_transformer_layers = self.model_dict.get('num_lower_transformer_layers', 1)
        self.lower_transformer_nhead = self.model_dict.get('lower_transformer_nhead', 4)
        self.lower_transformer_dim_feedforward = self.model_dict.get('lower_transformer_dim_feedforward', 2048)
        self.lower_transformer_dropout = self.model_dict.get('lower_transformer_dropout', 0.1)

        self.lower_transformer = TransformerBlock(
            d_model=self.embedding_output_dim,
            nhead=self.lower_transformer_nhead,
            num_layers=self.num_lower_transformer_layers,
            dim_feedforward=self.lower_transformer_dim_feedforward,
            dropout=self.lower_transformer_dropout,
            is_encoder=True  # Initialize as encoder
        )

        # GAT Layers
        self.num_gat_layers = self.model_dict.get('N_LAYERS', 5)
        self.num_heads = self.model_dict.get('NUM_HEADS', [1])[0]

        self.head_dim = self.embedding_output_dim // self.num_heads

        self.gat_layers = nn.ModuleList([
            MultiHeadAttModel(d=self.embedding_output_dim, dv=self.head_dim, d_out=self.embedding_output_dim, nv=self.num_heads, suffix=f'layer_{i}')
            for i in range(self.num_gat_layers)
        ])

        self.skip_connections = nn.ModuleList([
            nn.Linear(self.embedding_output_dim, self.embedding_output_dim)
            for _ in range(self.num_gat_layers)
        ])

        self.layer_norms = nn.ModuleList([
            nn.LayerNorm(self.embedding_output_dim)
            for _ in range(self.num_gat_layers)
        ])

        # Upper Transformer (Decoder)
        self.num_upper_transformer_layers = self.model_dict.get('num_upper_transformer_layers', 1)
        self.upper_transformer_nhead = self.model_dict.get('upper_transformer_nhead', 4)
        self.upper_transformer_dim_feedforward = self.model_dict.get('upper_transformer_dim_feedforward', 2048)
        self.upper_transformer_dropout = self.model_dict.get('upper_transformer_dropout', 0.1)

        self.upper_transformer = TransformerBlock(
            d_model=self.embedding_output_dim,
            nhead=self.upper_transformer_nhead,
            num_layers=self.num_upper_transformer_layers,
            dim_feedforward=self.upper_transformer_dim_feedforward,
            dropout=self.upper_transformer_dropout,
            is_encoder=False  # Initialize as decoder
        )

        # Output layer
        self.output_layer = nn.Linear(self.embedding_output_dim, self.action_space.n)

        self.relu = nn.ReLU(inplace=False)

    def forward(self, x, edge_index, train=True, tgt_mask=None, memory_mask=None):
        print(f"Input x shape: {x.shape}")
        print(f"Edge index shape: {edge_index.shape}")

        h = self.embedding_MLP(x, train)
        h = h.unsqueeze(0)  # Shape: (1, num_nodes, embedding_dim)

        # Prepare memory for Decoder (use lower transformer as memory)
        memory = self.lower_transformer(h)  # Encoder part
        memory = memory.squeeze(0)  # Shape: (num_nodes, embedding_dim)
        h = h.squeeze(0)  # Back to (num_nodes, embedding_dim)

        # GAT layers
        attention_matrices = []
        importance_scores = torch.tensor([self.intersection_importance.get(i, 0.5) for i in range(h.size(0))], dtype=torch.float32, device=h.device)

        for layer_idx in range(self.num_gat_layers):
            h_new, att = self.gat_layers[layer_idx](h, edge_index, train)
            h_skip = self.skip_connections[layer_idx](h)

            # Blend new and skip connections based on importance score
            blend_factor = importance_scores.unsqueeze(1) * ((layer_idx + 1) / self.num_gat_layers)
            h = blend_factor * h_new + (1 - blend_factor) * h_skip

            h = self.layer_norms[layer_idx](h)
            h = self.relu(h)

            attention_matrices.append(att)

        h = h.unsqueeze(0)  # Shape: (1, num_nodes, embedding_dim)
        memory = memory.unsqueeze(0)  # Shape: (1, num_nodes, embedding_dim)

        # Decoder Transformer using h as target and memory from lower transformer
        h = self.upper_transformer(h, memory=memory, tgt_mask=tgt_mask, memory_mask=memory_mask)
        h = h.squeeze(0)

        # Proceed with output_layer
        h_importance = h * (1 + importance_scores.unsqueeze(1)) / self.softmax_feather
        q_values = self.output_layer(h_importance)

        return q_values, attention_matrices


class Embedding_MLP(nn.Module):
    def __init__(self, in_size, layers):
        super(Embedding_MLP, self).__init__()
        constructor_dict = OrderedDict()
        for l_idx, l_size in enumerate(layers):
            name = f"node_embedding_{l_idx}"
            if l_idx == 0:
                h = nn.Linear(in_size, l_size)
                constructor_dict.update({name: h})
            else:
                h = nn.Linear(layers[l_idx - 1], l_size)
                constructor_dict.update({name: h})
            name = f"n_relu_{l_idx}"
            constructor_dict.update({name: nn.ReLU()})
        self.embedding_node = nn.Sequential(constructor_dict)

    def _forward(self, x):
        x = self.embedding_node(x)
        return x

    def forward(self, x, train=True):
        if train:
            return self._forward(x)
        else:
            with torch.no_grad():
                return self._forward(x)


class MultiHeadAttModel(MessagePassing):
    def __init__(self, d, dv, d_out, nv, suffix):
        super(MultiHeadAttModel, self).__init__(aggr='add')
        self.d = d
        self.dv = dv
        self.d_out = d_out
        self.nv = nv
        self.suffix = suffix
        self.W_target = nn.Linear(d, dv * nv)
        self.W_source = nn.Linear(d, dv * nv)
        self.hidden_embedding = nn.Linear(d, dv * nv)
        self.out = nn.Linear(dv * nv, d_out)
        self.att = None

    def forward(self, x, edge_index, train=True):
        if train:
            return self._forward(x, edge_index)
        else:
            with torch.no_grad():
                return self._forward(x, edge_index)

    def _forward(self, x, edge_index):
        edge_index, _ = add_self_loops(edge_index=edge_index)
        aggregated = self.propagate(edge_index=edge_index, x=x)
        out = self.out(aggregated)
        out = F.relu(out)
        return out, self.att

    def message(self, x_i, x_j, edge_index_i):
        h_target = F.relu(self.W_target(x_i))
        h_target = h_target.view(-1, self.nv, self.dv)
        agent_repr = h_target.permute(1, 0, 2)

        h_source = F.relu(self.W_source(x_j))
        h_source = h_source.view(-1, self.nv, self.dv)
        neighbor_repr = h_source.permute(1, 0, 2)
        index = edge_index_i

        e_i = torch.mul(agent_repr, neighbor_repr).sum(-1)
        max_node = torch_scatter.scatter_max(e_i, index=index)[0]
        max_i = max_node[:, index]
        ec_i = e_i - max_i
        ecexp_i = torch.exp(ec_i)
        norm_node = torch_scatter.scatter_sum(ecexp_i, index=index)
        norm_node = norm_node[:, index]
        normst_i = norm_node + 1e-12

        alpha_i = ecexp_i / normst_i
        alpha_i_expand = alpha_i.unsqueeze(-1)
        alpha_i_expand = alpha_i_expand.permute(1, 2, 0)

        hidden_neighbor = F.relu(self.hidden_embedding(x_j))
        hidden_neighbor = hidden_neighbor.view(-1, self.nv, self.dv)
        hidden_neighbor_repr = hidden_neighbor.permute(1, 0, 2)

        # Debug prints
        print(f"alpha_i_expand shape: {alpha_i_expand.shape}")
        print(f"hidden_neighbor_repr shape: {hidden_neighbor_repr.shape}")

        # Ensure alpha_i_expand has the correct shape
        if alpha_i_expand.shape[2] != hidden_neighbor_repr.shape[1]:
            alpha_i_expand = alpha_i_expand.expand(-1, -1, hidden_neighbor_repr.shape[1])

        out = (hidden_neighbor_repr * alpha_i_expand).mean(0)

        self.att = alpha_i
        return out

    def get_att(self):
        if self.att is None:
            print('invalid att')
        return self.att


class TransformerBlock(nn.Module):
    def __init__(self, d_model, nhead, num_layers, dim_feedforward=2048, dropout=0.1, is_encoder=True):
        super(TransformerBlock, self).__init__()
        
        if is_encoder:
            encoder_layer = nn.TransformerEncoderLayer(
                d_model=d_model,
                nhead=nhead,
                dim_feedforward=dim_feedforward,
                dropout=dropout,
                batch_first=True
            )
            self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        else:
            decoder_layer = nn.TransformerDecoderLayer(
                d_model=d_model,
                nhead=nhead,
                dim_feedforward=dim_feedforward,
                dropout=dropout,
                batch_first=True
            )
            self.transformer = nn.TransformerDecoder(decoder_layer, num_layers=num_layers)
        self.is_encoder = is_encoder
    
    def forward(self, input, memory=None, tgt_mask=None, memory_mask=None,
                tgt_key_padding_mask=None, memory_key_padding_mask=None):
        if self.is_encoder:
            # For encoder, memory is not used
            output = self.transformer(
                src=input, 
                mask=tgt_mask,
                src_key_padding_mask=tgt_key_padding_mask
            )
            return output
        else:
            # For decoder, input is tgt, and memory comes from encoder
            output = self.transformer(
                tgt=input,
                memory=memory,
                tgt_mask=tgt_mask,
                memory_mask=memory_mask,
                tgt_key_padding_mask=tgt_key_padding_mask,
                memory_key_padding_mask=memory_key_padding_mask
            )
            return output"

"root@520ef09ed306:/DaRL/LibSignal# python run.py
Using device: cuda
Using device: cuda:0 for BASETRAIN---------------------------------------------------------
building world...
creating intersections...
intersections created.
parsing roads...
roads parsed.
world built.
<agent.testlight.TestLightAgent object at 0x7fc7d565b430>
Using device: cuda:0 for training---------------------------------------------------------
2024-09-16 19:00:00 (INFO): step:3600/3600, q_loss:0, rewards:-40.19110870361328, queue:2.50451397895813, delay:0.13135415315628052, throughput:1471
2024-09-16 19:00:00 (INFO): episode:0/200, real avg travel time:227.81602172437204
Using device: cuda:0 for training test at episode 0
Traceback (most recent call last):
  File "/DaRL/LibSignal/run.py", line 88, in <module>
    test.run()
  File "/DaRL/LibSignal/run.py", line 82, in run
    self.task.run()
  File "/DaRL/LibSignal/task/task.py", line 48, in run
    raise e
  File "/DaRL/LibSignal/task/task.py", line 43, in run
    self.trainer.train()
  File "/DaRL/LibSignal/trainer/tsc_trainer.py", line 201, in train
    self.train_test(e)
  File "/DaRL/LibSignal/trainer/tsc_trainer.py", line 224, in train_test
    actions.append(ag.get_action(obs[idx], phases[idx], test=True))
  File "/DaRL/LibSignal/agent/testlight.py", line 226, in get_action
    q_values, attention_matrices = self.model(x=observation, edge_index=edge_index, train=False, tgt_mask=tgt_mask, memory_mask=memory_mask)
  File "/usr/local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/DaRL/LibSignal/agent/testlight.py", line 461, in forward
    h_new, att = self.gat_layers[layer_idx](h, edge_index, train)
  File "/usr/local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/DaRL/LibSignal/agent/testlight.py", line 534, in forward
    return self._forward(x, edge_index)
  File "/DaRL/LibSignal/agent/testlight.py", line 538, in _forward
    aggregated = self.propagate(edge_index=edge_index, x=x)
  File "/tmp/agent.testlight_MultiHeadAttModel_propagate_v1lmjcxv.py", line 176, in propagate
    out = self.message(
  File "/DaRL/LibSignal/agent/testlight.py", line 569, in message
    out = (hidden_neighbor_repr * alpha_i_expand).mean(0)
RuntimeError: The size of tensor a (51) must match the size of tensor b (5) at non-singleton dimension 2"