fix this so that it works :"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
from torch.nn import TransformerEncoder, TransformerEncoderLayer


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))
        sorted(observation_generators, 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))
        sorted(rewarding_generators, 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))
        sorted(queues, 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))
        sorted(delays, 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))
        sorted(phasing_generators, 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):
        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))
        sorted(observation_generators, key=lambda x: x[0])  # now generator's order is according to its index in graph
        self.ob_generator = observation_generators

        #  get reward generator for CoLightAgent
        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))
        sorted(rewarding_generators, key=lambda x: x[0])  # now generator's order is according to its index in graph
        self.reward_generator = rewarding_generators

        #  phase generator
        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))
        sorted(phasing_generators, key=lambda x: x[0])  # now generator's order is according to its index in graph
        self.phase_generator = phasing_generators

        # queue metric
        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))
        # now generator's order is according to its index in graph
        sorted(queues, key=lambda x: x[0])
        self.queue = queues

        # delay metric
        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))
        # now generator's order is according to its index in graph
        sorted(delays, 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)
        
        q_values, attention_matrices = self.model(x=observation, edge_index=edge_index, train=False)
        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 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)

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

        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)
        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]
        
        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.spatial_attention = nn.Parameter(torch.Tensor(1, self.num_heads, self.embedding_output_dim))
        self.temporal_attention = nn.Parameter(torch.Tensor(1, self.num_heads, self.embedding_output_dim))

        # Lower Transformer Encoder
        self.lower_transformer_layer = TransformerEncoderLayer(
            d_model=self.embedding_output_dim, 
            nhead=self.num_heads
        )
        self.lower_transformer = TransformerEncoder(
            self.lower_transformer_layer, 
            num_layers=1
        )

        # Upper Transformer Encoder
        self.upper_transformer_layer = TransformerEncoderLayer(
            d_model=self.embedding_output_dim, 
            nhead=self.num_heads
        )
        self.upper_transformer = TransformerEncoder(
            self.upper_transformer_layer, 
            num_layers=1
        )

        # Actor-Critic Networks
        self.actor = nn.Sequential(
            nn.Linear(self.embedding_output_dim, 128),
            nn.ReLU(),
            nn.Linear(128, output_dim)
        )
        self.critic = nn.Sequential(
            nn.Linear(self.embedding_output_dim, 128),
            nn.ReLU(),
            nn.Linear(128, 1)
        )

    def forward(self, x, edge_index, train=True):
        # Embedding
        x = self.embedding_MLP(x)

        # Lower Transformer
        x = self.lower_transformer(x)

        # GAT Layers with Skip Connections
        for gat_layer, skip_connection in zip(self.gat_layers, self.skip_connections):
            x = gat_layer(x, edge_index) + skip_connection(x)

        # Upper Transformer
        x = self.upper_transformer(x)

        # Actor-Critic Outputs
        action_probs = self.actor(x)
        state_value = self.critic(x)

        return action_probs, state_value

class Embedding_MLP(nn.Module):
    def __init__(self, input_dim, layers):
        super(Embedding_MLP, self).__init__()
        self.layers = nn.ModuleList()
        for i in range(len(layers)):
            self.layers.append(nn.Linear(input_dim if i == 0 else layers[i-1], layers[i]))
            self.layers.append(nn.ReLU())

    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        return x

class MultiHeadAttModel(MessagePassing):
    def __init__(self, d, dv, d_out, nv, suffix):
        super(MultiHeadAttModel, self).__init__(aggr='add')
        self.nv = nv
        self.d = d
        self.dv = dv
        self.d_out = d_out
        self.suffix = suffix
        self.att = nn.Parameter(torch.Tensor(1, nv, d))
        self.lin = nn.Linear(d, d_out)
        self.reset_parameters()

    def reset_parameters(self):
        nn.init.xavier_uniform_(self.att)
        nn.init.xavier_uniform_(self.lin.weight)

    def forward(self, x, edge_index):
        return self.propagate(edge_index, x=x)

    def message(self, x_j, edge_index, size):
        x_j = x_j.view(-1, self.nv, self.d)
        alpha = (x_j * self.att).sum(dim=-1)
        alpha = F.leaky_relu(alpha, negative_slope=0.2)
        alpha = F.softmax(alpha, dim=1)
        return x_j * alpha.unsqueeze(-1)

    def update(self, aggr_out):
        return self.lin(aggr_out.view(-1, self.nv * self.dv))

# A3C Training Loop
def a3c_train(agent, env, optimizer, gamma=0.99):
    state = env.reset()
    done = False
    while not done:
        action_probs, state_value = agent.model(torch.tensor(state, dtype=torch.float32).to(device))
        action = torch.multinomial(F.softmax(action_probs, dim=-1), 1).item()
        next_state, reward, done, _ = env.step(action)
        next_action_probs, next_state_value = agent.model(torch.tensor(next_state, dtype=torch.float32).to(device))

        td_target = reward + gamma * next_state_value * (1 - done)
        td_error = td_target - state_value

        actor_loss = -torch.log(action_probs[action]) * td_error
        critic_loss = td_error.pow(2)

        loss = actor_loss + critic_loss

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        state = next_state"
DESCRIPTION:"### Model Idea and Architecture

The updated `TestLightAgent` model combines the original features of the `TestLight` model with the advanced concepts from the provided image. The main enhancements include the integration of Lower and Upper Transformers and the replacement of the DQN-based action choice with an A3C (Asynchronous Advantage Actor-Critic) method. Here is a detailed description of the combined model idea and architecture:

#### 1. **Model Overview**
The model is designed to manage traffic signal control using reinforcement learning. It leverages graph-based representations of intersections and their connections, and it incorporates advanced neural network architectures to process and learn from the data.

#### 2. **Lower Transformer**
- **Purpose**: The Lower Transformer is used to enhance the embedding of the input features. It captures the temporal dynamics and interactions between the features of each intersection.
- **Implementation**: The Lower Transformer is applied immediately after the initial embedding MLP (Multi-Layer Perceptron). It processes the embedded features to produce a more informative representation.

#### 3. **Graph Attention Network (GAT) Layers**
- **Purpose**: The GAT layers are used to model the interactions between intersections. They allow the model to focus on the most relevant neighboring intersections when making decisions.
- **Implementation**: Multiple GAT layers are used, each with skip connections to preserve the original features. These layers aggregate information from neighboring nodes (intersections) using attention mechanisms.

#### 4. **Upper Transformer**
- **Purpose**: The Upper Transformer is used to further process the output of the GAT layers. It captures higher-level temporal and spatial dependencies in the data.
- **Implementation**: The Upper Transformer is applied after the GAT layers. It processes the output of the GAT layers to produce the final representation used for decision-making.

#### 5. **Actor-Critic Networks**
- **Purpose**: The Actor-Critic networks are used for the A3C method. The actor network selects actions based on the current state, while the critic network evaluates the value of the current state.
- **Implementation**: The actor network outputs action probabilities, and the critic network outputs a state value. These networks are trained using the A3C algorithm, which optimizes both the policy (actor) and the value function (critic).

#### 6. **A3C Training Loop**
- **Purpose**: The A3C training loop is used to train the model. It updates the actor and critic networks based on the observed rewards and state transitions.
- **Implementation**: The training loop samples actions from the actor network, observes the resulting state and reward, and updates the networks using the TD (Temporal Difference) error.

### Combined Architecture

1. **Input Features**: The input features include observations from lane vehicle generators, phase generators, queue metrics, and delay metrics for each intersection.

2. **Embedding MLP**: The input features are first processed by an embedding MLP to produce an initial representation.

3. **Lower Transformer**: The embedded features are then processed by the Lower Transformer to capture temporal dynamics.

4. **GAT Layers**: The output of the Lower Transformer is processed by multiple GAT layers with skip connections to model interactions between intersections.

5. **Upper Transformer**: The output of the GAT layers is further processed by the Upper Transformer to capture higher-level dependencies.

6. **Actor-Critic Networks**: The final representation is used by the actor and critic networks to select actions and evaluate state values.

7. **A3C Training**: The model is trained using the A3C algorithm, which updates the actor and critic networks based on observed rewards and state transitions.

### Summary

The combined model leverages the strengths of both the original `TestLight` model and the advanced concepts from the provided image. The Lower and Upper Transformers enhance the model's ability to capture temporal and spatial dependencies, while the A3C method provides a more robust and efficient way to learn optimal policies for traffic signal control. This architecture aims to improve the model's performance and generalization capabilities in managing complex traffic scenarios."

WORKING OLD VERSION OF CODE WITHOUT NEW FEATURES (working):"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))
        sorted(observation_generators, 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))
        sorted(rewarding_generators, 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))
        sorted(queues, 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))
        sorted(delays, 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))
        sorted(phasing_generators, 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):
        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))
        sorted(observation_generators, key=lambda x: x[0])  # now generator's order is according to its index in graph
        self.ob_generator = observation_generators

        #  get reward generator for CoLightAgent
        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))
        sorted(rewarding_generators, key=lambda x: x[0])  # now generator's order is according to its index in graph
        self.reward_generator = rewarding_generators

        #  phase generator
        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))
        sorted(phasing_generators, key=lambda x: x[0])  # now generator's order is according to its index in graph
        self.phase_generator = phasing_generators

        # queue metric
        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))
        # now generator's order is according to its index in graph
        sorted(queues, key=lambda x: x[0])
        self.queue = queues

        # delay metric
        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))
        # now generator's order is according to its index in graph
        sorted(delays, 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)
        
        q_values, attention_matrices = self.model(x=observation, edge_index=edge_index, train=False)
        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 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)

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

        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)
        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]
        
        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)
        ])
        
        self.output_layer = nn.Linear(self.embedding_output_dim, self.action_space.n)
        
        self.relu = nn.ReLU(inplace=False)

        # Add new parameters


    def forward(self, x, edge_index, train=True):
        h = self.embedding_MLP(x, train)
        
        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)
        
        # Apply intersection importance after all GAT layers, but more gently
        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, 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):
        h_target = F.relu(self.W_target(x_i))
        h_target = h_target.view(h_target.shape[:-1][0], 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(h_source.shape[:-1][0], self.nv, self.dv)

        neighbor_repr = h_source.permute(1, 0, 2)
        index = edge_index[1]

        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_select(1, index)
        ec_i = torch.add(e_i, -max_i)
        ecexp_i = torch.exp(ec_i)
        norm_node = torch_scatter.scatter_sum(ecexp_i, index=index)
        normst_node = torch.add(norm_node, 1e-12)
        normst_i = normst_node.index_select(1, index)

        alpha_i = ecexp_i / normst_i
        alpha_i_expand = alpha_i.repeat(self.dv, 1, 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(hidden_neighbor.shape[:-1][0], self.nv, self.dv)
        hidden_neighbor_repr = hidden_neighbor.permute(1, 0, 2)
        out = torch.mul(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"