this is my dynamiclight implementation for LibSignal:"from . import RLAgent
import random
import numpy as np
from collections import deque
import os
import torch
import torch.nn as nn
import torch.nn.functional as F
from common.registry import Registry
import gym
from generator import LaneVehicleGenerator, IntersectionPhaseGenerator
from torch.nn.utils import clip_grad_norm_
from agent.utils import SharedEpsGreedy
from pfrl.agents import DQN
from pfrl import replay_buffers
from pfrl.utils.contexts import evaluating
from agent import utils
import copy

'''DynamicLight is default set Shared Agent'''

@Registry.register_model('dynamiclight')
class DynamicLightAgent(RLAgent):
    def __init__(self, world, rank):
        super().__init__(world, world.intersection_ids[rank])
        self.world = world
        self.rank = rank
        self.id = self.world.intersection_ids[self.rank]
        self.inter = self.world.id2intersection[self.id]

        # Load configurations
        self.dic_agent_conf = Registry.mapping['model_mapping']['setting'].param
        model_conf = self.dic_agent_conf
        self.gamma = model_conf.get('gamma', 0.8)
        self.epsilon = model_conf.get('epsilon', 0.8)
        self.epsilon_decay = model_conf.get('epsilon_decay', 0.95)
        self.eps_min = model_conf.get('eps_min', 0.2)
        self.learning_rate = model_conf.get('learning_rate', 0.001)
        self.batch_size = model_conf.get('batch_size', 128)
        self.grad_clip = model_conf.get('grad_clip', 5.0)
        self.max_memory_len = model_conf.get('max_memory_len', 12000)
        self.num_action_dur = model_conf.get('num_action_dur', 7)
        self.num_feat1 = model_conf.get('num_feat1', 3)
        self.sround = model_conf.get('sround', 80)
        self.sround2 = model_conf.get('sround2', 160)
        self.sround3 = model_conf.get('sround3', 110)
        self.new_learning_rate = model_conf.get('new_learning_rate', 0.0001)
        self.phase = model_conf.get('phase', True)
        self.one_hot = model_conf.get('one_hot', False)

        # Load world configurations
        self.dic_world_conf = Registry.mapping['world_mapping']['setting'].param
        self.network_name = self.dic_world_conf.get('network', 'default')
        network_config = self.dic_world_conf['signal_config'][self.network_name]
        self.phase_pairs = network_config['phase_pairs']
        self.num_phases = len(self.phase_pairs)
        self.max_lane = network_config.get('max_lane', 16)
        self.comp_mask = self.relation()
        self.valid_acts = network_config.get('valid_acts', None)
        self.phase_map = self.phase_pairs  # For consistency with DynamicLightNet

        # Initialize other attributes
        self.memory = None  # We will use pfrl's replay buffer
        self.cnt_round = 0
        self.sub_agents = 1  # Control only one intersection per agent
        self.action_space = gym.spaces.Discrete(self.num_phases)

        # Initialize explorer
        trainer_conf = Registry.mapping['trainer_mapping']['setting'].param
        episodes = trainer_conf.get('episodes', 200) * 0.8
        steps = trainer_conf.get('steps', 1000)
        action_interval = trainer_conf.get('action_interval', 10)
        total_steps = episodes * steps / action_interval
        self.explorer = SharedEpsGreedy(
            self.epsilon,
            self.eps_min,
            total_steps,
            lambda: np.random.randint(self.num_phases),
        )

        # Initialize replay buffer using pfrl's ReplayBuffer
        from pfrl import replay_buffers
        self.replay_buffer = replay_buffers.ReplayBuffer(self.max_memory_len)

        # Initialize agents_iner
        self.agents_iner = self._build_model()

        # Initialize generators
        '''
        reset
        Reset information, including ob_generator, phase_generator, queue, delay, etc.

        :param: None
        :return: None
        '''
        observation_generators = []
        for inter in self.world.intersections:
            node_id = inter.id
            node_idx = self.world.id2idx[node_id]
            node_obj = self.world.id2intersection[node_id]
            tmp_generator = LaneVehicleGenerator(self.world, node_obj, ['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

        rewarding_generators = []
        for inter in self.world.intersections:
            node_id = inter.id
            node_idx = self.world.id2idx[node_id]
            node_obj = self.world.id2intersection[node_id]
            tmp_generator = LaneVehicleGenerator(self.world, node_obj, ["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

        phasing_generators = []
        for inter in self.world.intersections:
            node_id = inter.id
            node_idx = self.world.id2idx[node_id]
            node_obj = self.world.id2intersection[node_id]
            tmp_generator = IntersectionPhaseGenerator(self.world, node_obj, ['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

        queues = []
        for inter in self.world.intersections:
            node_id = inter.id
            node_idx = self.world.id2idx[node_id]
            node_obj = self.world.id2intersection[node_id]
            tmp_generator = LaneVehicleGenerator(self.world, node_obj, ["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
            node_idx = self.world.id2idx[node_id]
            node_obj = self.world.id2intersection[node_id]
            tmp_generator = LaneVehicleGenerator(self.world, node_obj, ["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

    def __repr__(self):
        return f"DynamicLightAgent(id={self.id})"

    def _build_model(self):
        '''
        _build_model
        Build a DynamicLightNet agent and optimizer.

        :param: None
        :return agents: DynamicLight_InerAgent consists of agent and optimizer
        '''
        # Initialize the model
        self.model = DynamicLightNet(
            num_phases=self.num_phases,
            num_action_dur=self.num_action_dur,
            max_lane=self.max_lane,
            num_feat1=self.num_feat1,
            phase_map=self.phase_map
        )
        # Initialize the target model
        self.target_model = copy.deepcopy(self.model)

        # Initialize the optimizer and criterion
        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=self.learning_rate)
        self.criterion = nn.MSELoss()

        # Create the agent
        agents = DynamicLight_InerAgent(
            self.model,
            self.optimizer,
            self.replay_buffer,
            self.gamma,
            self.explorer,
            minibatch_size=self.batch_size,
            replay_start_size=self.batch_size,
            phi=lambda x: np.asarray(x, dtype=np.float32),
            target_update_interval=self.dic_agent_conf.get("update_q_bar_freq", 100),
            update_interval=1
        )
        return agents

    def reset(self):
        '''
        reset
        Reset information, including ob_generator, phase_generator, queue, delay, etc.

        :param: None
        :return: None
        '''
        observation_generators = []
        for inter in self.world.intersections:
            node_id = inter.id
            node_idx = self.world.id2idx[node_id]
            node_obj = self.world.id2intersection[node_id]
            tmp_generator = LaneVehicleGenerator(self.world, node_obj, ['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

        rewarding_generators = []
        for inter in self.world.intersections:
            node_id = inter.id
            node_idx = self.world.id2idx[node_id]
            node_obj = self.world.id2intersection[node_id]
            tmp_generator = LaneVehicleGenerator(self.world, node_obj, ["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

        phasing_generators = []
        for inter in self.world.intersections:
            node_id = inter.id
            node_idx = self.world.id2idx[node_id]
            node_obj = self.world.id2intersection[node_id]
            tmp_generator = IntersectionPhaseGenerator(self.world, node_obj, ['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

        queues = []
        for inter in self.world.intersections:
            node_id = inter.id
            node_idx = self.world.id2idx[node_id]
            node_obj = self.world.id2intersection[node_id]
            tmp_generator = LaneVehicleGenerator(self.world, node_obj, ["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
            node_idx = self.world.id2idx[node_id]
            node_obj = self.world.id2intersection[node_id]
            tmp_generator = LaneVehicleGenerator(self.world, node_obj, ["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

    def relation(self):
        '''
        relation
        Get the phase competition relation between traffic movements.

        :param: None
        :return comp_mask: matrix of phase competition relation
        '''
        comp_mask = []
        for i in range(len(self.phase_pairs)):
            zeros = np.zeros(len(self.phase_pairs) - 1, dtype=np.int64)
            cnt = 0
            for j in range(len(self.phase_pairs)):
                if i == j:
                    continue
                pair_a = self.phase_pairs[i]
                pair_b = self.phase_pairs[j]
                if len(list(set(pair_a + pair_b))) == 3:
                    zeros[cnt] = 1
                cnt += 1
            comp_mask.append(zeros)
        comp_mask = torch.from_numpy(np.asarray(comp_mask))
        return comp_mask

    def get_ob(self):
        '''
        get_ob
        Get observation from environment.

        :param: None
        :return x_obs: observation generated by ob_generator, the shape is [sub_agents,lane_nums]
        '''
        x_obs = []  # sub_agents * lane_nums,
        for i in range(len(self.ob_generator)):
            tmp = self.ob_generator[i][1].generate()
            if self.ob_order != None:
                tt = []
                if self.ob_generator[i][1].I.id[:3] == 'GS_':
                    name = self.ob_generator[i][1].I.id[3:]
                else:
                    name = self.ob_generator[i][1].I.id
                for i in range(12):
                    # padding to 12 dims
                    if i in self.ob_order[name].keys():
                        tt.append(tmp[self.ob_order[name][i]])
                    else:
                        tt.append(0.)
                x_obs.append(np.array(tt))     
            else:
                x_obs.append(tmp)
            
        if self.ob_order != None:
            x_obs = np.array(x_obs, dtype=np.float32)
        return x_obs

    def get_reward(self):
        '''
        get_reward
        Get reward from environment.

        :param: None
        :return rewards: rewards generated by reward_generator
        '''
        # TODO: test output
        rewards = []  # sub_agents
        for i in range(len(self.reward_generator)):
            rewards.append(self.reward_generator[i][1].generate())
        rewards = np.squeeze(np.array(rewards))
        return rewards

    def get_phase(self):
        '''
        get_phase
        Get current phase of intersection(s) from environment.

        :param: None
        :return phase: current phase generated by phase_generator, the shape is [sub_agents,]
        '''
        phase = []  # sub_agents
        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):
        '''
        get_queue
        Get queue length of intersection.

        :param: None
        :return queue: value(one intersection) or [intersections,](multiple intersections)
        '''
        queue = []
        for i in range(len(self.queue)):
            queue.append((self.queue[i][1].generate()))
        maxsize = max(len(x) for x in queue)
        for i in range(len(queue)):
            queue[i] = np.pad(queue[i], (0, maxsize - queue[i].shape[-1]))
        tmp_queue = np.squeeze(np.array(queue))
        if self.sub_agents == 1:
            queue = np.sum(tmp_queue)
        else:
            queue = [np.sum(x) for x in tmp_queue]
        # queue = np.sum(tmp_queue, axis=1 if len(tmp_queue.shape)==1 else 0)
        return queue # [intersections,]

    def get_delay(self):
        '''
        get_delay
        Get delay of intersection.

        :param: None
        :return: total delay
        '''
        delay = []
        for i in range(len(self.delay)):
            delay.append((self.delay[i][1].generate()))
        delay = np.squeeze(np.array(delay))
        return delay # [intersections,]

    def get_action(self, ob, phase, test=False):
        '''
        get_action
        Generate action.

        :param ob: current observation, the shape is [agents, ob_length]
        :param phase: current phase, the shape is [agents,]
        :param test: boolean, decide whether it is a training or testing process
        :return acts: action taken by environment, the shape is [agents,]
        '''

        # input are np.array here
        # TODO: support irregular input in the future

        # concate ob and phase
        if self.phase:
            if self.one_hot:
                obs = np.concatenate([utils.idx2onehot(phase, self.action_space.n), ob], axis=1)
            else:
                obs = [np.concatenate([np.array([phase[i]]), ob[i]]) for i in range(self.sub_agents)]
        else:
            obs = ob

        if self.reverse_valid is None and self.valid_acts is not None:
            self.reverse_valid = dict()
            for signal_id in self.valid_acts:
                self.reverse_valid[signal_id] = {v: k for k, v in self.valid_acts[signal_id].items()}

        batch_obs = obs
        if self.valid_acts is None:
            batch_valid = None
            batch_reverse = None
        else:
            # TODO debug
            # batch_valid = [self.valid_acts.get(agent_id) for agent_id in ob.keys()]
            # batch_reverse = [self.reverse_valid.get(agent_id) for agent_id in ob.keys()]
            dic = {index: key for index, key in enumerate(self.ob_order.keys())}
            batch_valid = [self.valid_acts.get(dic[i]) if dic[i] in self.valid_acts.keys() else self.valid_acts.get('GS_'+dic[i]) for i in range(ob.shape[0])]
            batch_reverse = [self.reverse_valid.get(dic[i]) if dic[i] in self.reverse_valid.keys() else self.reverse_valid.get('GS_'+dic[i])  for i in range(ob.shape[0])]
        batch_acts = self.agents_iner.act(batch_obs,
                                valid_acts=batch_valid,
                                reverse_valid=batch_reverse, test=test)
        acts = np.array(batch_acts)
        return acts

    def remember(self, state_input, cur_phase, action, reward, next_state_input, next_cur_phase, done):
        '''
        remember
        Store the experience in memory for training later.

        :param state_input: current state input
        :param cur_phase: current phase
        :param action: action taken
        :param reward: reward received
        :param next_state_input: next state input
        :param next_cur_phase: next current phase
        :param done: whether the episode is done
        :return: None
        '''
        # Prepare the experience tuple
        experience = {
            'state': (state_input, cur_phase),
            'action': action,
            'reward': reward,
            'next_state': (next_state_input, next_cur_phase),
            'is_state_terminal': done,
        }
        self.replay_buffer.append(experience)

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

        # Sample a batch from the replay buffer
        batch = self.replay_buffer.sample(self.batch_size)

        # Extract elements from the batch
        state_batch = []
        action_batch = []
        reward_batch = []
        next_state_batch = []
        done_batch = []

        for experience in batch:
            state_input, cur_phase = experience['state']
            next_state_input, next_cur_phase = experience['next_state']
            action = experience['action']
            reward = experience['reward']
            done = experience['is_state_terminal']

            # For state and next_state, ensure they are properly concatenated if needed
            # Adjust according to your model's requirements
            state = state_input.reshape(1, -1)
            next_state = next_state_input.reshape(1, -1)

            state_batch.append(state)
            cur_phase_batch = np.array([cur_phase])
            next_state_batch.append(next_state)
            next_cur_phase_batch = np.array([next_cur_phase])
            action_batch.append(action[0])  # action is a tuple (phase_action, duration_action)
            reward_batch.append(reward)
            done_batch.append(done)

        # Convert to tensors
        state_batch = torch.tensor(np.vstack(state_batch), dtype=torch.float32)
        cur_phase_batch = torch.tensor(cur_phase_batch, dtype=torch.long)
        action_batch = torch.tensor(action_batch, dtype=torch.long)
        reward_batch = torch.tensor(reward_batch, dtype=torch.float32)
        next_state_batch = torch.tensor(np.vstack(next_state_batch), dtype=torch.float32)
        next_cur_phase_batch = torch.tensor(next_cur_phase_batch, dtype=torch.long)
        done_batch = torch.tensor(done_batch, dtype=torch.float32)

        # Compute target Q-values
        with torch.no_grad():
            next_phase_scores, _ = self.target_model(next_state_batch, next_cur_phase_batch)
            next_q_values = next_phase_scores.max(1)[0]
            target_q_values = reward_batch + self.gamma * next_q_values * (1 - done_batch)

        # Compute predicted Q-values
        phase_scores, _ = self.model(state_batch, cur_phase_batch)
        predicted_q_values = phase_scores.gather(1, action_batch.unsqueeze(1)).squeeze()

        # Compute loss
        loss = self.criterion(predicted_q_values, target_q_values)

        # Optimize the model
        self.optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.grad_clip)
        self.optimizer.step()

        # Update epsilon
        self.epsilon = max(self.epsilon * self.epsilon_decay, self.eps_min)
        self.cnt_round += 1

        # Update target network periodically
        if self.cnt_round % self.dic_agent_conf.get("update_q_bar_freq", 100) == 0:
            self.update_target_network()

        # Special training rounds (if applicable)
        # ...

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

    def load_model(self, e):
        '''
        load_model
        Load model params of an episode.

        :param e: specified episode
        :return: None
        '''
        model_name = os.path.join(Registry.mapping['logger_mapping']['path'].path,
                                  'model', f'{e}_{self.rank}.pt')
        checkpoint = torch.load(model_name)
        self.model.load_state_dict(checkpoint['model_state_dict'])
        self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
        self.target_model.load_state_dict(self.model.state_dict())

    def save_model(self, e):
        '''
        save_model
        Save model params of an episode.

        :param e: specified episode, used for file name
        :return: None
        '''
        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({
            'model_state_dict': self.model.state_dict(),
            'optimizer_state_dict': self.optimizer.state_dict(),
        }, model_name)


class DynamicLight_InerAgent(DQN):
    '''
    DynamicLight_InerAgent is based on DynamicLightNet.
    '''
    def __init__(self, q_function, optimizer, replay_buffer,
                 gamma, explorer, minibatch_size, replay_start_size, phi,
                 target_update_interval, update_interval):
        super().__init__(q_function, optimizer, replay_buffer, gamma, explorer,
                         minibatch_size=minibatch_size, replay_start_size=replay_start_size, phi=phi,
                         target_update_interval=target_update_interval, update_interval=update_interval)
        self.explorer = explorer  # Ensure the explorer is accessible

    def act(self, obs, test=False):
        '''
        act
        Get action.

        :param obs: observation, shape [1, ob_length]
        :param test: boolean, decide whether it is a training or testing process
        :return: action taken by the agent, as a list with one element
        '''
        state_input_tensor = torch.tensor(obs, dtype=torch.float32)

        if test or np.random.rand() >= self.explorer.epsilon:
            with torch.no_grad(), evaluating(self.model):
                phase_scores, duration_scores = self.model(state_input_tensor, torch.tensor([0]))  # Assuming cur_phase=0
                phase_action = torch.argmax(phase_scores, dim=1).item()
                duration_action = torch.argmax(duration_scores, dim=1).item()
        else:
            phase_action = np.random.randint(0, self.model.num_phases)
            duration_action = np.random.randint(0, self.model.num_action_dur)

        return [(phase_action, duration_action)]  # Return as a list with one action tuple
class DynamicLightNet(nn.Module):
    def __init__(self, num_phases, num_action_dur, max_lane, num_feat1, phase_map):
        super(DynamicLightNet, self).__init__()
        self.num_phases = num_phases
        self.num_action_dur = num_action_dur
        self.max_lane = max_lane
        self.num_feat1 = num_feat1
        self.phase_map = phase_map

        uni_dim = 4
        dims2 = 16

        # Embedding layers
        self.phase_emb = nn.Embedding(num_phases, uni_dim)
        self.feat_embs = nn.ModuleList([
            nn.Linear(1, uni_dim) for _ in range(num_feat1)
        ])
        self.feats_dense = nn.Linear(uni_dim * (num_feat1 + 1), dims2)

        # Attention layers
        self.mha1 = nn.MultiheadAttention(dims2, num_heads=4, batch_first=True)
        self.mha2 = nn.MultiheadAttention(dims2, num_heads=4, batch_first=True)

        # Phase action network
        self.dense_p_hidden1 = nn.Linear(dims2, 20)
        self.dense_p_hidden2 = nn.Linear(20, 20)
        self.dense_p = nn.Linear(20, num_phases)

        # Duration action network (Dueling DQN)
        self.dense_d_hidden1 = nn.Linear(dims2, 20)
        self.dense_d_hidden2 = nn.Linear(20, 20)
        self.dense_value = nn.Linear(20, 1)
        self.dense_advantage = nn.Linear(20, num_action_dur)

    def forward(self, state_input, cur_phase):
        batch_size = state_input.size(0)

        # Embedding
        phase_emb = torch.sigmoid(self.phase_emb(cur_phase))
        feat_emb_list = []
        for i in range(self.num_feat1):
            feat_i = state_input[:, :, i].unsqueeze(-1)  # [batch_size, max_lane, 1]
            feat_emb_i = torch.sigmoid(self.feat_embs[i](feat_i))  # [batch_size, max_lane, uni_dim]
            feat_emb_list.append(feat_emb_i)
        feats = torch.cat(feat_emb_list + [phase_emb.unsqueeze(1).repeat(1, self.max_lane, 1)], dim=-1)
        feats = torch.relu(self.feats_dense(feats))  # [batch_size, max_lane, dims2]

        # Process each phase
        phase_feats = []
        for indices in self.phase_map:
            indices_tensor = torch.tensor(indices, device=state_input.device)
            tmp_feat_1 = feats[:, indices_tensor, :]  # [batch_size, len(indices), dims2]
            tmp_feat_2, _ = self.mha1(tmp_feat_1, tmp_feat_1, tmp_feat_1)  # [batch_size, len(indices), dims2]
            tmp_feat_3 = tmp_feat_2.mean(dim=1, keepdim=True)  # [batch_size, 1, dims2]
            phase_feats.append(tmp_feat_3)
        phase_feat_all = torch.cat(phase_feats, dim=1)  # [batch_size, num_phases, dims2]

        # Attention over phases
        att_encoding, _ = self.mha2(phase_feat_all, phase_feat_all, phase_feat_all)

        # Phase action prediction
        hidden_p = torch.relu(self.dense_p_hidden1(att_encoding.mean(dim=1)))
        hidden_p = torch.relu(self.dense_p_hidden2(hidden_p))
        phase_scores = self.dense_p(hidden_p)

        # Duration action prediction (Dueling DQN)
        selected_phase_feat = att_encoding.mean(dim=1)
        hidden_d = torch.relu(self.dense_d_hidden1(selected_phase_feat))
        hidden_d = torch.relu(self.dense_d_hidden2(hidden_d))
        value = self.dense_value(hidden_d)
        advantage = self.dense_advantage(hidden_d)
        advantage = advantage - advantage.mean(dim=1, keepdim=True)
        duration_scores = value + advantage

        return phase_scores, duration_scores
"

original dynamiclight :""""
DynamicLight under feature fusion method 1
Input shape: [batch, max_lane*4]
Created by Liang Zhang
"""
from tensorflow.keras.layers import Input, Dense, Reshape,  Lambda,  Subtract, Add, MultiHeadAttention, Activation, Embedding
from tensorflow.keras.models import Model
from tensorflow.keras.optimizers import Adam
from .network_agent import NetworkAgent
from tensorflow.keras import backend as K
import numpy as np
import random
import tensorflow as tf
from tensorflow.keras.utils import to_categorical
from tensorflow.keras.callbacks import EarlyStopping
import copy, os


class DynamicLightAgent(NetworkAgent):
    
    def build_network(self):
        ins0 = Input(shape=(self.max_lane, self.num_feat1), name="input_total_features") 
        ins1 = Input(shape=(self.max_lane,), name="input_cur_phase")
        ins2 = Input(shape=(1, self.num_phases))
        uni_dim, dims2 = 4, 16
        phase_emb = Activation('sigmoid')(Embedding(2, uni_dim, input_length=self.max_lane)(ins1))
        feat_list = tf.split(ins0, self.num_feat1, axis=2)
        feat_embs = [Dense(uni_dim, activation='sigmoid')(feat_list[i]) for i in range(self.num_feat1)]
        feat_embs.append(phase_emb)    
        feats = tf.concat(feat_embs, axis=2)
        feats = Dense(dims2, activation="relu")(feats)
        lane_feats = tf.split(feats, self.max_lane, axis=1)
        phase_feats = []
        MHA1 = MultiHeadAttention(4, 16, attention_axes=1)
        Mean1 = Lambda(lambda x: K.mean(x, axis=1, keepdims=True))
        for i in range(self.num_phases):
            tmp_feat_1 = tf.concat([lane_feats[idx] for idx in self.phase_map[i]], axis=1)
            tmp_feat_2 = MHA1(tmp_feat_1, tmp_feat_1)
            tmp_feat_3 = Mean1(tmp_feat_2)
            phase_feats.append(tmp_feat_3)
        phase_feat_all = tf.concat(phase_feats, axis=1)
        att_encoding = MultiHeadAttention(4, 8, attention_axes=1)(phase_feat_all, phase_feat_all)
        hidden = Dense(20, activation="relu")(att_encoding)
        hidden = Dense(20, activation="relu")(hidden)
        phase_feature_final = Dense(1, activation="linear", name="beformerge")(hidden)
        pscore = Reshape((4,))(phase_feature_final)
        selected_phase_feat = Lambda(lambda x: tf.matmul(x[0], x[1]))([ins2, phase_feat_all])
        selected_phase_feat = Reshape((dims2, ))(selected_phase_feat)
        hidden2 = Dense(20, activation="relu")(selected_phase_feat)
        hidden2 = Dense(20, activation="relu")(hidden2)
        dscore = self.dueling_block(hidden2)
        network = Model(inputs=[ins0, ins1, ins2],
                        outputs=[pscore, dscore])
        network.compile(optimizer=Adam(lr=self.dic_agent_conf["LEARNING_RATE"], epsilon=1e-08),
                        loss=self.dic_agent_conf["LOSS_FUNCTION"])
        network.summary()
        return network
    
    def freze_layers(self):
        for layer in self.q_network.layers[0:21]:
            layer.trainable = False
        self.q_network.compile(optimizer=Adam(lr=self.dic_agent_conf["LEARNING_RATE"], epsilon=1e-08),
                               loss=self.dic_agent_conf["LOSS_FUNCTION"])

    def unfreze_layers(self):
        for layer in self.q_network.layers[0:21]:
            layer.trainable = True
        self.q_network.compile(optimizer=Adam(lr=self.dic_agent_conf["LEARNING_RATE2"], epsilon=1e-08),
                               loss=self.dic_agent_conf["LOSS_FUNCTION"])

    def resetlr(self):
        self.q_network.compile(optimizer=Adam(lr=self.dic_agent_conf["LEARNING_RATE3"], epsilon=1e-08),
                               loss=self.dic_agent_conf["LOSS_FUNCTION"])
    
    def train_network(self):
        epochs = self.dic_agent_conf["EPOCHS"]
        val_spl = 0.2
        early_stopping = EarlyStopping(
            monitor='val_loss', patience=self.dic_agent_conf["PATIENCE"], verbose=0, mode='min')
        batch_size = min(self.dic_agent_conf["BATCH_SIZE1"], len(self.Ys[0]))
        self.q_network.fit(self.Xs, self.Ys, batch_size=batch_size, epochs=epochs, shuffle=1,
                           verbose=2, validation_split=val_spl, callbacks=[early_stopping])
        print("=== training two model ===")

    def dueling_block(self, inputs):
        tmp_v = Dense(20, activation="relu", name="dense_values")(inputs)
        value = Dense(1, activation="linear", name="dueling_values")(tmp_v)
        tmp_a = Dense(20, activation="relu", name="dense_a")(inputs)
        a = Dense(self.num_action_dur, activation="linear", name="dueling_advantages")(tmp_a)
        mean = Lambda(lambda x: K.mean(x, axis=1, keepdims=True))(a)
        advantages = Subtract()([a, mean])
        q_values = Add(name='dueling_q_values')([value, advantages])
        return q_values
    
    def choose_action(self, states, list_need):
        dic_state_feature_arrays = {}
        used_feature = copy.deepcopy(self.dic_traffic_env_conf["LIST_STATE_FEATURE_1"])
        cur_phase = []
        for feature_name in used_feature:
            dic_state_feature_arrays[feature_name] = []
        for s in states:
            for feature_name in used_feature:
                if feature_name == "phase_total":
                    cur_phase.append(s[feature_name])
                else:
                    dic_state_feature_arrays[feature_name].append(s[feature_name])
        used_feature.remove("phase_total")
        state_input = [np.array(dic_state_feature_arrays[feature_name]).reshape(len(states), self.max_lane, -1) for feature_name in
                       used_feature]
        state_input = np.concatenate(state_input, axis=-1)
        
        # phase action 
        tmp_p, _ = self.q_network.predict([state_input, np.array(cur_phase), np.random.rand(len(states),1, 4)])
        if self.cnt_round < self.dic_agent_conf["SROUND"]:
            paction = self.epsilon_choice_one(tmp_p)
        else:
            paction = np.argmax(tmp_p, axis=1)
        # duration action
        phase_idx =  np.array(paction).reshape(len(paction), 1, 1)
        phase_matrix = self.phase_index2matrix(phase_idx)
        _, tmp_d = self.q_network.predict([state_input, np.array(cur_phase), phase_matrix])
        if self.cnt_round < self.dic_agent_conf["SROUND"]:
            daction = [1] * len(states)
        elif self.cnt_round < self.dic_agent_conf["SROUND2"]:
            daction = self.epsilon_choice_two(tmp_d)
        else:
            daction = np.argmax(tmp_d, axis=1)
        return paction, daction

    def phase_index2matrix(self, phase_index):
        # [batch, 1] -> [batch, 1, num_phase]
        lab = to_categorical(phase_index, num_classes=self.num_phases)
        return lab
    
    def epsilon_choice_one(self, q_values):
        max_1 = np.expand_dims(np.argmax(q_values, axis=-1), axis=-1)
        rand_1 = np.random.randint(self.num_phases, size=(len(q_values), 1))
        _p = np.concatenate([max_1, rand_1], axis=-1)
        select = np.random.choice([0, 1], size=len(q_values), p=[1 - self.dic_agent_conf["EPSILON"],
                                                                 self.dic_agent_conf["EPSILON"]])
        act = _p[np.arange(len(q_values)), select]
        return act
    
    def epsilon_choice_two(self, q_values):
        max_1 = np.expand_dims(np.argmax(q_values, axis=-1), axis=-1)
        rand_1 = np.random.randint(self.num_action_dur, size=(len(q_values), 1))
        _p = np.concatenate([max_1, rand_1], axis=-1)
        select = np.random.choice([0, 1], size=len(q_values), p=[1 - self.dic_agent_conf["EPSILON"],
                                                                 self.dic_agent_conf["EPSILON"]])
        act = _p[np.arange(len(q_values)), select]
        return act
    
    def prepare_Xs_Y(self, memory):
        """ used for update phase control model """
        ind_end = len(memory)
        print("memory size before forget: {0}".format(ind_end))
        # use all the samples to pretrain, i.e., without forgetting
        ind_sta = max(0, ind_end - self.dic_agent_conf["MAX_MEMORY_LEN"])
        memory_after_forget = memory[ind_sta: ind_end]
        print("memory size after forget:", len(memory_after_forget))
        # sample the memory
        sample_size = min(self.dic_agent_conf["SAMPLE_SIZE1"], len(memory_after_forget))
        sample_slice = random.sample(memory_after_forget, sample_size)
        print("memory samples number:", sample_size)
        used_feature = self.dic_traffic_env_conf["LIST_STATE_FEATURE_1"][1:]
        _phase1, _phase2 = [], []
        _state = [[] for _ in used_feature]
        _next_state = [[] for _ in used_feature]
        _action1 = []
        _action2 = []
        _action3 = []
        _reward = []
        for i in range(len(sample_slice)):
            state, action1, action2, next_state, reward, _ = sample_slice[i]
            for feat_idx, feat_name in enumerate(used_feature):
                _state[feat_idx].append(state[feat_name])
                _next_state[feat_idx].append(next_state[feat_name])
            _phase1.append(state["phase_total"])
            _phase2.append(next_state["phase_total"])
            _action1.append([[action1]])
            _action3.append(action1)
            _action2.append(action2)
            _reward.append(reward)
        _state2 = np.concatenate([np.array(ss).reshape(len(ss), self.max_lane, -1) for ss in _state], axis=-1)
        _next_state2 = np.concatenate([np.array(ss).reshape(len(ss), self.max_lane, -1) for ss in _next_state], axis=-1)
        phase_matrix = self.phase_index2matrix(np.array(_action1))
        
        cur_p, cur_d = self.q_network.predict([_state2, np.array(_phase1), phase_matrix])
        next_p, next_d = self.q_network_bar.predict([_next_state2, np.array(_phase2), phase_matrix])
        
        target1 = np.copy(cur_p)
        target2 = np.copy(cur_d)
        if self.cnt_round < self.dic_agent_conf["SROUND"]:
            for i in range(len(sample_slice)):
                target1[i, _action3[i]] = _reward[i] / self.dic_agent_conf["NORMAL_FACTOR"] + self.dic_agent_conf["GAMMA"] * np.max(next_p[i, :])
        elif self.cnt_round < self.dic_agent_conf["SROUND2"]:
            for i in range(len(sample_slice)):
                target2[i, _action2[i]] = _reward[i] / self.dic_agent_conf["NORMAL_FACTOR"] + self.dic_agent_conf["GAMMA"] * np.max(next_d[i, :])
        else:
            for i in range(len(sample_slice)):
                target1[i, _action3[i]] = _reward[i] / self.dic_agent_conf["NORMAL_FACTOR"] + self.dic_agent_conf["GAMMA"] *np.max(next_p[i, :])
            for i in range(len(sample_slice)):
                target2[i, _action2[i]] = _reward[i] / self.dic_agent_conf["NORMAL_FACTOR"] + self.dic_agent_conf["GAMMA"] * np.max(next_d[i, :])
            
        self.Xs = [_state2, np.array(_phase1), phase_matrix]
        self.Ys = [target1, target2]

    def save_network(self, file_name):
        if self.cnt_round == self.dic_agent_conf["SROUND"]-1:
            self.freze_layers()
        if self.cnt_round == self.dic_agent_conf["SROUND2"]-1:
            self.unfreze_layers()
        if self.cnt_round == self.dic_agent_conf["SROUND3"]-1:
            self.resetlr()
        self.q_network.save(os.path.join(self.dic_path["PATH_TO_MODEL"], "%s.h5" % file_name))"

network:"from tensorflow.keras.models import model_from_json, load_model
from tensorflow.keras.optimizers import Adam
import os
from .agent import Agent
import traceback


class NetworkAgent(Agent):
    def __init__(self, dic_agent_conf, dic_traffic_env_conf, dic_path, cnt_round, intersection_id="0"):
        super(NetworkAgent, self).__init__(
            dic_agent_conf, dic_traffic_env_conf, dic_path)
        # ===== check num actions == num phases ============
        self.num_actions = len(dic_traffic_env_conf["PHASE"])
        self.num_phases = len(dic_traffic_env_conf["PHASE"])
        self.num_action_dur = len(dic_traffic_env_conf["ACTION_DURATION"])
        self.memory = self.build_memory()
        self.cnt_round = cnt_round

        self.num_intersections = dic_traffic_env_conf["NUM_INTERSECTIONS"]

        self.cyclicInd = [[0] * self.num_phases for _ in range(self.num_intersections)]
        self.cyclicInd2 = [0] * self.num_intersections

        self.Xs, self.Ys, self.Xs_one, self.Y_one = None, None, None, None

        self.num_lane = dic_traffic_env_conf["NUM_LANE"]
        self.max_lane = dic_traffic_env_conf["MAX_LANE"]
        self.phase_map = dic_traffic_env_conf["PHASE_MAP"]
        len_feat1, len_feat2 = self.cal_input_len()
        self.num_feat1 = int(len_feat1/self.max_lane)
        self.num_feat2 = int(len_feat2/self.max_lane)

        if cnt_round == 0:
            if os.listdir(self.dic_path["PATH_TO_MODEL"]):
                self.load_network("round_0_inter_{0}".format(intersection_id))
            else:
                self.q_network = self.build_network()
            self.q_network_bar = self.build_network_from_copy(self.q_network)
        else:
            try:
                self.load_network("round_{0}_inter_{1}".format(cnt_round-1, self.intersection_id))

                if "UPDATE_Q_BAR_EVERY_C_ROUND" in self.dic_agent_conf:
                    if self.dic_agent_conf["UPDATE_Q_BAR_EVERY_C_ROUND"]:
                        self.load_network_bar("round_{0}_inter_{1}".format(
                                max((cnt_round - 1) // self.dic_agent_conf["UPDATE_Q_BAR_FREQ"] *
                                    self.dic_agent_conf["UPDATE_Q_BAR_FREQ"], 0), self.intersection_id))
                    else:
                        self.load_network_bar("round_{0}_inter_{1}".format(
                            max(cnt_round - self.dic_agent_conf["UPDATE_Q_BAR_FREQ"], 0),
                            self.intersection_id))
                else:
                    self.load_network_bar("round_{0}_inter_{1}".format(
                            max(cnt_round - self.dic_agent_conf["UPDATE_Q_BAR_FREQ"], 0), self.intersection_id))
            except Exception:
                print('traceback.format_exc():\n%s' % traceback.format_exc())

        # decay the epsilon
        decayed_epsilon = self.dic_agent_conf["EPSILON"] * pow(self.dic_agent_conf["EPSILON_DECAY"], cnt_round)
        self.dic_agent_conf["EPSILON"] = max(decayed_epsilon, self.dic_agent_conf["MIN_EPSILON"])
    
    def cal_input_len(self):
        N1, N2 = 0, 0
        used_feature1 = self.dic_traffic_env_conf["LIST_STATE_FEATURE_1"]
        used_feature2 = self.dic_traffic_env_conf["LIST_STATE_FEATURE_2"]
        for feat_name1 in used_feature1:
            if "num_in_deg" in feat_name1:
                N1 += self.max_lane*4
            elif "phase_total" in feat_name1:
                N1 += 0
            else:
                N1 += self.max_lane
        for feat_name2 in used_feature2:
            if "num_in_deg" in feat_name2:
                N2 += self.max_lane*4
            elif "phase_total" in feat_name2:
                N2 += 0
            else:
                N2 += self.max_lane
        return N1, N2
    
    def load_network(self, file_name, file_path=None):
        if file_path is None:
            file_path = self.dic_path["PATH_TO_MODEL"]
        self.q_network = load_model(os.path.join(file_path,  "%s.h5" % file_name))
        print("succeed in loading model %s" % file_name)

    def load_network_bar(self, file_name, file_path=None):
        if file_path is None:
            file_path = self.dic_path["PATH_TO_MODEL"]
        self.q_network_bar = load_model(os.path.join(file_path,  "%s.h5" % file_name))
        print("succeed in loading model %s" % file_name)

    def save_network(self, file_name):
        self.q_network.save(os.path.join(self.dic_path["PATH_TO_MODEL"], "%s.h5" % file_name))

    @staticmethod
    def build_memory():
        return []

    def build_network_from_copy(self, network_copy):
        """Initialize a Q network from a copy"""
        network_structure = network_copy.to_json()
        network_weights = network_copy.get_weights()
        network = model_from_json(network_structure)
        network.set_weights(network_weights)
        network.compile(optimizer=Adam(lr=self.dic_agent_conf["LEARNING_RATE"]),
                        loss=self.dic_agent_conf["LOSS_FUNCTION"])
        return network

    def train_network(self):
        pass
"

""""
Specialized environment to corresponding to dynamic phase duration
"""
import pickle
import numpy as np
import json
import sys
import pandas as pd
import os
import cityflow as engine
import time
from multiprocessing import Process


class Intersection:
    def __init__(self, inter_id, dic_traffic_env_conf, eng, light_id_dict, path_to_log, lanes_length_dict):
        self.inter_id = inter_id
        self.inter_name = "intersection_{0}_{1}".format(inter_id[0], inter_id[1])
        self.eng = eng
        self.dic_traffic_env_conf = dic_traffic_env_conf
        self.lane_length = lanes_length_dict
        self.obs_length = dic_traffic_env_conf["OBS_LENGTH"]
        # newl add one obs_length for queue vehicle to realize precise observation
#         self.obs_length_q = dic_traffic_env_conf["OBS_LENGTH_Q"]
        self.num_lane = dic_traffic_env_conf["NUM_LANE"]
        self.max_lane = dic_traffic_env_conf["MAX_LANE"]
        self.padding = self.num_lane < self.max_lane
        self.padding1 = [0 for _ in range((self.max_lane-self.num_lane) * 4)]
        self.padding2 = [0 for _ in range(self.max_lane-self.num_lane)]

        self.list_approachs = ["W", "E", "N", "S"]
        # corresponding exiting lane for entering lanes
        self.dic_approach_to_node = {"W": 0, "E": 2, "S": 1, "N": 3}
        self.dic_entering_approach_to_edge = {"W": "road_{0}_{1}_0".format(inter_id[0] - 1, inter_id[1])}
        self.dic_entering_approach_to_edge.update({"E": "road_{0}_{1}_2".format(inter_id[0] + 1, inter_id[1])})
        self.dic_entering_approach_to_edge.update({"N": "road_{0}_{1}_3".format(inter_id[0], inter_id[1] + 1)})
        self.dic_entering_approach_to_edge.update({"S": "road_{0}_{1}_1".format(inter_id[0], inter_id[1] - 1)})
        self.dic_exiting_approach_to_edge = {
            approach: "road_{0}_{1}_{2}".format(inter_id[0], inter_id[1], self.dic_approach_to_node[approach]) for
            approach in self.list_approachs}
        self.list_phases = dic_traffic_env_conf["PHASE"]

        # generate all lanes
        self.list_entering_lanes = []
        for (approach, lane_number) in zip(self.list_approachs, dic_traffic_env_conf["NUM_LANES"]):
            self.list_entering_lanes += [self.dic_entering_approach_to_edge[approach] + "_" + str(i) for i in
                                         range(lane_number)]
        self.list_exiting_lanes = []
        for (approach, lane_number) in zip(self.list_approachs, dic_traffic_env_conf["NUM_LANES"]):
            self.list_exiting_lanes += [self.dic_exiting_approach_to_edge[approach] + "_" + str(i) for i in
                                        range(lane_number)]

        self.list_lanes = self.list_entering_lanes + self.list_exiting_lanes

        self.adjacency_row = light_id_dict["adjacency_row"]
        self.neighbor_ENWS = light_id_dict["neighbor_ENWS"]

        # ========== record previous & current feats ==========
        self.dic_lane_vehicle_previous_step = {}
        self.dic_lane_vehicle_previous_step_in = {}
        self.dic_lane_waiting_vehicle_count_previous_step = {}
        self.dic_vehicle_speed_previous_step = {}
        self.dic_vehicle_distance_previous_step = {}

        # in [entering_lanes] out [exiting_lanes]
        self.dic_lane_vehicle_current_step_in = {}
        self.dic_lane_vehicle_current_step = {}
        self.dic_lane_waiting_vehicle_count_current_step = {}
        self.dic_vehicle_speed_current_step = {}
        self.dic_vehicle_distance_current_step = {}

        self.list_lane_vehicle_previous_step_in = []
        self.list_lane_vehicle_current_step_in = []

        self.dic_vehicle_arrive_leave_time = dict()  # cumulative

        self.dic_feature = {}  # this second
        self.dic_feature_previous_step = {}  # this second

        # =========== signal info set ================
        # -1: all yellow, -2: all red, -3: none
        self.all_yellow_phase_index = -1
        self.all_red_phase_index = -2

        self.current_phase_index = 1
        self.previous_phase_index = 1
        self.eng.set_tl_phase(self.inter_name, self.current_phase_index)
        path_to_log_file = os.path.join(path_to_log, "signal_inter_{0}.txt".format(self.inter_name))
        df = [self.get_current_time(), self.current_phase_index]
        df = pd.DataFrame(df)
        df = df.transpose()
        df.to_csv(path_to_log_file, mode="a", header=False, index=False)

        self.next_phase_to_set_index = None
        self.current_phase_duration = -1
        self.all_red_flag = False
        self.all_yellow_flag = False
        self.flicker = 0

    def set_signal(self, action, action_pattern, yellow_time, path_to_log):
        if self.all_yellow_flag:
            # in yellow phase
            self.flicker = 0
            if self.current_phase_duration >= yellow_time:  # yellow time reached
                self.current_phase_index = self.next_phase_to_set_index
                self.eng.set_tl_phase(self.inter_name, self.current_phase_index)  # if multi_phase, need more adjustment
                path_to_log_file = os.path.join(path_to_log, "signal_inter_{0}.txt".format(self.inter_name))
                df = [self.get_current_time(), self.current_phase_index]
                df = pd.DataFrame(df)
                df = df.transpose()
                df.to_csv(path_to_log_file, mode="a", header=False, index=False)
                self.all_yellow_flag = False
        else:
            # determine phase
            if action_pattern == "switch":  # switch by order
                if action == 0:  # keep the phase
                    self.next_phase_to_set_index = self.current_phase_index
                elif action == 1:  # change to the next phase
                    self.next_phase_to_set_index = (self.current_phase_index + 1) % len(self.list_phases)
                    # if multi_phase, need more adjustment
                else:
                    sys.exit("action not recognized\n action must be 0 or 1")

            elif action_pattern == "set":  # set to certain phase
                # self.next_phase_to_set_index = self.DIC_PHASE_MAP[action] # if multi_phase, need more adjustment
                self.next_phase_to_set_index = action + 1
            # set phase
            if self.current_phase_index == self.next_phase_to_set_index:
                # the light phase keeps unchanged
                pass
            else:  # the light phase needs to change
                # change to yellow first, and activate the counter and flag
                self.eng.set_tl_phase(self.inter_name, 0)  # !!! yellow, tmp
                path_to_log_file = os.path.join(path_to_log, "signal_inter_{0}.txt".format(self.inter_name))
                df = [self.get_current_time(), self.current_phase_index]
                df = pd.DataFrame(df)
                df = df.transpose()
                df.to_csv(path_to_log_file, mode="a", header=False, index=False)
                self.current_phase_index = self.all_yellow_phase_index
                self.all_yellow_flag = True
                self.flicker = 1

    # update inner measurements
    def update_previous_measurements(self):
        self.previous_phase_index = self.current_phase_index
        self.dic_lane_vehicle_previous_step = self.dic_lane_vehicle_current_step
        self.dic_lane_vehicle_previous_step_in = self.dic_lane_vehicle_current_step_in
        self.dic_lane_waiting_vehicle_count_previous_step = self.dic_lane_waiting_vehicle_count_current_step
        self.dic_vehicle_speed_previous_step = self.dic_vehicle_speed_current_step
        self.dic_vehicle_distance_previous_step = self.dic_vehicle_distance_current_step

    def update_current_measurements(self, simulator_state):
        def _change_lane_vehicle_dic_to_list(dic_lane_vehicle):
            list_lane_vehicle = []
            for value in dic_lane_vehicle.values():
                list_lane_vehicle.extend(value)
            return list_lane_vehicle

        if self.current_phase_index == self.previous_phase_index:
            self.current_phase_duration += 1
        else:
            self.current_phase_duration = 1

        self.dic_lane_vehicle_current_step = {}
        self.dic_lane_vehicle_current_step_in = {}
        self.dic_lane_waiting_vehicle_count_current_step = {}
        for lane in self.list_entering_lanes:
            self.dic_lane_vehicle_current_step_in[lane] = simulator_state["get_lane_vehicles"][lane]

        for lane in self.list_lanes:
            self.dic_lane_vehicle_current_step[lane] = simulator_state["get_lane_vehicles"][lane]
            self.dic_lane_waiting_vehicle_count_current_step[lane] = simulator_state["get_lane_waiting_vehicle_count"][lane]

        self.dic_vehicle_speed_current_step = simulator_state["get_vehicle_speed"]
        self.dic_vehicle_distance_current_step = simulator_state["get_vehicle_distance"]

        # get vehicle list
        self.list_lane_vehicle_current_step_in = _change_lane_vehicle_dic_to_list(self.dic_lane_vehicle_current_step_in)
        self.list_lane_vehicle_previous_step_in = _change_lane_vehicle_dic_to_list(self.dic_lane_vehicle_previous_step_in)

        list_vehicle_new_arrive = list(set(self.list_lane_vehicle_current_step_in) - set(self.list_lane_vehicle_previous_step_in))
        # can't use empty set to - real set
        if not self.list_lane_vehicle_previous_step_in:  # previous step is empty
            list_vehicle_new_left = list(set(self.list_lane_vehicle_current_step_in) -
                                         set(self.list_lane_vehicle_previous_step_in))
        else:
            list_vehicle_new_left = list(set(self.list_lane_vehicle_previous_step_in) -
                                         set(self.list_lane_vehicle_current_step_in))
        # update vehicle arrive and left time
        self._update_arrive_time(list_vehicle_new_arrive)
        self._update_left_time(list_vehicle_new_left)
        # update feature
        self._update_feature()

    def _update_leave_entering_approach_vehicle(self):
        list_entering_lane_vehicle_left = []
        # update vehicles leaving entering lane
        if not self.dic_lane_vehicle_previous_step:  # the dict is not empty
            for _ in self.list_entering_lanes:
                list_entering_lane_vehicle_left.append([])
        else:
            last_step_vehicle_id_list = []
            current_step_vehilce_id_list = []
            for lane in self.list_entering_lanes:
                last_step_vehicle_id_list.extend(self.dic_lane_vehicle_previous_step[lane])
                current_step_vehilce_id_list.extend(self.dic_lane_vehicle_current_step[lane])

            list_entering_lane_vehicle_left.append(
                list(set(last_step_vehicle_id_list) - set(current_step_vehilce_id_list))
            )
        return list_entering_lane_vehicle_left

    def _update_arrive_time(self, list_vehicle_arrive):
        ts = self.get_current_time()
        # get dic vehicle enter leave time
        for vehicle in list_vehicle_arrive:
            if vehicle not in self.dic_vehicle_arrive_leave_time:
                self.dic_vehicle_arrive_leave_time[vehicle] = {"enter_time": ts, "leave_time": np.nan}

    def _update_left_time(self, list_vehicle_left):
        ts = self.get_current_time()
        # update the time for vehicle to leave entering lane
        for vehicle in list_vehicle_left:
            try:
                self.dic_vehicle_arrive_leave_time[vehicle]["leave_time"] = ts
            except KeyError:
                print("vehicle not recorded when entering")
                sys.exit(-1)

    def _update_feature(self):
        dic_feature = dict()
        dic_feature["cur_phase"] = [self.current_phase_index]
#         dic_feature["time_this_phase"] = [self.current_phase_duration]
    
        if self.current_phase_index > 0:
            dic_feature["phase_total"] = self._get_total_phase_infos(self.current_phase_index)
    
        dic_feature["lane_num_vehicle_in"] = self._get_lane_num_vehicles(self.list_entering_lanes)
        dic_feature["lane_num_vehicle_out"] = self._get_lane_num_vehicles(self.list_exiting_lanes)

        dic_feature["lane_queue_vehicle_in"] = self._get_lane_queue_length(self.list_entering_lanes)
        dic_feature["lane_queue_vehicle_out"] = self._get_lane_queue_length(self.list_exiting_lanes)
        dic_feature["traffic_movement_pressure"] = self._get_traffic_movement_pressure_general(
            dic_feature["lane_queue_vehicle_in"],
            dic_feature["lane_queue_vehicle_out"])
        
        _, _, _, _, lane_run_in_part, _ = self._get_part_observations_first()
        # -------- run vehicle --------------
        dic_feature["lane_run_in_part"] = self._get_effective_running(lane_run_in_part)
        
        dic_feature["num_in_deg"] = self._orgnize_several_segments2()  # [None, 12*4]
#         dic_feature["num_out_deg"] = self._orgnize_several_segments3()  # [None, 12*4]
        # -------- reward------------------
        dic_feature["pressure"] = self._get_pressure(dic_feature["lane_queue_vehicle_in"],
                                                     dic_feature["lane_queue_vehicle_out"])
        dic_feature["adjacency_matrix"] = self._get_adjacency_row()

        self.dic_feature = dic_feature
    
    def _get_effective_running(self, lane_run_in_part):
        res = lane_run_in_part
        if self.padding:
            res = res + self.padding2
        return res
    
    
    def _get_total_phase_infos(self, current_phase_index):
#         print(self.dic_traffic_env_conf['PHASETOTAL'])
        res =  self.dic_traffic_env_conf['PHASETOTAL'][current_phase_index]
        if self.padding:
            res = res + self.padding2
        return res
    
    # get the part observation with the first length
    def _get_part_observations_first(self):
        """
        return: lane_num_in_part
                lane_num_out_part
                lane_queue_in_part
                lane_queue_out_part
                lane_run_in_part
        """
        f_p_num, l_p_num, l_p_q = self._get_part_observations(lane_vehicles=self.dic_lane_vehicle_current_step,
                                                              vehicle_distance=self.dic_vehicle_distance_current_step,
                                                              vehicle_speed=self.dic_vehicle_speed_current_step,
                                                              lane_length=self.lane_length,
                                                              obs_length=self.obs_length,
                                                              list_lanes=self.list_lanes)
        # queue vehicles
        lane_queue_in_part = [len(l_p_q[lane]) for lane in self.list_entering_lanes]
        lane_queue_out_part = [len(l_p_q[lane]) for lane in self.list_exiting_lanes]

        # num vehicles [first part and last part]
        # last part
        lane_num_in_part_l = [len(l_p_num[lane]) for lane in self.list_entering_lanes]
        lane_num_out_part_l = [len(l_p_num[lane]) for lane in self.list_exiting_lanes]

        # first part
        lane_num_in_part_f = [len(f_p_num[lane]) for lane in self.list_entering_lanes]
        lane_num_out_part_f = [len(f_p_num[lane]) for lane in self.list_exiting_lanes]

        # lane part total
        lane_num_in_part_total = list(np.array(lane_num_in_part_f)+np.array(lane_num_in_part_l))
        lane_num_out_part_total = list(np.array(lane_num_out_part_f) + np.array(lane_num_out_part_l))

        # running vehicles
        lane_run_in_part = list(np.array(lane_num_in_part_l) - np.array(lane_queue_in_part))
        return lane_queue_in_part, lane_queue_out_part, lane_num_in_part_total, lane_num_out_part_total, lane_run_in_part, lane_num_in_part_l

    @staticmethod
    def _get_part_observations(lane_vehicles, vehicle_distance, vehicle_speed,
                               lane_length, obs_length, list_lanes):
        """
            Input: lane_vehicles :      Dict{lane_id    :   [vehicle_ids]}
                   vehicle_distance:    Dict{vehicle_id :   float(dist)}
                   vehicle_speed:       Dict{vehicle_id :   float(speed)}
                   lane_length  :       Dict{lane_id    :   float(length)}
                   obs_length   :       The part observation length
                   list_lanes   :       List[lane_ids at the intersection]
        :return:
                    part_vehicles:      Dict{ lane_id, [vehicle_ids]}
        """
        # get vehicle_ids and speeds
        first_part_num_vehicle = {}
        first_part_queue_vehicle = {}  # useless, at the begin of lane, there is no waiting vechiles
        last_part_num_vehicle = {}
        last_part_queue_vehicle = {}

        for lane in list_lanes:
            first_part_num_vehicle[lane] = []
            first_part_queue_vehicle[lane] = []
            last_part_num_vehicle[lane] = []
            last_part_queue_vehicle[lane] = []
            last_part_obs_length = lane_length[lane] - obs_length
            for vehicle in lane_vehicles[lane]:
                """ get the first part of obs
                    That is vehicle_distance <= obs_length 
                """
                # set as num_vehicle
                if "shadow" in vehicle:  # remove the shadow
                    vehicle = vehicle[:-7]
                temp_v_distance = vehicle_distance[vehicle]
                if temp_v_distance <= obs_length:
                    first_part_num_vehicle[lane].append(vehicle)
                    # analyse if waiting
                    if vehicle_speed[vehicle] <= 0.1:
                        first_part_queue_vehicle[lane].append(vehicle)

                """ get the last part of obs
                    That is  lane_length-obs_length <= vehicle_distance <= lane_length 
                """
                if temp_v_distance >= last_part_obs_length:
                    last_part_num_vehicle[lane].append(vehicle)
                    # analyse if waiting
                    if vehicle_speed[vehicle] <= 0.1:
                        last_part_queue_vehicle[lane].append(vehicle)

        return first_part_num_vehicle, last_part_num_vehicle, last_part_queue_vehicle

    
    
    def _get_traffic_movement_pressure_efficient(self, enterings, exitings):

        list_approachs = ["W", "E", "N", "S"]
        if self.num_lane == 8:
            index_maps = {
                "W": [0, 1],
                "E": [2, 3],
                "N": [4, 5],
                "S": [6, 7],
                "WN": [0, 1, 4, 5],
                "SW": [0, 1, 6, 7],
                "ES": [2, 3, 6, 7],
                "NE": [2, 3, 4, 5]

            }
            turn_maps = ["S", "WN",
                         "N", "ES",
                         "W", "NE",
                         "E", "SW"]
        elif self.num_lane == 10:
            index_maps = {
                "W": [0, 1, 2],
                "E": [3, 4, 5],
                "N": [6, 7],
                "S": [8, 9],
                "NE": [6, 7, 3, 4, 5],
                "SW": [8, 9, 0, 1, 2]
            }
            turn_maps = ["S", "W", "N",
                         "N", "E", "S",
                         "W", "NE",
                         "E", "SW"]
        elif self.num_lane == 12:
            index_maps = {
                "W": [0, 1, 2],
                "E": [3, 4, 5],
                "N": [6, 7, 8],
                "S": [9, 10, 11]
            }
            turn_maps = ["S", "W", "N",
                         "N", "E", "S",
                         "W", "N", "E",
                         "E", "S", "W"]
        elif self.num_lane == 16:
            index_maps = {
                "W": [0, 1, 2, 3],
                "E": [4, 5, 6, 7],
                "N": [8, 9, 10, 11],
                "S": [12, 13, 14, 15]
            }
            turn_maps = ["S", "W", "W", "N",
                         "N", "E", "E", "S",
                         "W", "N", "N", "E",
                         "E", "S", "S", "W"]

        # vehicles in exiting road
        outs_maps = {}
        for approach in index_maps.keys():
            outs_maps[approach] = np.mean([exitings[i] for i in index_maps[approach]])
        # turn_maps = ["S", "W", "N", "N", "E", "S", "W", "N", "E", "E", "S", "W"]
        t_m_p = [enterings[j] - outs_maps[turn_maps[j]] for j in range(self.num_lane)]
        if self.padding:
            t_m_p = t_m_p + self.padding2

        return t_m_p
    
    def _get_traffic_movement_pressure_general(self, enterings, exitings):

        list_approachs = ["W", "E", "N", "S"]
        if self.num_lane == 8:
            index_maps = {
                "W": [0, 1],
                "E": [2, 3],
                "N": [4, 5],
                "S": [6, 7],
                "WN": [0, 1, 4, 5],
                "SW": [0, 1, 6, 7],
                "ES": [2, 3, 6, 7],
                "NE": [2, 3, 4, 5]

            }
            turn_maps = ["S", "WN",
                         "N", "ES",
                         "W", "NE",
                         "E", "SW"]
        elif self.num_lane == 10:
            index_maps = {
                "W": [0, 1, 2],
                "E": [3, 4, 5],
                "N": [6, 7],
                "S": [8, 9],
                "NE": [6, 7, 3, 4, 5],
                "SW": [8, 9, 0, 1, 2]
            }
            turn_maps = ["S", "W", "N",
                         "N", "E", "S",
                         "W", "NE",
                         "E", "SW"]
        elif self.num_lane == 12:
            index_maps = {
                "W": [0, 1, 2],
                "E": [3, 4, 5],
                "N": [6, 7, 8],
                "S": [9, 10, 11]
            }
            turn_maps = ["S", "W", "N",
                         "N", "E", "S",
                         "W", "N", "E",
                         "E", "S", "W"]
        elif self.num_lane == 16:
            index_maps = {
                "W": [0, 1, 2, 3],
                "E": [4, 5, 6, 7],
                "N": [8, 9, 10, 11],
                "S": [12, 13, 14, 15]
            }
            turn_maps = ["S", "W", "W", "N",
                         "N", "E", "E", "S",
                         "W", "N", "N", "E",
                         "E", "S", "S", "W"]

        # vehicles in exiting road
        outs_maps = {}
        for approach in index_maps.keys():
            outs_maps[approach] = sum([exitings[i] for i in index_maps[approach]])
        if self.num_lane == 16:
            t_m_p = []
            for i in range(self.num_lane):
                if i in [0, 3, 4, 7, 8, 11, 12, 15]:
                    t_m_p.append(enterings[i]-outs_maps[turn_maps[i]])
                else:
                    t_m_p.append(enterings[i] - outs_maps[turn_maps[i]]/2)
        else:
            t_m_p = [enterings[j] - outs_maps[turn_maps[j]] for j in range(self.num_lane)]

        if self.padding:
            t_m_p = t_m_p + self.padding2

        return t_m_p

    
    
    def _orgnize_several_segments2(self):
        part1, part2, part3, part4 = self._get_several_segments(lane_vehicles=self.dic_lane_vehicle_current_step,
                                                                vehicle_distance=self.dic_vehicle_distance_current_step,
                                                                vehicle_speed=self.dic_vehicle_speed_current_step,
                                                                lane_length=self.lane_length,
                                                                list_lanes=self.list_lanes)
        num_in_part1 = [len(part1[lane]) for lane in self.list_entering_lanes]
        num_in_part2 = [len(part2[lane]) for lane in self.list_entering_lanes]
        num_in_part3 = [len(part3[lane]) for lane in self.list_entering_lanes]
        num_in_part4 = [len(part4[lane]) for lane in self.list_entering_lanes]

        total_in = []
        for i in range(len(self.list_entering_lanes)):
            total_in.extend([num_in_part1[i], num_in_part2[i], num_in_part3[i], num_in_part4[i]])
        if self.padding:
            total_in = total_in + self.padding1
        return total_in

    def _get_several_segments(self, lane_vehicles, vehicle_distance, vehicle_speed,
                              lane_length, list_lanes):
        # get four segments [100, 200, 300, 400] for segment
        obs_length = 100
        part1, part2, part3, part4 = {}, {}, {}, {}
        for lane in list_lanes:
            part1[lane], part2[lane], part3[lane], part4[lane] = [], [], [], []
            for vehicle in lane_vehicles[lane]:
                # set as num_vehicle
                if "shadow" in vehicle:  # remove the shadow
                    vehicle = vehicle[:-7]
                    continue
                temp_v_distance = vehicle_distance[vehicle]
                if temp_v_distance > lane_length[lane] - obs_length:
                    part1[lane].append(vehicle)
                    #  running vehicles
                    # if vehicle_speed[vehicle] > 0.1:
                    #     part1[lane].append(vehicle)
                elif lane_length[lane] - 2 * obs_length < temp_v_distance <= lane_length[lane] - obs_length:
                    part2[lane].append(vehicle)
                elif lane_length[lane] - 3 * obs_length < temp_v_distance <= lane_length[lane] - 2 * obs_length:
                    part3[lane].append(vehicle)
                elif lane_length[lane] - 4 * obs_length < temp_v_distance <= lane_length[lane] - 3 * obs_length:
                    part4[lane].append(vehicle)
        return part1, part2, part3, part4
    
    def _orgnize_several_segments3(self):
        part1, part2, part3, part4 = self._get_several_segments2(lane_vehicles=self.dic_lane_vehicle_current_step,
                                                                vehicle_distance=self.dic_vehicle_distance_current_step,
                                                                vehicle_speed=self.dic_vehicle_speed_current_step,
                                                                lane_length=self.lane_length,
                                                                list_lanes=self.list_lanes)
        num_in_part1 = [len(part1[lane]) for lane in self.list_exiting_lanes]
        num_in_part2 = [len(part2[lane]) for lane in self.list_exiting_lanes]
        num_in_part3 = [len(part3[lane]) for lane in self.list_exiting_lanes]
        num_in_part4 = [len(part4[lane]) for lane in self.list_exiting_lanes]

        total_in = []
        for i in range(len(self.list_entering_lanes)):
            total_in.extend([num_in_part1[i], num_in_part2[i], num_in_part3[i], num_in_part4[i]])
        if self.padding:
            total_in = total_in + self.padding1
        return total_in

    def _get_several_segments2(self, lane_vehicles, vehicle_distance, vehicle_speed,
                              lane_length, list_lanes):
        # get four segments [100, 200, 300, 400] for segment
        obs_length = 100
        part1, part2, part3, part4 = {}, {}, {}, {}
        for lane in list_lanes:
            part1[lane], part2[lane], part3[lane], part4[lane] = [], [], [], []
            for vehicle in lane_vehicles[lane]:
                # set as num_vehicle
                if "shadow" in vehicle:  # remove the shadow
                    vehicle = vehicle[:-7]
                    continue
                temp_v_distance = vehicle_distance[vehicle]
                if temp_v_distance <= obs_length:
                    part1[lane].append(vehicle)
                elif obs_length*2>= temp_v_distance > obs_length:
                    part2[lane].append(vehicle)
                elif obs_length*3>= temp_v_distance > obs_length*2:
                    part3[lane].append(vehicle)
                elif obs_length*4>= temp_v_distance > obs_length*3:
                    part4[lane].append(vehicle)
                
        return part1, part2, part3, part4
    

    def _get_pressure(self, l_in, l_out):
        return list(np.array(l_in)-np.array(l_out))

    def _get_pressure_map(self):
        return [self.dic_lane_waiting_vehicle_count_current_step[lane] for lane in self.list_entering_lanes] + \
               [-self.dic_lane_waiting_vehicle_count_current_step[lane] for lane in self.list_exiting_lanes]

    def _get_lane_queue_length(self, list_lanes):
        """
        queue length for each lane
        """
        res = [self.dic_lane_waiting_vehicle_count_current_step[lane] for lane in list_lanes]
        if self.padding:
            res = res + self.padding2
        return res

    def _get_lane_num_vehicles(self, list_lanes):
        res = [len(self.dic_lane_vehicle_current_step[lane]) for lane in list_lanes]
        if self.padding:
            res = res + self.padding2
        return res

    def _get_lane_num_vehicle_entring(self):
        """
        vehicle number for each lane
        """
        res = [len(self.dic_lane_vehicle_current_step[lane]) for lane in self.list_entering_lanes]
        if self.padding:
            res = res + self.padding2
        return res

    def _get_lane_num_vehicle_downstream(self):
        """
        vehicle number for each lane, exiting
        """
        return [len(self.dic_lane_vehicle_current_step[lane]) for lane in self.list_exiting_lanes]

    # ================= get functions from outside ======================
    def get_current_time(self):
        return self.eng.get_current_time()

    def get_dic_vehicle_arrive_leave_time(self):
        return self.dic_vehicle_arrive_leave_time

    def get_feature(self):
        return self.dic_feature

    def get_state(self, list_state_features):
        dic_state = {state_feature_name: self.dic_feature[state_feature_name] for
                     state_feature_name in list_state_features}
        return dic_state

    def _get_adjacency_row(self):
        return self.adjacency_row

    def get_reward(self, dic_reward_info):
        dic_reward = dict()
        # dic_reward["sum_lane_queue_length"] = None
        dic_reward["pressure"] = np.absolute(np.sum(self.dic_feature["pressure"]))
        dic_reward["queue_length"] = np.absolute(np.sum(self.dic_feature["lane_queue_vehicle_in"]))
        reward = 0
        for r in dic_reward_info:
            if dic_reward_info[r] != 0:
                reward += dic_reward_info[r] * dic_reward[r]
        return reward


class CityFlowEnv:

    def __init__(self, path_to_log, path_to_work_directory, dic_traffic_env_conf):
        self.path_to_log = path_to_log
        self.path_to_work_directory = path_to_work_directory
        self.dic_traffic_env_conf = dic_traffic_env_conf
        self.num_intersection = dic_traffic_env_conf["NUM_INTERSECTIONS"]

        self.current_time = None
        self.id_to_index = None
        self.traffic_light_node_dict = None
        self.eng = None
        self.list_intersection = None
        self.list_inter_log = None
        self.list_lanes = None
        self.system_states = None
        self.lane_length = None
        self.list_action = None
        self.list_need_action = None
        self.list_memory = None
        # check min action time
        if self.dic_traffic_env_conf["MIN_ACTION_TIME"] <= self.dic_traffic_env_conf["YELLOW_TIME"]:
            """ include the yellow time in action time """
            print("MIN_ACTION_TIME should include YELLOW_TIME")
            sys.exit()

        # touch new inter_{}.pkl (if exists, remove)
        for inter_ind in range(self.dic_traffic_env_conf["NUM_INTERSECTIONS"]):
            path_to_log_file = os.path.join(self.path_to_log, "inter_{0}.pkl".format(inter_ind))
            f = open(path_to_log_file, "wb")
            f.close()

    def reset(self):
        print(" ============= self.eng.reset() to be implemented ==========")
        cityflow_config = {
            "interval": self.dic_traffic_env_conf["INTERVAL"],
            "seed": 0,
            "laneChange": True,
            "dir": self.path_to_work_directory+"/",
            "roadnetFile": self.dic_traffic_env_conf["ROADNET_FILE"],
            "flowFile": self.dic_traffic_env_conf["TRAFFIC_FILE"],
            "rlTrafficLight": True,
            "saveReplay": False,
            "roadnetLogFile": "frontend/web/roadnetLogFile.json",
            "replayLogFile": "frontend/web/replayLogFile.txt"
        }
        # print(cityflow_config)
        with open(os.path.join(self.path_to_work_directory, "cityflow.config"), "w") as json_file:
            json.dump(cityflow_config, json_file)

        self.eng = engine.Engine(os.path.join(self.path_to_work_directory, "cityflow.config"), thread_num=1)

        # get adjacency
        self.traffic_light_node_dict = self._adjacency_extraction()

        # get lane length
        _, self.lane_length = self.get_lane_length()

        # initialize intersections (grid)
        self.list_intersection = [Intersection((i+1, j+1), self.dic_traffic_env_conf, self.eng,
                                               self.traffic_light_node_dict["intersection_{0}_{1}".format(i+1, j+1)],
                                               self.path_to_log,
                                               self.lane_length)
                                  for i in range(self.dic_traffic_env_conf["NUM_COL"])
                                  for j in range(self.dic_traffic_env_conf["NUM_ROW"])]
        self.list_inter_log = [[] for _ in range(self.dic_traffic_env_conf["NUM_COL"] *
                                                 self.dic_traffic_env_conf["NUM_ROW"])]
        self.list_action = [[] for _ in range(self.num_intersection)]
        self.list_need_action = [i for i in range(self.dic_traffic_env_conf["NUM_INTERSECTIONS"])]
        # [phase-action, duration-action, cur_state, reward]
        self.list_memory = [[[], [], [], []] for _ in range(self.num_intersection)]
        self.id_to_index = {}
        count = 0
        for i in range(self.dic_traffic_env_conf["NUM_COL"]):
            for j in range(self.dic_traffic_env_conf["NUM_ROW"]):
                self.id_to_index["intersection_{0}_{1}".format(i+1, j+1)] = count
                count += 1

        self.list_lanes = []
        for inter in self.list_intersection:
            self.list_lanes += inter.list_lanes
        self.list_lanes = np.unique(self.list_lanes).tolist()
        # get new measurements
        self.system_states = {"get_lane_vehicles": self.eng.get_lane_vehicles(),
                              "get_lane_waiting_vehicle_count": self.eng.get_lane_waiting_vehicle_count(),
                              "get_vehicle_speed": self.eng.get_vehicle_speed(),
                              "get_vehicle_distance": self.eng.get_vehicle_distance(),
                              }
        for inter in self.list_intersection:
            inter.update_current_measurements(self.system_states)
        state = self.get_state()
        return state, self.get_current_time(), self.list_need_action

    def step(self, action1, action2):
        """
        action1: set the phase, 多个数值, list
        action2: set the duration， 多个数值, list
        inter_idx: which intersection
        """
        # 补充action, 记录 action， cur_state
        before_action_state = self.get_state2(self.list_need_action)
        for i, inter_id in enumerate(self.list_need_action):
            tmp_duration = self.dic_traffic_env_conf["ACTION_DURATION"][action2[i]]
            self.list_action[inter_id].extend([action1[i]] * tmp_duration)
            # 记录 action， cur_state
            self.list_memory[inter_id][0] = action1[i]
            self.list_memory[inter_id][1] = action2[i]
            self.list_memory[inter_id][2] = before_action_state[i]
        # 寻找最短的action
        list_len = []
        for i in range(self.num_intersection):
            list_len.append(len(self.list_action[i]))
        min_duartion = min(list_len)
        # 设置合适的 action
        list_action_in_sec = [[] for _ in range(min_duartion)]
        for i in range(min_duartion):
            list_action_in_sec[i].extend([self.list_action[id][i] for id in range(self.num_intersection)])
        # 更新 记录的 list—action
        new_need = []
        for i in range(self.num_intersection):
            self.list_action[i] = self.list_action[i][min_duartion:]
            if len(self.list_action[i]) == 0:
                new_need.append(i)
        # 更新 list-need-action, 记录下一步需要 action的inter
        self.list_need_action = new_need
        step_start_time = time.time()
        for i in range(min_duartion):
            action_in_sec = list_action_in_sec[i]
            instant_time = self.get_current_time()
            self.current_time = self.get_current_time()
            # state = self.get_state()
            if i == 0:
                print("time: {0}".format(instant_time))
            self._inner_step(action_in_sec)
            #  record every reward for average reward
            tmp_reward = self.get_reward()
            for inter_id in range(self.num_intersection):
                self.list_memory[inter_id][3].append(tmp_reward[inter_id])

        final_reward = self.get_reward2(self.list_need_action)
        next_state = self.get_state2(self.list_need_action)
        self.log(after_action_state=next_state,
                 final_reward=final_reward,
                 )
        step_time = self.get_current_time()
        print("Step time: ", time.time() - step_start_time)
        return next_state, step_time, self.list_need_action

    def select_phase_feature(self, state):
        """
        state: [None, 12]
        return: [None, 2]
        """
        feat = state[self.dic_traffic_env_conf["LIST_STATE_FEATURE"][1]]
        return state

    def _get_phase_reward(self, list_state):
        """
        get the reward for each specific phase
        assert the first feat is for reward
        """
        phase_rewards = []
        reward_name = self.dic_traffic_env_conf["LIST_STATE_FEATURE"][0]
        for i in range(self.list_intersection):
            tmp_state = list_state[i][reward_name]
            tmp_action = self.list_action[i][0]
            if tmp_action == 0:
                p_reward = tmp_state[1] + tmp_state[4]
            elif tmp_action == 1:
                p_reward = tmp_state[7] + tmp_state[10]
            elif tmp_action == 2:
                p_reward = tmp_state[0] + tmp_state[3]
            elif tmp_action == 3:
                p_reward = tmp_state[6] + tmp_state[9]
            phase_rewards.append(-p_reward)
        return phase_rewards

    def _inner_step(self, action):
        # copy current measurements to previous measurements
        for inter in self.list_intersection:
            inter.update_previous_measurements()
        # set signals
        # multi_intersection decided by action {inter_id: phase}
        for inter_ind, inter in enumerate(self.list_intersection):
            inter.set_signal(
                action=action[inter_ind],
                action_pattern=self.dic_traffic_env_conf["ACTION_PATTERN"],
                yellow_time=self.dic_traffic_env_conf["YELLOW_TIME"],
                path_to_log=self.path_to_log
            )
        # run one step
        for i in range(int(1/self.dic_traffic_env_conf["INTERVAL"])):
            self.eng.next_step()

        self.system_states = {"get_lane_vehicles": self.eng.get_lane_vehicles(),
                              "get_lane_waiting_vehicle_count": self.eng.get_lane_waiting_vehicle_count(),
                              "get_vehicle_speed": self.eng.get_vehicle_speed(),
                              "get_vehicle_distance": self.eng.get_vehicle_distance()
                              }

        for inter in self.list_intersection:
            inter.update_current_measurements(self.system_states)

    def get_feature(self):
        list_feature = [inter.get_feature() for inter in self.list_intersection]
        return list_feature

    def get_state2(self, inter_idx):
        list_state = [self.list_intersection[idx].get_state(self.dic_traffic_env_conf["LIST_STATE_FEATURE"])
                      for idx in inter_idx]
        return list_state

    def get_state(self):
        list_state = [inter.get_state(self.dic_traffic_env_conf["LIST_STATE_FEATURE"]) for inter in self.list_intersection]
        return list_state

    def get_reward2(self, inter_idx):
        list_reward = [self.list_intersection[idx].get_reward(self.dic_traffic_env_conf["DIC_REWARD_INFO"]) for
                       idx in inter_idx]
        return list_reward

    def get_reward(self):
        list_reward = [inter.get_reward(self.dic_traffic_env_conf["DIC_REWARD_INFO"]) for inter in self.list_intersection]
        return list_reward

    def get_current_time(self):
        return self.eng.get_current_time()

    def log(self, after_action_state, final_reward):
        # [cur-state, phase-action, duration-action, next-state, final-reward, average_reward]
        for i, inter_ind in enumerate(self.list_need_action):
            # action1 = self.list_memory[inter_ind][0]
            average_reward = np.mean(self.list_memory[inter_ind][3])
            self.list_memory[inter_ind][3] = []
            self.list_inter_log[inter_ind].append([self.select_phase_feature(self.list_memory[inter_ind][2]),
                                                   self.list_memory[inter_ind][0],  # phase action
                                                   self.list_memory[inter_ind][1],  # duration action
                                                   self.select_phase_feature(after_action_state[i]),
                                                   final_reward[i],
                                                   average_reward
                                                   ])

    def batch_log_2(self):
        """
        Used for model test, only log the vehicle_inter_.csv
        """
        for inter_ind in range(self.dic_traffic_env_conf["NUM_INTERSECTIONS"]):
            # changed from origin
            if int(inter_ind) % 100 == 0:
                print("Batch log for inter ", inter_ind)
            path_to_log_file = os.path.join(self.path_to_log, "vehicle_inter_{0}.csv".format(inter_ind))
            dic_vehicle = self.list_intersection[inter_ind].get_dic_vehicle_arrive_leave_time()
            df = pd.DataFrame.from_dict(dic_vehicle, orient="index")
            df.to_csv(path_to_log_file, na_rep="nan")

    def batch_log(self, start, stop):
        """
        only log inter_{}.pkl
        """
        for inter_ind in range(start, stop):
            # changed from origin
            if int(inter_ind) % 100 == 0:
                print("Batch log for inter ", inter_ind)
            # path_to_log_file = os.path.join(self.path_to_log, "vehicle_inter_{0}.csv".format(inter_ind))
            # dic_vehicle = self.list_intersection[inter_ind].get_dic_vehicle_arrive_leave_time()
            # df = pd.DataFrame.from_dict(dic_vehicle, orient="index")
            # df.to_csv(path_to_log_file, na_rep="nan")
            path_to_log_file = os.path.join(self.path_to_log, "inter_{0}.pkl".format(inter_ind))
            f = open(path_to_log_file, "wb")
            pickle.dump(self.list_inter_log[inter_ind], f)
            f.close()

    def bulk_log_multi_process(self, batch_size=100):
        assert len(self.list_intersection) == len(self.list_inter_log)
        if batch_size > len(self.list_intersection):
            batch_size_run = len(self.list_intersection)
        else:
            batch_size_run = batch_size
        process_list = []
        for batch in range(0, len(self.list_intersection), batch_size_run):
            start = batch
            stop = min(batch + batch_size, len(self.list_intersection))
            p = Process(target=self.batch_log, args=(start, stop))
            print("before")
            p.start()
            print("end")
            process_list.append(p)
        print("before join")

        for t in process_list:
            t.join()
        print("end join")

    def _adjacency_extraction(self):
        traffic_light_node_dict = {}
        file = os.path.join(self.path_to_work_directory, self.dic_traffic_env_conf["ROADNET_FILE"])
        with open("{0}".format(file)) as json_data:
            net = json.load(json_data)
            for inter in net["intersections"]:
                if not inter["virtual"]:
                    traffic_light_node_dict[inter["id"]] = {"location": {"x": float(inter["point"]["x"]),
                                                                         "y": float(inter["point"]["y"])},
                                                            "total_inter_num": None, "adjacency_row": None,
                                                            "inter_id_to_index": None,
                                                            "neighbor_ENWS": None}

            top_k = self.dic_traffic_env_conf["TOP_K_ADJACENCY"]
            total_inter_num = len(traffic_light_node_dict.keys())
            inter_id_to_index = {}

            edge_id_dict = {}
            for road in net["roads"]:
                if road["id"] not in edge_id_dict.keys():
                    edge_id_dict[road["id"]] = {}
                edge_id_dict[road["id"]]["from"] = road["startIntersection"]
                edge_id_dict[road["id"]]["to"] = road["endIntersection"]

            index = 0
            for i in traffic_light_node_dict.keys():
                inter_id_to_index[i] = index
                index += 1

            for i in traffic_light_node_dict.keys():
                location_1 = traffic_light_node_dict[i]["location"]

                row = np.array([0]*total_inter_num)
                # row = np.zeros((self.dic_traffic_env_conf["NUM_ROW"],self.dic_traffic_env_conf["NUM_col"]))
                for j in traffic_light_node_dict.keys():
                    location_2 = traffic_light_node_dict[j]["location"]
                    dist = self._cal_distance(location_1, location_2)
                    row[inter_id_to_index[j]] = dist
                if len(row) == top_k:
                    adjacency_row_unsorted = np.argpartition(row, -1)[:top_k].tolist()
                elif len(row) > top_k:
                    adjacency_row_unsorted = np.argpartition(row, top_k)[:top_k].tolist()
                else:
                    adjacency_row_unsorted = [k for k in range(total_inter_num)]
                adjacency_row_unsorted.remove(inter_id_to_index[i])
                traffic_light_node_dict[i]["adjacency_row"] = [inter_id_to_index[i]]+adjacency_row_unsorted
                traffic_light_node_dict[i]["total_inter_num"] = total_inter_num

            for i in traffic_light_node_dict.keys():
                traffic_light_node_dict[i]["total_inter_num"] = inter_id_to_index
                traffic_light_node_dict[i]["neighbor_ENWS"] = []
                for j in range(4):
                    road_id = i.replace("intersection", "road")+"_"+str(j)
                    if edge_id_dict[road_id]["to"] not in traffic_light_node_dict.keys():
                        traffic_light_node_dict[i]["neighbor_ENWS"].append(None)
                    else:
                        traffic_light_node_dict[i]["neighbor_ENWS"].append(edge_id_dict[road_id]["to"])

        return traffic_light_node_dict

    @staticmethod
    def _cal_distance(loc_dict1, loc_dict2):
        a = np.array((loc_dict1["x"], loc_dict1["y"]))
        b = np.array((loc_dict2["x"], loc_dict2["y"]))
        return np.sqrt(np.sum((a-b)**2))

    @staticmethod
    def end_cityflow():
        print("============== cityflow process end ===============")

    def get_lane_length(self):
        """
        newly added part for get lane length
        Read the road net file
        Return: dict{lanes} normalized with the min lane length
        """
        file = os.path.join(self.path_to_work_directory, self.dic_traffic_env_conf["ROADNET_FILE"])
        with open(file) as json_data:
            net = json.load(json_data)
        roads = net['roads']
        lanes_length_dict = {}
        lane_normalize_factor = {}

        for road in roads:
            points = road["points"]
            road_length = abs(points[0]['x'] + points[0]['y'] - points[1]['x'] - points[1]['y'])
            for i in range(4):
                lane_id = road['id'] + "_{0}".format(i)
                lanes_length_dict[lane_id] = road_length
        min_length = min(lanes_length_dict.values())

        for key, value in lanes_length_dict.items():
            lane_normalize_factor[key] = value / min_length
        return lane_normalize_factor, lanes_length_dict
"
"from models.dynamiclight import DynamicLightAgent
from models.dynamiclightft import DynamicLightAgentFT
from models.phase_dynamic_0_1 import DynamicAgent01
from models.phase_dynamic_0_2 import DynamicAgent02
from models.phase_dynamic_0_3 import DynamicAgent03
from models.phase_dynamic_0_5 import DynamicAgent05
from models.phase_dynamic_mp import DynamicAgentMP
from models.phase_dynamic_rand import DynamicAgentRand
from models.phase_dynamic_ft import DynamicAgentFT
from models.phase_dynamic_lite import DynamicAgentLite
from models.phase_dynamic_multi import DynamicAgentMulti



DIC_AGENTS = {
    "dynamiclight": DynamicLightAgent,
    "dynamiclightft": DynamicLightAgentFT,
    "Dyn01": DynamicAgent01,
    "Dyn02": DynamicAgent02,
    "Dyn03": DynamicAgent03,
    "Dyn05": DynamicAgent05,
    "DynC": DynamicAgentFT,
    "DynMP": DynamicAgentMP,
    "DynLite": DynamicAgentLite,
    "DynM": DynamicAgentMulti,
    "DynRand": DynamicAgentRand
}

DIC_PATH = {
    "PATH_TO_MODEL": "model/default",
    "PATH_TO_WORK_DIRECTORY": "records/default",
    "PATH_TO_DATA": "data/template",
    "PATH_TO_PRETRAIN_MODEL": "model/default",
    "PATH_TO_ERROR": "errors/default",
}

dic_traffic_env_conf = {
    "FORGET_ROUND": 1,
    "RUN_COUNTS": 3600,
    "MODEL_NAME": None,
    "TOP_K_ADJACENCY": 5,
    "ACTION_PATTERN": "set",
    "NUM_INTERSECTIONS": 1,
    "MAX_LANE": 16,
    "OBS_LENGTH": 167,
    "MIN_ACTION_TIME": 15,
    "MEASURE_TIME": 15,
    "BINARY_PHASE_EXPANSION": True,
    "YELLOW_TIME": 5,
    "ALL_RED_TIME": 0,
    "NUM_PHASES": 4,
    "NUM_LANES": [3, 3, 3, 3],
    "INTERVAL": 1,
    "ACTION_DURATION": {
                0: 10,
                1: 15,
                2: 20,
                3: 25,
                4: 30,
                5: 35,
                6: 40
            },

    "LIST_STATE_FEATURE": [
        "cur_phase",
        "time_this_phase",
        "lane_num_vehicle",
        "lane_num_vehicle_downstream",
        "traffic_movement_pressure_num",
        "traffic_movement_pressure_queue",
        "traffic_movement_pressure_queue_efficient",
        "pressure",
        "adjacency_matrix"
    ],
    "DIC_REWARD_INFO": {
        "queue_length": 0,
        "pressure": 0,
    },
    "PHASE": {
        1: [0, 1, 0, 1, 0, 0, 0, 0],
        2: [0, 0, 0, 0, 0, 1, 0, 1],
        3: [1, 0, 1, 0, 0, 0, 0, 0],
        4: [0, 0, 0, 0, 1, 0, 1, 0]
        },
    "PHASETOTAL": {
        1: [0, 1, 1, 0, 1, 1, 0, 0, 1, 0, 0, 1],
        2: [0, 0, 1, 0, 0, 1, 0, 1, 1, 0, 1, 1],
        3: [1, 0, 1, 1, 0, 1, 0, 0, 1, 0, 0, 1],
        4: [0, 0, 1, 0, 0, 1, 1, 0, 1, 1, 0, 1]
        },
    "list_lane_order": ["WL", "WT", "EL", "ET", "NL", "NT", "SL", "ST"],
    "PHASE_LIST": ['WT_ET', 'NT_ST', 'WL_EL', 'NL_SL'],

}

DIC_BASE_AGENT_CONF = {
    "SROUND": 80,
    "SROUND2": 160,
    "SROUND3": 110,
    "D_DENSE": 20,
    "LEARNING_RATE": 0.001,
    "LEARNING_RATE2": 0.0002,
    "LEARNING_RATE3": 0.0005,
    "PATIENCE": 10,
    "BATCH_SIZE1": 128,
    "EPOCHS": 100,
    "SAMPLE_SIZE1": 4000,
    "MAX_MEMORY_LEN": 12000,

    "UPDATE_Q_BAR_FREQ": 5,
    "UPDATE_Q_BAR_EVERY_C_ROUND": False,

    "GAMMA": 0.8,
    "NORMAL_FACTOR": 20,

    "EPSILON": 0.8,
    "EPSILON_DECAY": 0.95,
    "MIN_EPSILON": 0.2,
    "LOSS_FUNCTION": "mean_squared_error",
}

DIC_FIXEDTIME_AGENT_CONF = {
    "FIXED_TIME": [15, 15, 15, 15]
}

DIC_MAXPRESSURE_AGENT_CONF = {
    "FIXED_TIME": [15, 15, 15, 15]
}
"
"from .config import DIC_AGENTS
from .cityflow_env import CityFlowEnv
import time
import os
import copy


class Generator:
    def __init__(self, cnt_round, cnt_gen, dic_path, dic_agent_conf, dic_traffic_env_conf):

        self.cnt_round = cnt_round
        self.cnt_gen = cnt_gen
        self.dic_path = dic_path
        self.dic_agent_conf = copy.deepcopy(dic_agent_conf)
        self.dic_traffic_env_conf = dic_traffic_env_conf
        self.agents = [None]*dic_traffic_env_conf['NUM_AGENTS']
        self.path_to_log = os.path.join(self.dic_path["PATH_TO_WORK_DIRECTORY"], "train_round",
                                        "round_"+str(self.cnt_round), "generator_"+str(self.cnt_gen))
        if not os.path.exists(self.path_to_log):
            os.makedirs(self.path_to_log)
            start_time = time.time()
            for i in range(dic_traffic_env_conf['NUM_AGENTS']):
                agent_name = self.dic_traffic_env_conf["MODEL_NAME"]
                agent = DIC_AGENTS[agent_name](
                    dic_agent_conf=self.dic_agent_conf,
                    dic_traffic_env_conf=self.dic_traffic_env_conf,
                    dic_path=self.dic_path,
                    cnt_round=self.cnt_round,
                    intersection_id=str(i),
                )
                self.agents[i] = agent
            print("Create intersection agent time: ", time.time()-start_time)

        self.env = CityFlowEnv(
            path_to_log=self.path_to_log,
            path_to_work_directory=self.dic_path["PATH_TO_WORK_DIRECTORY"],
            dic_traffic_env_conf=self.dic_traffic_env_conf
        )

    def generate(self):
        reset_env_start_time = time.time()
        state, step_time, list_need = self.env.reset()
        reset_env_time = time.time() - reset_env_start_time
        running_start_time = time.time()
        while step_time < self.dic_traffic_env_conf["RUN_COUNTS"]:
            step_start_time = time.time()
            for i in range(self.dic_traffic_env_conf["NUM_AGENTS"]):
                phase_action, duration_action = self.agents[i].choose_action(state, list_need)
            # print(phase_action, duration_action)
            next_state, step_time, list_need = self.env.step(phase_action, duration_action)

            print("time: {0}, running_time: {1}".format(self.env.get_current_time(),
                                                        time.time()-step_start_time))

            state = next_state

        running_time = time.time() - running_start_time
        log_start_time = time.time()
        print("start logging.......................")
        self.env.bulk_log_multi_process()
        log_time = time.time() - log_start_time
        self.env.end_cityflow()
        print("reset_env_time: ", reset_env_time)
        print("running_time_total: ", running_time)
        print("log_time: ", log_time)
"
"from .config import DIC_AGENTS
import pickle
import os
import time
import traceback
import random


class Updater:

    def __init__(self, cnt_round, dic_agent_conf, dic_traffic_env_conf, dic_path):
        self.cnt_round = cnt_round
        self.dic_path = dic_path
        self.dic_traffic_env_conf = dic_traffic_env_conf
        self.dic_agent_conf = dic_agent_conf
        self.agents = []
        self.sample_set_list = []
        self.sample_indexes1, self.sample_indexes2 = None, None
        print("Number of agents: ", dic_traffic_env_conf['NUM_AGENTS'])
        for i in range(dic_traffic_env_conf['NUM_AGENTS']):
            agent_name = self.dic_traffic_env_conf["MODEL_NAME"]
            agent = DIC_AGENTS[agent_name](
                self.dic_agent_conf, self.dic_traffic_env_conf,
                self.dic_path, self.cnt_round, str(i))
            self.agents.append(agent)

    def load_sample_with_forget(self, i):
        sample_set = []
        try:
            sample_file = open(os.path.join(self.dic_path["PATH_TO_WORK_DIRECTORY"], "train_round",
                                            "total_samples_inter_{0}".format(i) + ".pkl"), "rb")
            try:
                cur_sample_set = []
                while True:
                    cur_sample_set += pickle.load(sample_file)
                sample_file.close()
            except EOFError:
                print("===== load samples finished =====")
                sample_file.close()

            ind_end = len(cur_sample_set)
            ind_sta = max(0, ind_end - self.dic_agent_conf["MAX_MEMORY_LEN"])
            # forget
            memory_after_forget = cur_sample_set[ind_sta: ind_end]
            print("==== memory size after forget ====:", len(memory_after_forget))
            if self.cnt_round % self.dic_traffic_env_conf["FORGET_ROUND"] == 0:
                with open(os.path.join(self.dic_path["PATH_TO_WORK_DIRECTORY"], "train_round",
                                       "total_samples_inter_{0}".format(i) + ".pkl"), "wb+") as f:
                    pickle.dump(memory_after_forget, f, -1)
            
            sample_size1 = min(self.dic_agent_conf["SAMPLE_SIZE1"], len(memory_after_forget))
            self.sample_indexes1 = random.sample(range(len(memory_after_forget)), sample_size1)
            sample_set1 = [memory_after_forget[k] for k in self.sample_indexes1]
            sample_set2 = []
            print("==== memory samples number =====:",  sample_size1)
            if self.cnt_round == self.dic_agent_conf["SROUND"]-1 or self.cnt_round == self.dic_agent_conf["SROUND2"]-1:
                with open(os.path.join(self.dic_path["PATH_TO_WORK_DIRECTORY"], "train_round",
                                       "total_samples_inter_{0}".format(i) + ".pkl"), "wb+") as f:
                    pickle.dump([], f, -1)
        except:
            error_dir = os.path.join(self.dic_path["PATH_TO_WORK_DIRECTORY"]).replace("records", "errors")
            if not os.path.exists(error_dir):
                os.makedirs(error_dir)
            f = open(os.path.join(error_dir, "error_info_inter_{0}.txt".format(i)), "a")
            f.write("Fail to load samples for inter {0}\n".format(i))
            f.write('traceback.format_exc():\n%s\n' % traceback.format_exc())
            f.close()
            print('traceback.format_exc():\n%s' % traceback.format_exc())
            pass
        if i % 100 == 0:
            print("load_sample for inter {0}".format(i))
        return sample_set1, sample_set2

    def load_sample_for_agents(self):
        start_time = time.time()
        print("Start load samples at", start_time)
        sample_set_all1,  sample_set_all2 = [], []
        for i in range(self.dic_traffic_env_conf['NUM_INTERSECTIONS']):
            sample_set1, sample_set2 = self.load_sample_with_forget(i)
            sample_set_all1.extend(sample_set1)
            sample_set_all2.extend(sample_set2)
        self.agents[0].prepare_Xs_Y(sample_set_all1)
            
    def update_network(self, i):
        print('update agent %d' % i)
        self.agents[i].train_network()
        self.agents[i].save_network("round_{0}_inter_{1}".format(self.cnt_round, self.agents[i].intersection_id))

    def update_network_for_agents(self):
        print("update_network_for_agents", self.dic_traffic_env_conf['NUM_AGENTS'])
        for i in range(self.dic_traffic_env_conf['NUM_AGENTS']):
            self.update_network(i)
"

this is libsignal's guide :"# Customize Agent

This documentation overviews how to create a new agent, including importing an agent from other libs.

## Create a New Agent Class

To begin with, we should create a new model implemented from **RLAgent** for RL-based agents or **BaseAgent** for other types of agents.

Here we take the RL-based agent as an example. We would like to develop a new model, named `NewModel`, for traffic signal control tasks.

First please create a new file `newmodel.py` in the directory `LibSignal/agent/` and write the following code into the file.

```
from . import RLAgent
from common.registry import Registry

@Registry.register_model('NewModel')
class NewModelAgent(RLAgent):
    def __init__(self, world, rank):
        pass
    
    def reset(self):
        pass
    
    def __repr__(self):
        pass
        
    def get_ob(self):
        pass
    
    def get_reward(self):
        pass
        
    def get_phase(self):
        pass
        
    def get_queue(self):
        pass
        
    def get_delay(self):
        pass
        
    def get_action(self):
        pass
        
    def sample(self):
        pass
        
    def remember(self):
        pass
        
    def build_model(self):
        pass
        
    def update_target_network(self):
        pass
        
    def train(self):
        pass
        
    def load_model(self):
        pass
        
    def save_model(self):
        pass

class NewModelNet():
    def __init__(self, **kwargs):
        pass
    
    def forward(self, x):
        pass
```

You can see that we use `Registry` to register the new model, define the `NewModelNet` class, and list the methods that must appear in the class.
In the following sections, we take the configuration of DQN model as an example.

## Implement \_\_init\_\_()

Then we implement `__init__()` method, which is used to define the model structure according to the parameters of input and configuration information from `Registry`.

The input parameters of `__init__()` are `world` and this intersection's `rank`. The configuration information from `Registry` contains various model parameters, input parameters required by creating generators and so on.

You can define `__init__()` like this:

```

@Registry.register_model('NewModel')
class NewModelAgent(RLAgent):
    def __init__(self, world, rank):
        # section 1: get configs
        self.world = world
        self.rank = rank
        self.sub_agents = len(self.world.intersections)
        self.inter_id = self.world.intersection_ids[self.rank]
        self.inter = self.world.id2intersection[self.inter_id]

        self.phase = Registry.mapping['model_mapping']['setting'].param['phase']
        self.one_hot = Registry.mapping['model_mapping']['setting'].param['one_hot']
        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.learning_rate = Registry.mapping['model_mapping']['setting'].param['learning_rate']
        self.vehicle_max = Registry.mapping['model_mapping']['setting'].param['vehicle_max']
        self.batch_size = Registry.mapping['model_mapping']['setting'].param['batch_size']

        # section 2: create generators for each Agent
        self.ob_generator = LaneVehicleGenerator(self.world,  self.inter, ['lane_count'], in_only=True, average=None)
        self.phase_generator = IntersectionPhaseGenerator(world,  self.inter, ["phase"],
                                                          targets=["cur_phase"], negative=False)
        self.reward_generator = LaneVehicleGenerator(self.world,  self.inter, ["lane_waiting_count"],
                                                     in_only=True, average='all', negative=True)

        # section 3: set action space and ob_length
        self.action_space = gym.spaces.Discrete(len(self.inter.phases))
        if self.phase:
            if self.one_hot:
                self.ob_length = self.ob_generator.ob_length + len(self.inter.phases)
            else:
                self.ob_length = self.ob_generator.ob_length + 1
        else:
            self.ob_length = self.ob_generator.ob_length

        # section 4: create model, target model and others
        self.model = self._build_model()
        self.target_model = self._build_model()
        self.update_target_network()
        self.criterion = nn.MSELoss(reduction='mean')
        self.optimizer = optim.RMSprop(self.model.parameters(),
                                       lr=self.learning_rate,
                                       alpha=0.9, centered=False, eps=1e-7)
```

You can see that in the section 1 of the code, we take the necessary parameters from `Registry` and input parameters.

In the section 2 of the code, we create observation, reward and phase generators for the model. Notes that different models have different input parameters of generators. Since the definitions of queue and delay are unified, queue and delay generators are directly created in the **RLAgent**.

In the section 3 of the code, we set action space list and observation length.

In the section 4 of the code, we create model, target model and others, including a criterion and an optimizer.

.. Note::
For Multi-Agent that shares information among agents, the code for creating generators is as follows:

```
# get generators for Agent
observation_generators = []
for inter in self.world.intersections:
    node_id = inter.id
    node_idx = self.world.id2idx[node_id]
    node_obj = self.world.id2intersection[node_id]
    tmp_generator = LaneVehicleGenerator(self.world, node_obj, ['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
rewarding_generators = []
for inter in self.world.intersections:
    node_id = inter.id
    node_idx = self.world.id2idx[node_id]
    node_obj = self.world.id2intersection[node_id]
    tmp_generator = LaneVehicleGenerator(self.world, node_obj, ["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

#  get phase generator
phasing_generators = []
for inter in self.world.intersections:
    node_id = inter.id
    node_idx = self.world.id2idx[node_id]
    node_obj = self.world.id2intersection[node_id]
    tmp_generator = IntersectionPhaseGenerator(self.world, node_obj, ['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

#  get queue generator
queues = []
for inter in self.world.intersections:
    node_id = inter.id
    node_idx = self.world.id2idx[node_id]
    node_obj = self.world.id2intersection[node_id]
    tmp_generator = LaneVehicleGenerator(self.world, node_obj, ["lane_waiting_count"], 
                                            in_only=True, negative=False)
    queues.append((node_idx, tmp_generator))
sorted(queues, key=lambda x: x[0])
self.queue = queues

#  get delay generator
delays = []
for inter in self.world.intersections:
    node_id = inter.id
    node_idx = self.world.id2idx[node_id]
    node_obj = self.world.id2intersection[node_id]
    tmp_generator = LaneVehicleGenerator(self.world, node_obj, ["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
```


## Implement reset()

Then we implement `reset()` method, `reset()` is used to reset information, including ob_generator, phase_generator, reward_generator, queue, delay, etc.
For example, you can define `reset()` like this:
```
@Registry.register_model('NewModel')
class NewModelAgent(RLAgent):
    def reset(self):
        inter_id = self.world.intersection_ids[self.rank]
        inter_obj = self.world.id2intersection[inter_id]
        self.ob_generator = LaneVehicleGenerator(self.world, inter_obj, ['lane_count'], in_only=True, average=None)
        self.phase_generator = IntersectionPhaseGenerator(self.world, inter_obj, ["phase"],
                                                          targets=["cur_phase"], negative=False)
        self.reward_generator = LaneVehicleGenerator(self.world, inter_obj, ["lane_waiting_count"],
                                                     in_only=True, average='all', negative=True)
        self.queue = LaneVehicleGenerator(self.world, inter_obj,
                                                     ["lane_waiting_count"], in_only=True,
                                                     negative=False)
        self.delay = LaneVehicleGenerator(self.world, inter_obj,
                                                     ["lane_delay"], in_only=True, average="all",
                                                     negative=False)     
```

Note: The above codes are applied for Single-Agent. For Multi-Agent generators, the user must reset by the way of traversing the list.

## Implement build_model()

Then we implement `build_model()` method, `build_model()` is used to create a model(network). This method will be called by `__init__()` twice for creating the model and target model.
For example, you can define `build_model()` like this:
```
@Registry.register_model('NewModel')
class NewModelAgent(RLAgent):
    def build_model(self):
        model = NewModelNet(self.ob_length, self.action_space.n)
        return model     
```

## Implement NewModelNet

Then we implement `NewModelNet` class, which is the core of the model. `NewModelNet` class can inherit both the base `nn.Module` class and classes already implemented in other libs.

For example, you can define `NewModelNet` like this:
```
class NewModelNet(nn.Module):
    def __init__(self):
        super(NewModelNet, self).__init__()
        pass

    def _forward(self, x):
        pass
```

## Implement \_\_repr\_\_()

Then we implement `repr()` method, `repr()` is used to return the `self.model` structure.
For example, you can define `__repr__()` like this:
```
@Registry.register_model('NewModel')
class NewModelAgent(RLAgent):
    def __repr__(self):
        return self.model.__repr__()   
```

## Implement get_ob()

Then we implement `get_ob()` method, `get_ob()` is used to get observation from environment.
For Single-Agent, you can define `get_ob()` like this:
```
@Registry.register_model('NewModel')
class NewModelAgent(RLAgent):
    def get_ob(self):
        x_obs = []
        x_obs.append(self.ob_generator.generate())
        x_obs = np.array(x_obs, dtype=np.float32)
        return x_obs  
```

For Multi-Agent, you can define `get_ob()` like this:
```
@Registry.register_model('NewModel')
class NewModelAgent(RLAgent):
    def get_ob(self):
        # sub_agents * lane_nums,
        x_obs = []  
        for i in range(len(self.ob_generator)):
            x_obs.append(self.ob_generator[i][1].generate())
        length = set([len(i) for i in x_obs])
        if len(length) == 1:
            x_obs = np.array(x_obs, dtype=np.float32)
        else:
            x_obs = [np.expand_dims(x,axis=0) for x in x_obs]
        return x_obs
```

## Implement get_reward()

Then we implement `get_reward()` method, `get_reward()` is used to get reward from environment.
For Single-Agent, you can define `get_reward()` like this:
```
@Registry.register_model('NewModel')
class NewModelAgent(RLAgent):
    def get_reward(self):
        rewards = []
        rewards.append(self.reward_generator.generate())
        rewards = np.squeeze(np.array(rewards))
        return rewards    
```

For Multi-Agent, you can define `get_reward()` like this:
```
@Registry.register_model('NewModel')
class NewModelAgent(RLAgent):
    def get_reward(self):
        rewards = []  # sub_agents
        for i in range(len(self.reward_generator)):
            rewards.append(self.reward_generator[i][1].generate())
        rewards = np.squeeze(np.array(rewards))
        return rewards
```

## Implement get_phase()

Then we implement `get_phase()` method, `get_phase()` is used to get current phase of intersection(s) from environment.
For Single-Agent, you can define `get_phase()` like this:
```
@Registry.register_model('NewModel')
class NewModelAgent(RLAgent):
    def get_phase(self):
        phase = []
        phase.append(self.phase_generator.generate())
        phase = (np.concatenate(phase)).astype(np.int8)
        return phase      
```

For Multi-Agent, you can define `get_phase()` like this:
```
@Registry.register_model('NewModel')
class NewModelAgent(RLAgent):
    def get_phase(self):
        phase = []  # sub_agents
        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
```


## Implement get_queue()

Then we implement `get_queue()` method, `get_queue()` is used to get queue length of intersection.
For Single-Agent, you can define `get_queue()` like this:
```
@Registry.register_model('NewModel')
class NewModelAgent(RLAgent):
    def get_queue(self):
        queue = []
        queue.append(self.queue.generate())
        # sum of lane nums
        queue = np.sum(np.squeeze(np.array(queue)))
        return queue    
```

For Multi-Agent, you can define `get_queue()` like this:
```
@Registry.register_model('NewModel')
class NewModelAgent(RLAgent):
    def get_queue(self):
        queue = []
        for i in range(len(self.queue)):
            queue.append((self.queue[i][1].generate()))
        tmp_queue = np.squeeze(np.array(queue))
        queue = np.sum(tmp_queue, axis=1 if len(tmp_queue.shape)==2 else 0)
        return queue  
```

## Implement get_delay()

Then we implement `get_delay()` method, `get_delay()` is used to get delay of intersection.
For Single-Agent, you can define `get_delay()` like this:
```
@Registry.register_model('NewModel')
class NewModelAgent(RLAgent):
    def get_delay(self):
        delay = []
        delay.append(self.delay.generate())
        delay = np.sum(np.squeeze(np.array(delay)))
        return delay   
```

For Multi-Agent, you can define `get_delay()` like this:
```
@Registry.register_model('NewModel')
class NewModelAgent(RLAgent):
    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))
        return delay # [intersections,]       
```

## Implement get_action()

Then we implement `get_action()` method, `get_action()` is used to generate action according to features. Different models have different features, users only need to input the required parameters. 
For DQN, it requires observation and phase as a feature, so you can define `get_action()` like this:
```
@Registry.register_model('NewModel')
class NewModelAgent(RLAgent):
    def get_action(self, ob, phase, test=False):
        if not test:
            if np.random.rand() <= self.epsilon:
                return self.sample()
        if self.phase:
            if self.one_hot:
                feature = np.concatenate([ob, utils.idx2onehot(phase, self.action_space.n)], axis=1)
            else:
                feature = np.concatenate([ob, phase], axis=1)
        else:
            feature = ob
        observation = torch.tensor(feature, dtype=torch.float32)
        actions = self.model(observation, train=False)
        actions = actions.clone().detach().numpy()
        return np.argmax(actions, axis=1)   
```

## Implement sample()

Then we implement `sample()` method, `sample()` is used to sample action randomly.
For example, you can define `sample()` like this:
```
@Registry.register_model('NewModel')
class NewModelAgent(RLAgent):
    def sample(self):
        return np.random.randint(0, self.action_space.n, self.sub_agents)     
```

## Implement remember()

Then we implement `remember()` method, `remember()` is used to put current step information into the replay buffer for training the agent later.
For example, you can define `remember()` like this:
```
@Registry.register_model('NewModel')
class NewModelAgent(RLAgent):
    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)))    
```

## Implement update_target_network()

Then we implement `update_target_network()` method, `update_target_network()` is used to update params of target network.
For example, you can define `update_target_network()` like this:
```
@Registry.register_model('NewModel')
class NewModelAgent(RLAgent):
    def update_target_network(self):
        weights = self.model.state_dict()
        self.target_model.load_state_dict(weights) 
```

## Implement train()

Then we implement `train()` method, `train()` is used to train the agent, and optimize the action generated by agent.
For example, you can define `train()` like this:
```
@Registry.register_model('NewModel')
class NewModelAgent(RLAgent):
    def train(self):
        #take batch-sized samples from the replay buffer randomly 
        samples = random.sample(self.replay_buffer, self.batch_size)
        # convert samples into corresponding formats
        b_t, b_tp, rewards, actions = self._batchwise(samples)
        # put the next_feature into target model
        out = self.target_model(b_tp, train=False)
        target = rewards + self.gamma * torch.max(out, dim=1)[0]
        # put the current_feature into target model
        target_f = self.model(b_t, train=False)
        for i, action in enumerate(actions):
            target_f[i][action] = target[i]
        loss = self.criterion(self.model(b_t, train=True), target_f)
        self.optimizer.zero_grad()
        loss.backward()
        clip_grad_norm_(self.model.parameters(), self.grad_clip)
        self.optimizer.step()
        if self.epsilon > self.epsilon_min:
            self.epsilon *= self.epsilon_decay
        return loss.clone().detach().numpy() 
```

## Implement load_model()

Then we implement `load_model()` method, `load_model()` is used to load model params of an episode.
For example, you can define `load_model()` like this:
```
@Registry.register_model('NewModel')
class NewModelAgent(RLAgent):
    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))   
```

## Implement save_model()

Then we implement `save_model()` method, `save_model()` is used to save model params of an episode.
For example, you can define `save_model()` like this:
```
@Registry.register_model('NewModel')
class NewModelAgent(RLAgent):
    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)
```

.. Note::
If the customized model needs more complex methods, then you can rewrite the corresponding interface mentioned above. If the model is imported from other libs, and some methods have already been included in the library, then users can use "pass" to skip the implementation of the corresponding interface.

## Import The Model

After adding the model, you need to modify the `__init__.py` file in `LibSignal/agent/__init__.py`.

Please add code like this:
```
from .newmodel import NewModelAgent
```


## Add Model Config

Then, you need to create the `LibSignal/configs/tsc/newmodel.yml` file, which is used to set the parameters of the model, trainer, world and logger.

For example, you can add codes like the following:
```
includes:
  - configs/tsc/base.yml

model:
  name: newmodel
  train_model: True
  epsilon: 0.1
  one_hot: True
  phase: True

trainer:
  learning_start: 1000

world:
  signal_config: {
    hz1x1: {
      phase_pairs: [[2, 6], [0, 4], [3, 7], [1, 5], [6, 7], [2, 3], [4, 5], [0, 1]],
      valid_acts: null
    }
  }

logger:
    attention: True
```

Note: The filename of the config and the value must be the same as the class name of the model you added. Just like the `NewModel` above.

</br>

Now that you have learned how to add a new model, try the following commands to run this model!

```
python run.py -a newmodel
```"

If my implementation is too bad, you can rework it from the ground up.

This is an agent example :"from . import RLAgent
import random
import numpy as np
from collections import deque
import os
import torch
import torch.nn as nn
import torch.nn.functional as F
from common.registry import Registry
import gym
from generator import LaneVehicleGenerator, IntersectionPhaseGenerator, IntersectionVehicleGenerator
from torch.nn.utils import clip_grad_norm_
from pfrl.q_functions import DiscreteActionValueHead
from agent.utils import SharedEpsGreedy
from pfrl import replay_buffers
from typing import Any, Sequence
from pfrl.agents import DQN
from pfrl.utils.contexts import evaluating
from agent import utils

'''MPLight is default set Shared Agent'''

@Registry.register_model('mplight')
class MPLightAgent(RLAgent):
    '''
    MPLightAgent consists of FRAP agent and methods for training agents, communicating with environment, etc.
    '''
    def __init__(self, world, rank):
        super().__init__(world,world.intersection_ids[rank])
        self.dic_agent_conf = Registry.mapping['model_mapping']['setting']
        self.dic_traffic_env_conf = Registry.mapping['world_mapping']['setting']
        self.dic_trainer_conf = Registry.mapping['trainer_mapping']['setting']
        
        self.gamma = self.dic_agent_conf.param["gamma"]
        self.grad_clip = self.dic_agent_conf.param["grad_clip"]
        self.learning_rate = self.dic_agent_conf.param["learning_rate"]
        self.batch_size = self.dic_agent_conf.param["batch_size"]
        self.buffer_size = Registry.mapping['trainer_mapping']['setting'].param['buffer_size']
        self.replay_buffer = replay_buffers.ReplayBuffer(self.buffer_size)


        self.world = world
        self.rank = rank
        self.sub_agents = len(self.world.intersections)
        self.inter_id = self.world.intersection_ids[self.rank]
        self.phase = self.dic_agent_conf.param['phase']
        self.one_hot = self.dic_agent_conf.param['one_hot']
        self.action_space_list = [gym.spaces.Discrete(len(x.phases)) for x in self.world.intersections]
        # create competition matrix
        map_name = self.dic_traffic_env_conf.param['network']
        self.phase_pairs = self.dic_traffic_env_conf.param['signal_config'][map_name]['phase_pairs']
        self.comp_mask = self.relation()
        self.valid_acts = self.dic_traffic_env_conf.param['signal_config'][map_name]['valid_acts']
        self.reverse_valid = None
        self.ob_order = None
        if 'lane_order' in self.dic_traffic_env_conf.param['signal_config'][map_name].keys():
            self.ob_order = self.dic_traffic_env_conf.param['signal_config'][map_name]['lane_order']
        self.model = None
        self.optimizer = None
        self.num_phases = len(self.phase_pairs)
        self.num_actions = len(self.phase_pairs)
        episodes = Registry.mapping['trainer_mapping']['setting'].param['episodes'] * 0.8
        steps = Registry.mapping['trainer_mapping']['setting'].param['steps']
        action_interval = Registry.mapping['trainer_mapping']['setting'].param['action_interval']
        total_steps = episodes * steps / action_interval
        self.explorer = SharedEpsGreedy(
                self.dic_agent_conf.param["eps_start"],
                self.dic_agent_conf.param["eps_end"],
                self.sub_agents*total_steps,
                lambda: np.random.randint(len(self.phase_pairs)),
            )
        self.agents_iner = self._build_model()
        
        # get generators for MPLightAgent
        observation_generators = []
        for inter in self.world.intersections:
            node_id = inter.id
            node_idx = self.world.id2idx[node_id]
            node_obj = self.world.id2intersection[node_id]
            tmp_generator = LaneVehicleGenerator(self.world, node_obj, ['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
        rewarding_generators = []
        for inter in self.world.intersections:
            node_id = inter.id
            node_idx = self.world.id2idx[node_id]
            node_obj = self.world.id2intersection[node_id]
            tmp_generator = LaneVehicleGenerator(self.world, node_obj, ["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

        #  get phase generator
        phasing_generators = []
        for inter in self.world.intersections:
            node_id = inter.id
            node_idx = self.world.id2idx[node_id]
            node_obj = self.world.id2intersection[node_id]
            tmp_generator = IntersectionPhaseGenerator(self.world, node_obj, ['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

        #  get queue generator
        queues = []
        for inter in self.world.intersections:
            node_id = inter.id
            node_idx = self.world.id2idx[node_id]
            node_obj = self.world.id2intersection[node_id]
            tmp_generator = LaneVehicleGenerator(self.world, node_obj, ["lane_waiting_count"], 
                                                 in_only=True, negative=False)
            queues.append((node_idx, tmp_generator))
        sorted(queues, key=lambda x: x[0])
        self.queue = queues

        #  get delay generator
        delays = []
        for inter in self.world.intersections:
            node_id = inter.id
            node_idx = self.world.id2idx[node_id]
            node_obj = self.world.id2intersection[node_id]
            tmp_generator = LaneVehicleGenerator(self.world, node_obj, ["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

        

        # TODO check ob_length, compared with presslight and original mplight and RESCO-mplight
        # this is extracted from presslight so far
        # if self.phase:
        #     if self.one_hot:
        #         self.ob_length = self.ob_generator.ob_length + len(self.inter_obj.phases) # 32
        #     else:
        #         self.ob_length = self.ob_generator.ob_length + 1 # 25
        # else:
        #     self.ob_length = self.ob_generator.ob_length # 24

    def __repr__(self):
        return self.agents_iner.__repr__()

    def reset(self):
        '''
        reset
        Reset information, including ob_generator, phase_generator, queue, delay, etc.

        :param: None
        :return: None
        '''
        observation_generators = []
        for inter in self.world.intersections:
            node_id = inter.id
            node_idx = self.world.id2idx[node_id]
            node_obj = self.world.id2intersection[node_id]
            tmp_generator = LaneVehicleGenerator(self.world, node_obj, ['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

        rewarding_generators = []
        for inter in self.world.intersections:
            node_id = inter.id
            node_idx = self.world.id2idx[node_id]
            node_obj = self.world.id2intersection[node_id]
            tmp_generator = LaneVehicleGenerator(self.world, node_obj, ["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

        phasing_generators = []
        for inter in self.world.intersections:
            node_id = inter.id
            node_idx = self.world.id2idx[node_id]
            node_obj = self.world.id2intersection[node_id]
            tmp_generator = IntersectionPhaseGenerator(self.world, node_obj, ['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

        queues = []
        for inter in self.world.intersections:
            node_id = inter.id
            node_idx = self.world.id2idx[node_id]
            node_obj = self.world.id2intersection[node_id]
            tmp_generator = LaneVehicleGenerator(self.world, node_obj, ["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
            node_idx = self.world.id2idx[node_id]
            node_obj = self.world.id2intersection[node_id]
            tmp_generator = LaneVehicleGenerator(self.world, node_obj, ["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
    
    def relation(self):
        '''
        relation
        Get the phase competition relation between traffic movements.

        :param: None
        :return comp_mask: matrix of phase competition relation
        '''
        comp_mask = []
        for i in range(len(self.phase_pairs)):
            zeros = np.zeros(len(self.phase_pairs) - 1, dtype=np.int64)
            cnt = 0
            for j in range(len(self.phase_pairs)):
                if i == j: continue
                pair_a = self.phase_pairs[i]
                pair_b = self.phase_pairs[j]
                if len(list(set(pair_a + pair_b))) == 3: zeros[cnt] = 1
                cnt += 1
            comp_mask.append(zeros)
        comp_mask = torch.from_numpy(np.asarray(comp_mask))
        return comp_mask

    def get_ob(self):
        '''
        get_ob
        Get observation from environment.

        :param: None
        :return x_obs: observation generated by ob_generator, the shape is [sub_agents,lane_nums]
        '''
        x_obs = []  # sub_agents * lane_nums,
        for i in range(len(self.ob_generator)):
            tmp = self.ob_generator[i][1].generate()
            if self.ob_order != None:
                tt = []
                if self.ob_generator[i][1].I.id[:3] == 'GS_':
                    name = self.ob_generator[i][1].I.id[3:]
                else:
                    name = self.ob_generator[i][1].I.id
                for i in range(12):
                    # padding to 12 dims
                    if i in self.ob_order[name].keys():
                        tt.append(tmp[self.ob_order[name][i]])
                    else:
                        tt.append(0.)
                x_obs.append(np.array(tt))     
            else:
                x_obs.append(tmp)
            
        if self.ob_order != None:
            x_obs = np.array(x_obs, dtype=np.float32)
        return x_obs

    def get_reward(self):
        '''
        get_reward
        Get reward from environment.

        :param: None
        :return rewards: rewards generated by reward_generator
        '''
        # TODO: test output
        rewards = []  # sub_agents
        for i in range(len(self.reward_generator)):
            rewards.append(self.reward_generator[i][1].generate())
        rewards = np.squeeze(np.array(rewards))
        return rewards

    def get_phase(self):
        '''
        get_phase
        Get current phase of intersection(s) from environment.

        :param: None
        :return phase: current phase generated by phase_generator, the shape is [sub_agents,]
        '''
        phase = []  # sub_agents
        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):
        '''
        get_queue
        Get queue length of intersection.

        :param: None
        :return queue: value(one intersection) or [intersections,](multiple intersections)
        '''
        queue = []
        for i in range(len(self.queue)):
            queue.append((self.queue[i][1].generate()))
        maxsize = max(len(x) for x in queue)
        for i in range(len(queue)):
            queue[i] = np.pad(queue[i], (0, maxsize - queue[i].shape[-1]))
        tmp_queue = np.squeeze(np.array(queue))
        if self.sub_agents == 1:
            queue = np.sum(tmp_queue)
        else:
            queue = [np.sum(x) for x in tmp_queue]
        # queue = np.sum(tmp_queue, axis=1 if len(tmp_queue.shape)==1 else 0)
        return queue # [intersections,]

    def get_delay(self):
        '''
        get_delay
        Get delay of intersection.

        :param: None
        :return: total delay
        '''
        delay = []
        for i in range(len(self.delay)):
            delay.append((self.delay[i][1].generate()))
        delay = np.squeeze(np.array(delay))
        return delay # [intersections,]

    def get_action(self, ob, phase, test=False):
        '''
        get_action
        Generate action.

        :param ob: current observation, the shape is [agents, ob_length]
        :param phase: current phase, the shape is [agents,]
        :param test: boolean, decide whether it is a training or testing process
        :return acts: action taken by environment, the shape is [agents,]
        '''

        # input are np.array here
        # TODO: support irregular input in the future

        # concate ob and phase
        if self.phase:
            if self.one_hot:
                obs = np.concatenate([utils.idx2onehot(phase, self.action_space.n), ob], axis=1)
            else:
                obs = [np.concatenate([np.array([phase[i]]), ob[i]]) for i in range(self.sub_agents)]
        else:
            obs = ob

        if self.reverse_valid is None and self.valid_acts is not None:
            self.reverse_valid = dict()
            for signal_id in self.valid_acts:
                self.reverse_valid[signal_id] = {v: k for k, v in self.valid_acts[signal_id].items()}

        batch_obs = obs
        if self.valid_acts is None:
            batch_valid = None
            batch_reverse = None
        else:
            # TODO debug
            # batch_valid = [self.valid_acts.get(agent_id) for agent_id in ob.keys()]
            # batch_reverse = [self.reverse_valid.get(agent_id) for agent_id in ob.keys()]
            dic = {index: key for index, key in enumerate(self.ob_order.keys())}
            batch_valid = [self.valid_acts.get(dic[i]) if dic[i] in self.valid_acts.keys() else self.valid_acts.get('GS_'+dic[i]) for i in range(ob.shape[0])]
            batch_reverse = [self.reverse_valid.get(dic[i]) if dic[i] in self.reverse_valid.keys() else self.reverse_valid.get('GS_'+dic[i])  for i in range(ob.shape[0])]
        batch_acts = self.agents_iner.act(batch_obs,
                                valid_acts=batch_valid,
                                reverse_valid=batch_reverse, test=test)
        acts = np.array(batch_acts)
        return acts
        

    def sample(self):
        pass
        # """Applicable to various traffic light plans at each intersection."""
        # ran_phase = []
        # for x in self.action_space_list:
        #     ran_phase.append(np.random.randint(0, x.n))
        # return ran_phase

    def remember(self, last_obs, last_phase, actions, actions_prob, rewards, obs, cur_phase, done, key):
        '''
        remember
        Put current step information into replay buffer for training agent later. 
        pfrl.dqn have automatically implemented this part, so there is no need to implement it.

        :param last_obs: last step observation
        :param last_phase: last step phase
        :param actions: actions executed by intersections
        :param actions_prob: the probability that the intersections execute the actions
        :param rewards: current step rewards
        :param obs: current step observation
        :param cur_phase: current step phase
        :param done: boolean, decide whether the process is done
        :param key: key to store this record, e.g., episode_step_agentid
        :return: None
        '''
        self.do_observe(obs, cur_phase, rewards, done)
    
    def do_observe(self, ob, phase, reward, done):
        '''
        do_observe
        Reconstruct the samples. Call iner_agent to store data and train.

        :param obs: current step observation
        :param phase: current step phase
        :param reward: current step rewards
        :param done: boolean, decide whether the process is done
        :return: None
        '''
        # TODO how to handle inregular lane length?
        # # padding ob length
        # max_len = max([len(x) for x in ob])
        # ob_uni = list(map(lambda l:list(l) + [0.]*(max_len - len(l)), ob))
        # last_ob_uni = list(map(lambda l:list(l) + [0.]*(max_len - len(l)), last_ob))

        # concate ob and phase
        if self.phase:
            if self.one_hot:
                obs = np.concatenate([utils.idx2onehot(phase, self.action_space.n), ob], axis=1)
            else:
                obs = [np.concatenate([np.array([phase[i]]), ob[i]]) for i in range(self.sub_agents)]
        else:
            obs = ob
        reset = [False] * self.sub_agents
        dones = [done] * self.sub_agents if isinstance(done, bool) else done
        rewards = [reward] if isinstance(reward, float) else reward
        self.agents_iner.observe(obs, rewards, dones, reset)

    def _build_model(self):
        '''
        _build_model
        Build a FRAP agent and optimizer.

        :param: None
        :return agents: MPLight_InerAgent consists of agent and optimizer
        '''
        self.model = FRAP(self.dic_agent_conf, self.num_actions, self.phase_pairs, self.comp_mask)
        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=self.learning_rate)
        agents = MPLight_InerAgent(self.model, self.optimizer, self.replay_buffer, self.gamma, self.explorer,
                        minibatch_size=self.batch_size, replay_start_size=self.batch_size, 
                        phi=lambda x: np.asarray(x, dtype=np.float32),
                        # TODO check what is TARGET_UPDATE, target_update_interval, update_interval
                        target_update_interval=self.dic_agent_conf.param["target_update"]*self.sub_agents, update_interval=self.sub_agents
                        )
        return agents

    def update_target_network(self):
        '''
        update_target_network
        Update params of target network. 
        pfrl.dqn have automatically implemented this part, so there is no need to implement it.

        :param: None
        :return: None
        '''
        pass

    def train(self):
        '''
        train
        Train the agent. 
        pfrl.dqn have automatically implemented this part, so there is no need to implement it.

        :param: None
        :return: value of loss
        '''
        result = self.agents_iner.get_statistics()
        return result[1][1]


    def load_model(self, e):
        '''
        load_model
        Load model params of an episode.

        :param e: specified episode
        :return: None
        '''
        model_name = os.path.join(Registry.mapping['logger_mapping']['path'].path,
                                  'model', f'{e}_{self.rank}.pt')
        self.agents_iner = self._build_model()
        # self.agents_iner.load_state_dict(torch.load(model_name))
        tmp_dict = {}
        tmp_dict.load_state_dict(torch.load(model_name))

    def save_model(self, e):
        '''
        save_model
        Save model params of an episode.

        :param e: specified episode, used for file name
        :return: None
        '''
        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)
        torch.save({
            'model_state_dict': self.model.state_dict(),
            'optimizer_state_dict': self.optimizer.state_dict(),
        }, model_name)

class MPLight_InerAgent(DQN):
    '''
    MPLight_InerAgent is based on FRAP and integrates pressure into state and reward design.
    '''
    def __init__(self, q_function, optimizer,replay_buffer, 
                 gamma, explorer, minibatch_size, replay_start_size, phi, 
                 target_update_interval, update_interval):
        super().__init__(q_function, optimizer, replay_buffer, gamma, explorer,
                         minibatch_size=minibatch_size, replay_start_size=replay_start_size, phi=phi,
                         target_update_interval=target_update_interval, update_interval=update_interval)
        # self.batch_last_state = None
        # self.batch_last_action = None

    def act(self, obs, valid_acts=None, reverse_valid=None, test=False):
        '''
        act
        Get action.

        :param obs: observation
        :param valid_acts: directory recording valid phases and the corresponding order of each intersection
        :param reverse_valid: reverse the key and value in valid_acts
        :param test: boolean, decide whether is test process
        :return: action taken by environment
        '''
        return self.batch_act(obs, valid_acts=valid_acts, reverse_valid=reverse_valid, test=test)

    def observe(self, obs, reward, done, reset):
        '''
        observe
        Call function of pfrl.dqn to store data and train.

        :param obs: current step observation
        :param reward: current step rewards
        :param done: boolean, decide whether the process is done
        :param reset: boolean, decide whether to reset the process
        :return: None
        '''
        # self.training = True
        self.batch_observe(obs, reward, done, reset)

    def batch_act(self, batch_obs, valid_acts=None, reverse_valid=None, test=False):
        '''
        batch_act
        Get action.

        :param batch_obs: observation
        :param valid_acts: directory recording valid phases and the corresponding order of each intersection
        :param reverse_valid: reverse the key and value in valid_acts
        :param test: boolean, decide whether is test process
        :return valid_batch_action: valid action taken by environment in batch form
        '''
        # self.training = not test
        if valid_acts is None: 
            return super(MPLight_InerAgent, self).batch_act(batch_obs)
        with torch.no_grad(), evaluating(self.model):
            batch_av = self._evaluate_model_and_update_recurrent_states(batch_obs)
            # TODO do not use gpu
            batch_qvals = batch_av.params[0].detach().cpu().numpy() # shape: [sub_agents, 1+ob_length(remove right_lane?)]
            batch_argmax = []
            for i in range(len(batch_obs)):
                batch_item = batch_qvals[i]
                max_val, max_idx = None, None
                for idx in valid_acts[i]:
                    batch_item_qval = batch_item[idx]
                    if max_val is None:
                        max_val = batch_item_qval
                        max_idx = idx
                    elif batch_item_qval > max_val:
                        max_val = batch_item_qval
                        max_idx = idx
                batch_argmax.append(max_idx)
            batch_argmax = np.asarray(batch_argmax)

        if self.training:
            batch_action = []
            for i in range(len(batch_obs)):
                av = batch_av[i : i + 1]
                greed = batch_argmax[i]
                act, greedy = self.explorer.select_action(self.t, lambda: greed, action_value=av, num_acts=len(valid_acts[i]))
                if not greedy:
                    act = reverse_valid[i][act] # get the real action(12 dims)
                batch_action.append(act)

            self.batch_last_obs = list(batch_obs)
            self.batch_last_action = list(batch_action)
        else:
            batch_action = batch_argmax

        valid_batch_action = []
        for i in range(len(batch_action)):
            valid_batch_action.append(valid_acts[i][batch_action[i]])
        return valid_batch_action        

        
class FRAP(nn.Module):
    '''
    FRAP captures the phase competition relation between traffic movements through a modified network structure.
    '''
    def __init__(self, dic_agent_conf, output_shape, phase_pairs, competition_mask):
        super(FRAP, self).__init__()
        self.oshape = output_shape
        self.phase_pairs = phase_pairs
        self.comp_mask = competition_mask
        self.demand_shape = dic_agent_conf.param['demand_shape']      # Allows more than just queue to be used
        self.one_hot = dic_agent_conf.param['one_hot']
        self.d_out = 4      # units in demand input layer
        self.p_out = 4      # size of phase embedding
        self.lane_embed_units = 16
        relation_embed_size = 4

        self.p = nn.Embedding(2, self.p_out)
        self.d = nn.Linear(self.demand_shape, self.d_out)

        self.lane_embedding = nn.Linear(self.p_out + self.d_out, self.lane_embed_units)

        self.lane_conv = nn.Conv2d(2*self.lane_embed_units, 20, kernel_size=(1, 1))

        self.relation_embedding = nn.Embedding(2, relation_embed_size)
        self.relation_conv = nn.Conv2d(relation_embed_size, 20, kernel_size=(1, 1))

        self.hidden_layer = nn.Conv2d(20, 20, kernel_size=(1, 1))
        self.before_merge = nn.Conv2d(20, 1, kernel_size=(1, 1))

        self.head = DiscreteActionValueHead()

    def forward(self, states):
        '''
        states: [agents, ob_length]
        ob_length:concat[len(one_phase),len(intersection_lane)]
        '''
        num_movements = int((states.size()[1]-1)/self.demand_shape) if not self.one_hot else int((states.size()[1]-len(self.phase_pairs))/self.demand_shape)
        batch_size = states.size()[0]
        acts = states[:, 0].to(torch.int64) if not self.one_hot else states[:, :len(self.phase_pairs)].to(torch.int64)
        states = states[:, 1:] if not self.one_hot else states[:, len(self.phase_pairs):]
        states = states.float()

        # Expand action index to mark demand input indices
        extended_acts = []
        if not self.one_hot:
            for i in range(batch_size):
                act_idx = acts[i]
                pair = self.phase_pairs[act_idx]
                zeros = torch.zeros(num_movements, dtype=torch.int64)
                zeros[pair[0]] = 1
                zeros[pair[1]] = 1
                extended_acts.append(zeros)
            extended_acts = torch.stack(extended_acts)
        else:
            extended_acts = acts
        phase_embeds = torch.sigmoid(self.p(extended_acts))

        phase_demands = []
        for i in range(num_movements):
            phase = phase_embeds[:, i]  # size 4
            demand = states[:, i:i+self.demand_shape] # order by NESW_RTL
            demand = torch.sigmoid(self.d(demand))    # size 4
            phase_demand = torch.cat((phase, demand), -1)
            phase_demand_embed = F.relu(self.lane_embedding(phase_demand))
            phase_demands.append(phase_demand_embed)
        phase_demands = torch.stack(phase_demands, 1)
        # phase_demands_old = torch.stack(phase_demands, 1)
        # # turn direction from NESW to ESWN
        # if num_movements == 8:
        #     phase_demands = torch.cat([phase_demands_old[:,2:,:],phase_demands_old[:,:2,:]],1)
        # elif num_movements == 12:
        #     phase_demands = torch.cat([phase_demands_old[:,3:,:],phase_demands_old[:,:3,:]],1)
        

        pairs = []
        for pair in self.phase_pairs:
            pairs.append(phase_demands[:, pair[0]] + phase_demands[:, pair[1]])

        rotated_phases = []
        for i in range(len(pairs)):
            for j in range(len(pairs)):
                if i != j: rotated_phases.append(torch.cat((pairs[i], pairs[j]), -1))
        rotated_phases = torch.stack(rotated_phases, 1)
        rotated_phases = torch.reshape(rotated_phases,
                                       (batch_size, self.oshape, self.oshape - 1, 2 * self.lane_embed_units))
        rotated_phases = rotated_phases.permute(0, 3, 1, 2)  # Move channels up
        rotated_phases = F.relu(self.lane_conv(rotated_phases))  # Conv-20x1x1  pair demand representation

        # Phase competition mask
        competition_mask = self.comp_mask.repeat((batch_size, 1, 1))
        relations = F.relu(self.relation_embedding(competition_mask))
        relations = relations.permute(0, 3, 1, 2)  # Move channels up
        relations = F.relu(self.relation_conv(relations))  # Pair demand representation

        # Phase pair competition
        combine_features = rotated_phases * relations
        combine_features = F.relu(self.hidden_layer(combine_features))  # Phase competition representation
        combine_features = self.before_merge(combine_features)  # Pairwise competition result

        # Phase score
        combine_features = torch.reshape(combine_features, (batch_size, self.oshape, self.oshape - 1))
        q_values = torch.sum(combine_features, dim=-1)
        return self.head(q_values)
" (mplight)