"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)
        
        # 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
        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
        )

        # 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
        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
        )

        # 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):
        h = self.embedding_MLP(x, train)

        # Apply Lower Transformer
        h = h.unsqueeze(0)  # shape [1, N, D]
        h = self.lower_transformer(h)
        h = h.squeeze(0)    # shape back to [N, D]

        # 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)

        # Apply Upper Transformer
        h = h.unsqueeze(0)  # shape [1, N, D]
        h = self.upper_transformer(h)
        h = h.squeeze(0)    # shape back to [N, D]

        # 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, 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

class TransformerBlock(nn.Module):
    def __init__(self, d_model, nhead, num_layers, dim_feedforward=2048, dropout=0.1):
        super(TransformerBlock, self).__init__()
        encoder_layer = nn.TransformerEncoderLayer(d_model=d_model,
                                                   nhead=nhead,
                                                   dim_feedforward=dim_feedforward,
                                                   dropout=dropout,
                                                   batch_first=True)
        self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)

    def forward(self, x):
        # x shape: [batch size, sequence length, embedding size]
        output = self.transformer_encoder(x)
        return output"
change transformer block from using an encoder layer to use a decoder layer