text: import math
  import random
  
  import networkx as nx
  import numba
  import numpy as np
  import pygame.image
  import copy
  from sko.GA import GA
  
  from scenarios.continuous_uav_base.agents import AgentGroup2D
  from utils.hparams import hparams
  from scenarios.continuous_uav_base.render_utils import ContinuousWorldRenderer
  
  N_AGENT = hparams['env_num_agent']
  N_POI = hparams['env_num_poi']
  N_MAJOR = hparams['env_num_major_point']
  MAJOR_POINT_MIN_DISTANCE_PERCENT = hparams['env_major_point_min_distant_percent']
  N_OBSTACLE_RECT = hparams['env_num_rect_obstacle']
  N_OBSTACLE_CIRCLE = hparams['env_num_circle_obstacle']
  NUM_GRIDS = hparams['env_num_grid']
  COVER_RANGE = hparams['env_coverage_range']
  OBS_RANGE = hparams['env_observe_range']
  COMM_RANGE = hparams['env_comm_range']
  COLLIDE_RANGE = hparams['env_collide_range']
  
  # non-pixel configs
  MAX_NEIGHBOR_AGENTS = hparams['env_num_max_neighbor_agent']
  MAX_NEIGHBOR_POI = hparams['env_num_max_neighbor_poi']
  SQUARE_OBS = hparams['env_use_square_obs']
  PIXEL_OBS = hparams['env_use_pixel_obs']
  STATIC_POI = hparams['env_static_poi']
  ENABLE_OBSTACLE = hparams['env_enable_obstacle']
  ENABLE_CHARGER = hparams['env_enable_charger']
  N_CHARGER = hparams['env_num_charger']
  BATTERY_CAP = hparams['env_battery_capacity']
  OBS_TEST = hparams['env_observation_test']
  
  
  def see_circle(x, y, x_offsets, y_offsets, map0, map1, map2):
      map0 = map0[x + x_offsets, y + y_offsets]
      map1 = map1[x + x_offsets, y + y_offsets]
      map2 = map2[x + x_offsets, y + y_offsets]
      return np.concatenate([map0, map1, map2])
  
  
  def see_local(x, y, agent_map, poi_map, obstacle_map, square=SQUARE_OBS, local_indices=None):
      if square:
          raw_agent_obs_map = agent_map[x - OBS_RANGE:x + OBS_RANGE + 1, y - OBS_RANGE:y + OBS_RANGE + 1]
          reshape_agent_obs_map = raw_agent_obs_map.reshape(
              [3, (2 * OBS_RANGE + 1) // 3, (2 * OBS_RANGE + 1) // 3, 3])
          reduced_agent_obs_map = reshape_agent_obs_map.sum(axis=0).sum(axis=-1).reshape([-1])
  
          raw_poi_obs_map = poi_map[x - OBS_RANGE:x + OBS_RANGE + 1, y - OBS_RANGE:y + OBS_RANGE + 1]
          reshape_poi_obs_map = raw_poi_obs_map.reshape([3, (2 * OBS_RANGE + 1) // 3, (2 * OBS_RANGE + 1) // 3, 3])
          reduced_poi_obs_map = reshape_poi_obs_map.sum(axis=0).sum(axis=-1).reshape([-1])
  
          raw_obstacle_obs_map = obstacle_map[x - OBS_RANGE:x + OBS_RANGE + 1, y - OBS_RANGE:y + OBS_RANGE + 1]
          reshape_obstacle_obs_map = raw_obstacle_obs_map.reshape(
              [3, (2 * OBS_RANGE + 1) // 3, (2 * OBS_RANGE + 1) // 3, 3])
          reduced_obstacle_obs_map = reshape_obstacle_obs_map.sum(axis=0).sum(axis=-1).reshape([-1])
  
          return np.concatenate([reduced_agent_obs_map, reduced_poi_obs_map, reduced_obstacle_obs_map])
      else:
          local_obs = see_circle(x, y, local_indices[0], local_indices[1], agent_map, poi_map, obstacle_map)
          return local_obs
  
  
  @numba.njit
  def cal_two_point_distance(x1, y1, x2, y2):
      return ((x1 - x2) ** 2 + (y1 - y2) ** 2) ** 0.5
  
  
  @numba.njit
  def decimal2binary_arr(x, length=10):
      h = []
      for _ in range(length):  # int类型的位置信息的二进制编码
          h.append(x % 2)
          x = x // 2
      return np.array(h)
  
  
  @numba.njit
  def cal_cover_num(poi_pos, group, agent_pos, i):
      individual_covered = 0
      total_covered = 0
      for j, (poi_x, poi_y) in enumerate(poi_pos):
          covered_by_i = False
          covered_by_others = False
          for agent_id in group:
              agent_x, agent_y = agent_pos[agent_id]
              if cal_two_point_distance(agent_x, agent_y, poi_x, poi_y) < COVER_RANGE:
                  if i == agent_id:
                      covered_by_i = True
                  else:
                      covered_by_others = True
          if covered_by_i and not covered_by_others:
              individual_covered += 1
          if covered_by_i or covered_by_others:
              total_covered += 1
      return individual_covered, total_covered
  
  
  @numba.njit
  def cal_poi_is_covered(poi_is_covered, poi_pos, agent_pos):
      for i, (poi_x, poi_y) in enumerate(poi_pos):
          for j, (agent_x, agent_y) in enumerate(agent_pos):
              if cal_two_point_distance(poi_x, poi_y, agent_x, agent_y) <= COVER_RANGE:
                  poi_is_covered[i] = 1
                  break
      return poi_is_covered
  
  
  @numba.njit
  def in_rectangle(x, y, min_x, min_y, height, width):
      return x > min_x and x < min_x + height and y > min_y and y < min_y + width
  
  
  @numba.njit
  def in_circle(x, y, x_core, y_core, radius):
      distance_to_core = ((x - x_core) ** 2 + (y - y_core) ** 2) ** 0.5
      return distance_to_core < radius
  
  
  @numba.njit
  def cal_obstacle_map(obstacle_map, rect_obstacles, circle_obstacles):
      world_scale = obstacle_map.shape
      for x in range(world_scale[0]):
          for y in range(world_scale[1]):
              if x < OBS_RANGE or x > NUM_GRIDS + OBS_RANGE:
                  obstacle_map[x, y] = 1
                  continue
              if y < OBS_RANGE or y > NUM_GRIDS + OBS_RANGE:
                  obstacle_map[x, y] = 1
                  continue
              for (x_min, y_min, height, width) in rect_obstacles:
                  if in_rectangle(x, y, x_min, y_min, height, width):
                      obstacle_map[x, y] = 1
                      break
              for (x_core, y_core, radius) in circle_obstacles:
                  if in_circle(x, y, x_core, y_core, radius):
                      obstacle_map[x, y] = 1
                      break
      return obstacle_map
  
  
  @numba.njit
  def generate_n_agent_not_in_map(n, map):
      out = []
      for _ in range(n):
          while True:
              x = random.random() * 0.8 + 0.1
              y = random.random() * 0.8 + 0.1
              x_idx = int(x * NUM_GRIDS + OBS_RANGE)
              y_idx = int(y * NUM_GRIDS + OBS_RANGE)
              if map[x_idx, y_idx] == 0:
                  out.append([x, y])
                  break
      return np.array(out, dtype=np.float32)
  
  
  @numba.njit
  def clip(x, x_min, x_max):
      if x > x_max:
          return x_max
      if x < x_min:
          return x_min
      return x
  
  
  @numba.njit
  def generate_poi_not_in_map(n_major, n_poi, map):
      pos_major = []
      for i in range(n_major):
          while True:
              x = random.random() * 0.7 + 0.15
              y = random.random() * 0.7 + 0.15
              x_idx = int(x * NUM_GRIDS + OBS_RANGE)
              y_idx = int(y * NUM_GRIDS + OBS_RANGE)
              if map[x_idx, y_idx] == 0:
                  if len(pos_major) > 0:
                      not_good = False
                      for pos in pos_major:
                          if cal_two_point_distance(x, y, pos[0], pos[1]) < MAJOR_POINT_MIN_DISTANCE_PERCENT:
                              not_good = True
                              break
                      if not_good:
                          continue
                      pos_major.append([x, y])
                      break
                  else:
                      pos_major.append([x, y])
                      break
      pos_major = np.array(pos_major)
      poi_pos = []
      for i in range(n_major):
          num_scattering_poi = n_poi // n_major
          pos_major_i = pos_major[i]
          x_major_i, y_major_i = pos_major_i
          for j in range(num_scattering_poi):
              while True:
                  x_diff = np.random.randn(1)[0] * 0.1
                  x_poi = clip(x_major_i + x_diff, 0, 1)
  
                  # y_diff = (random.random() - 0.5) * 2 * 0.3
                  y_diff = np.random.randn(1)[0] * 0.1
                  y_poi = clip(y_major_i + y_diff, 0, 1)
                  x_index = int(x_poi * NUM_GRIDS + OBS_RANGE)
                  y_index = int(y_poi * NUM_GRIDS + OBS_RANGE)
                  if map[x_index, y_index] == 0:
                      poi_pos.append([x_poi, y_poi])
                      break
      poi_pos = np.array(poi_pos, dtype=np.float32)
      return pos_major, poi_pos
  
  
  @numba.njit
  def generate_charger_not_in_map(n_charger, map, pos_major):
      out = []
      assert n_charger == len(pos_major)
      for i in range(len(pos_major)):
          pos_major_i = pos_major[i]
          x_major_i, y_major_i = pos_major_i
          while True:
              x_diff = np.random.randn(1)[0] * 0.05
              x_poi = clip(x_major_i + x_diff, 0, 1)
              y_diff = np.random.randn(1)[0] * 0.05
              y_poi = clip(y_major_i + y_diff, 0, 1)
              x_idx = int(x_poi * NUM_GRIDS + OBS_RANGE)
              y_idx = int(y_poi * NUM_GRIDS + OBS_RANGE)
              if map[x_idx, y_idx] == 0:
                  out.append([x_poi, y_poi])
                  break
      return np.array(out, dtype=np.float32)
  
  
  class AgentUAVMBS(AgentGroup2D):
      def __init__(self, n_agent, pos=None):
          self.mass = 1
          super(AgentUAVMBS, self).__init__(n_agent, pos=pos, vel=None, mass=self.mass)
          self.n_agent = n_agent
          self.max_force = 0.04  # 1m/s^2
          self.max_vel = 0.02  # 10m/s
          self.hover_energy_cost = 0.5
          self.move_energy_cost = 0.5
          self.energy_consumption = np.ones([n_agent]) * self.hover_energy_cost
          self.action2force = self.max_force * np.array([
              [1, 0],  # action 0 denotes 1 unit force downward
              [math.sqrt(2) / 2, -math.sqrt(2) / 2],  # left down side
              [0, -1],  # leftward
              [-math.sqrt(2) / 2, -math.sqrt(2) / 2],  # left upwards
              [-1, 0],  # upward
              [-math.sqrt(2) / 2, math.sqrt(2) / 2],  # right upward
              [0, 1],  # rightward
              [math.sqrt(2) / 2, math.sqrt(2) / 2],  # right downward
  
              [0.5, 0],  # action 0 denotes 1 unit force downward
              [math.sqrt(2) / 4, -math.sqrt(2) / 4],  # left down side
              [0, -0.5],  # leftward
              [-math.sqrt(2) / 4, -math.sqrt(2) / 4],  # left upwards
              [-0.5, 0],  # upward
              [-math.sqrt(2) / 4, math.sqrt(2) / 4],  # right upward
              [0, 0.5],  # rightward
              [math.sqrt(2) / 4, math.sqrt(2) / 4],  # right downward
              [0, 0]  # no throttle
          ])
          self.n_action = self.action2force.shape[0]
  
      def deploy_action(self, actions, obstacle_map, batteries, time_step=1):
          assert isinstance(actions, np.ndarray)
          actions = actions.reshape([-1]).astype(int)
          assert actions.size == self.n_agent
          force = self.action2force[actions]      # [n_agent, 2]
          acceleration = force / self.mass        # 加速度的取值范围
  
          # Update UAV Battery Status, using values of acceleration
          for i in range(self.n_agent):
              accel_value = (acceleration[i] ** 2).sum(axis=0) ** 0.5
              accel_value = (acceleration[i] ** 2).sum(axis=0) ** 0.5
              if abs(accel_value) < 0.00001:
                  batteries[i] -= 2
              elif abs(accel_value - 0.02) < 0.00001:
                  batteries[i] -= 2
              elif abs(accel_value - 0.04) < 0.00001:
                  batteries[i] -= 4
  
          self.vel = self.vel + acceleration * time_step
          vel_norm = (self.vel ** 2).sum(axis=-1, keepdims=True) ** 0.5
          vel_less_than_threshold_mask = np.array(vel_norm < self.max_vel, dtype=np.float32)
          vel_greater_than_threshold_mask = np.array(vel_norm > self.max_vel, dtype=np.float32)
          vel_greater_than_threshold_mask = vel_greater_than_threshold_mask * (
                  self.max_vel * np.ones_like(vel_norm) / (vel_norm + 1e-5))
          clip_vel_mask = vel_less_than_threshold_mask + vel_greater_than_threshold_mask
          self.vel = self.vel * clip_vel_mask
          clipped_vel_norm = (self.vel ** 2).sum(axis=-1, keepdims=True) ** 0.5
          self.energy_consumption = np.ones([self.n_agent]) * self.hover_energy_cost + (
                  clipped_vel_norm / self.max_vel * self.move_energy_cost).reshape([-1])
  
          collision_penalty = np.zeros([self.n_agent], dtype=np.float32)
          previous_pos = self.pos
          self.pos = self.pos + self.vel * time_step
          self.pos = np.clip(self.pos, 0, 1)
          for i, (x_i, y_i) in enumerate(self.pos):
              x_i = int(x_i * NUM_GRIDS + OBS_RANGE)
              y_i = int(y_i * NUM_GRIDS + OBS_RANGE)
              if obstacle_map[x_i, y_i] != 0:
                  self.pos[i] = previous_pos[i]
                  collision_penalty[i] = -1
          return collision_penalty, batteries
  
  
  class ContinuousMCS:
      def __init__(self, n_agent=N_AGENT, world_scale=(NUM_GRIDS + 2 * OBS_RANGE, NUM_GRIDS + 2 * OBS_RANGE),
                   heuristic_reward=True):
          self.world_scale = np.array(world_scale, dtype=np.int32)
          self.n_agent = n_agent
          self.pixel_obs = PIXEL_OBS
          self.static_poi = STATIC_POI
          self.square_obs = SQUARE_OBS
          self.heuristic_reward = heuristic_reward
          self.obs_test = OBS_TEST  # For Observation Testing Case
          self.agents = None
          self.agent_map = np.zeros(world_scale)
          self.num_major_points = N_MAJOR
          self.num_poi = N_POI
          self.major_point_pos = np.zeros([self.num_major_points, 2])
          self.poi_pos = np.zeros([self.num_major_points, 2])
          self.poi_map = np.zeros(world_scale)
          self.num_rect_obstacles = N_OBSTACLE_RECT if ENABLE_OBSTACLE else 0
          self.num_circle_obstacles = N_OBSTACLE_CIRCLE if ENABLE_OBSTACLE else 0
          self.obstacle_max_width = 10
          self.rect_obstacles = np.zeros([self.num_circle_obstacles, 3])  # represent as (x,y,r)
          self.circle_obstacles = np.zeros([self.num_rect_obstacles, 4])  # (x_min,y_min,x_max,y_max)
          self.obstacle_map = np.zeros(world_scale)
          self.local_indices_x = []
          self.local_indices_y = []
          for offset_x in range(-OBS_RANGE, OBS_RANGE):
              for offset_y in range(-OBS_RANGE, OBS_RANGE):
                  if (offset_x ** 2 + offset_y ** 2) ** 0.5 <= OBS_RANGE:
                      self.local_indices_x.append(offset_x)
                      self.local_indices_y.append(offset_y)
          self.local_indices_x = np.array(self.local_indices_x)
          self.local_indices_y = np.array(self.local_indices_y)
          if self.pixel_obs:
              if self.square_obs:
                  self.obs_dim = 20 + 2 + 3 * ((2 * OBS_RANGE + 1) // 3) ** 2
              else:
                  self.obs_dim = 30 + 2 + 3 * len(self.local_indices_x) + 2 + 1
          else:
              self.obs_dim = 4 + 4 * MAX_NEIGHBOR_AGENTS + 2 * MAX_NEIGHBOR_POI
          self.act_dim = None
  
          # Properties for Chargers and UAV Batteries
          self.enable_charger = ENABLE_CHARGER
          self.num_charger = N_CHARGER if ENABLE_CHARGER else 0
          self.charger_pos = np.zeros([self.num_charger, 2])
          self.charger_map = np.zeros(world_scale)
          self.is_charging = np.zeros(self.n_agent)
          self.batteries = np.full(self.n_agent, BATTERY_CAP, dtype=int)  # 每一轮episode都需要重新reset()
  
          self.poi_coverage_history = []
          self.energy_consumption_history = []
          self.num_neighbor_history = []
          self.timeslot = 0
          self.reset()
          self.renderer = None
  
      def render(self):
          if self.renderer is None:
              self.renderer = ContinuousWorldRenderer(num_grid=NUM_GRIDS, obs_range=OBS_RANGE, fps=10, scale_factor=2.5)
              self.agent_surface = pygame.image.load("scenarios/continuous_uav_base/UAV.png")
          render_info = self.get_render_info()
          self.renderer.render_mcs(render_info, agent_surface=self.agent_surface)
  
      def reset(self):
          self.timeslot = 0
          # generate rectangle obstacle
          xy_min = (np.random.rand(self.num_rect_obstacles, 2) * 0.8 + 0.1) * NUM_GRIDS + OBS_RANGE
          height_width = (np.random.rand(self.num_rect_obstacles, 2) + 1) / 2 * self.obstacle_max_width
  
          self.rect_obstacles = np.concatenate([xy_min, height_width], axis=-1)  # [n_obstacle, 4]
  
          # generate circle obstacle
          xy = (np.random.rand(self.num_circle_obstacles, 2) * 0.8 + 0.1) * NUM_GRIDS + OBS_RANGE
          radius = (np.random.rand(self.num_circle_obstacles, 1) + 1) / 2 * 0.5 * self.obstacle_max_width
          self.circle_obstacles = np.concatenate([xy, radius], axis=-1)
          self.refresh_obstacle_map()
  
          pos_major, poi_pos = generate_poi_not_in_map(n_major=self.num_major_points, n_poi=self.num_poi,
                                                       map=self.obstacle_map)
          self.major_point_pos = pos_major
          self.poi_pos = poi_pos
          self.refresh_poi_map()
  
          agent_pos = generate_n_agent_not_in_map(self.n_agent, self.obstacle_map)
          self.agents = AgentUAVMBS(self.n_agent, agent_pos)
          self.act_dim = self.agents.n_action
          self.refresh_agent_map()
  
          # generate charging stations
          pos_charger = generate_charger_not_in_map(self.num_charger, self.obstacle_map, pos_major)
          self.charger_pos = pos_charger
          self.refresh_charger_map()
  
          # reset UAV battery for every episode
          self.batteries = np.full(self.n_agent, BATTERY_CAP, dtype=int)
  
          adj = self.get_adj()
          obs = self.get_obs()
  
          self.poi_coverage_history = []
          self.energy_consumption_history = []
          self.num_neighbor_history = []
          self.poi_is_covered = cal_poi_is_covered(np.zeros([self.num_poi]), poi_pos, agent_pos)
  
          return obs, adj
  
      def init_from_map(self, poi_map, agent_map):
          self.reset()
          self.poi_pos = []
          agent_pos_lst = []
          for grid_x in range(poi_map.shape[0]):
              for grid_y in range(poi_map.shape[1]):
                  num_poi = int(poi_map[grid_x][grid_y])
                  for _ in range(num_poi):
                      self.poi_pos.append([grid_x / NUM_GRIDS, grid_y / NUM_GRIDS])
          for grid_x in range(agent_map.shape[0]):
              for grid_y in range(agent_map.shape[1]):
                  num_agent = int(agent_map[grid_x][grid_y])
                  for _ in range(num_agent):
                      agent_pos_lst.append([grid_x / NUM_GRIDS, grid_y / NUM_GRIDS])
          self.num_poi = len(self.poi_pos)
          self.poi_pos = np.array(self.poi_pos)
          agent_pos = np.array(agent_pos_lst).reshape([-1, 2])
          self.n_agent = len(agent_pos)
          self.agents = AgentUAVMBS(n_agent=self.n_agent, pos=agent_pos)
          self.refresh_agent_map()
          self.refresh_poi_map()
          return self.get_obs(), self.get_adj()
  
      def refresh_poi_map(self):
          self.poi_map = np.zeros(self.world_scale, dtype=np.int32)
          self.x_pois = np.clip(self.poi_pos[:, 0] * NUM_GRIDS + OBS_RANGE, OBS_RANGE, NUM_GRIDS + OBS_RANGE - 1).astype(
              np.int32)
          self.y_pois = np.clip(self.poi_pos[:, 1] * NUM_GRIDS + OBS_RANGE, OBS_RANGE, NUM_GRIDS + OBS_RANGE - 1).astype(
              np.int32)
          for x, y in zip(self.x_pois, self.y_pois):
              self.poi_map[x, y] += 1
  
      def refresh_obstacle_map(self):
          self.obstacle_map = cal_obstacle_map(np.zeros(self.world_scale, dtype=np.int32),
                                               self.rect_obstacles, self.circle_obstacles)
  
      def refresh_agent_map(self):
          self.agent_map = np.zeros(self.world_scale, dtype=np.int32)
          self.x_agents = np.clip(self.agents.x * NUM_GRIDS + OBS_RANGE, OBS_RANGE, NUM_GRIDS + OBS_RANGE - 1).astype(
              np.int32)
          self.y_agents = np.clip(self.agents.y * NUM_GRIDS + OBS_RANGE, OBS_RANGE, NUM_GRIDS + OBS_RANGE - 1).astype(
              np.int32)
          for x, y in zip(self.x_agents, self.y_agents):
              self.agent_map[x, y] += 1
  
      def refresh_charger_map(self):
          self.charger_map = np.zeros(self.world_scale, dtype=np.int32)
          self.x_charger = np.clip(self.charger_pos[:, 0] * NUM_GRIDS + OBS_RANGE, OBS_RANGE, NUM_GRIDS + OBS_RANGE - 1) \
              .astype(np.int32)
          self.y_charger = np.clip(self.charger_pos[:, 1] * NUM_GRIDS + OBS_RANGE, OBS_RANGE, NUM_GRIDS + OBS_RANGE - 1) \
              .astype(np.int32)
          for x, y in zip(self.x_charger, self.y_charger):
              self.charger_map[x, y] += 1
  
      def get_mazes(self):
          if not self.pixel_obs:
              print("warning: you are acquiring agent map in no-pixel-obs mode!")
          return self.poi_map, self.agent_map
  
      def get_pos(self):
          return self.poi_pos, self.agents.pos, self.rect_obstacles, self.circle_obstacles
  
      def get_render_info(self):
          render_info = {}
          render_info['poi_pos'] = self.poi_pos
          render_info['agent_pos'] = self.agents.pos
          render_info['rect_obstacles'] = self.rect_obstacles
          render_info['circle_obstacles'] = self.circle_obstacles
          render_info['poi_cover_id'] = self.poi_is_covered
          render_info['poi_cover_percent'] = sum(self.poi_is_covered) / self.num_poi
          c_t, f_t, _, _ = self.cal_episodic_coverage_and_fairness()
          render_info['episode_coverage_item'] = c_t
          render_info['episode_fairness_item'] = f_t
          render_info['energy_consumption'] = self.agents.energy_consumption[-1]
          render_info['timeslot'] = self.timeslot
          render_info['charger_pos'] = self.charger_pos
          render_info['battery'] = self.batteries
          return render_info
  
      def get_obs(self):
          if self.pixel_obs:
              obs = []
              for i, (x, y) in enumerate(zip(self.x_agents, self.y_agents)):
                  obs_i = []
                  # UAV Position (binary)
                  binary_x = decimal2binary_arr(int(x), length=10)
                  obs_i.append(binary_x)
                  binary_y = decimal2binary_arr(int(y), length=10)
                  obs_i.append(binary_y)
  
                  # Adjacent Observation: agents, pois & obstacles within UAV OBS_RANGE
                  if self.square_obs:
                      h_local = see_local(x, y, self.agent_map, self.poi_map, self.obstacle_map)
                  else:
                      h_local = see_local(x, y, self.agent_map, self.poi_map, self.obstacle_map,
                                      local_indices=[self.local_indices_x, self.local_indices_y])
                  obs_i.append(h_local)
  
                  # UAV Velocity (float [x_vel, y_vel])
                  vel = self.agents.vel[i] / self.agents.max_vel
                  obs_i.append(vel)
  
                  # UAV ID (binary)
                  binary_id = np.array(decimal2binary_arr(int(i), length=10))
                  obs_i.append(binary_id)
  
                  # UAV Battery (int)
                  bat = self.batteries[i]
                  obs_i.append([bat])
  
                  # Charging Station Observation (positive int)
                  # Relative Position between CLOSEST charging station and agents
                  charge_distance = []
                  for i in range(self.num_charger):
                      xc = self.charger_pos[i][0] * NUM_GRIDS + OBS_RANGE
                      yc = self.charger_pos[i][1] * NUM_GRIDS + OBS_RANGE
                      distance = math.sqrt(pow(xc - x, 2) + (pow(yc - y, 2)))
                      # distance = abs(xc - x) + abs(yc - y)
                      charge_distance.append(distance)
                  closest_charger_index = charge_distance.index(min(charge_distance))
                  xrel = self.charger_pos[closest_charger_index][0] * NUM_GRIDS + OBS_RANGE - x + NUM_GRIDS
                  yrel = self.charger_pos[closest_charger_index][1] * NUM_GRIDS + OBS_RANGE - y + NUM_GRIDS
                  obs_i.append([int(xrel)])
                  obs_i.append([int(yrel)])
  
                  obs_i = np.concatenate(obs_i)
                  obs.append(obs_i)
              obs = np.array(obs)
          else:
              agent_adj = self.adj_t  # [n_agent, n_agent]
              agent_pos = self.agents.pos  # [n_agent, 2]
              agent_vel = self.agents.vel  # [n_agent, 2]
              poi_pos = self.poi_pos  # [n_poi, 2]
              n_poi = self.num_poi
              n_agent = self.n_agent
              normalized_obs_range = OBS_RANGE / NUM_GRIDS
              dummy_value = -3
              n_max_neighbor_agent = MAX_NEIGHBOR_AGENTS
              n_max_neighbor_poi = MAX_NEIGHBOR_POI
  
              tmp_agent_i_pos = np.repeat(agent_pos[:, np.newaxis, :], repeats=n_agent, axis=1)  # [n_agent, n_poi, 2]
              tmp_agent_j_pos = np.repeat(agent_pos[np.newaxis, :, :], repeats=n_agent, axis=0)  # [n_agent, n_poi, 2]
              agent_diff = tmp_agent_j_pos - tmp_agent_i_pos
  
              tmp_agent_pos = np.repeat(agent_pos[:, np.newaxis, :], repeats=n_poi, axis=1)  # [n_agent, n_poi, 2]
              tmp_poi_pos = np.repeat(poi_pos[np.newaxis, :], repeats=n_agent, axis=0)  # [n_agent, n_poi, 2]
              agent_poi_diff = tmp_poi_pos - tmp_agent_pos  # [n_agent, n_poi, 2]
              agent_poi_distance = (agent_poi_diff ** 2).sum(axis=-1) ** 0.5  # [n_agent, n_poi]
              agent_poi_adj = (agent_poi_distance < normalized_obs_range).astype(np.float32)
  
              obs = []
              for i in range(n_agent):
                  obs_i = []
  
                  self_pos = agent_pos[i].reshape([-1])  # 2 bit
                  self_vel = agent_vel[i].reshape([-1])  # 2 bit
                  obs_i.append(self_pos)
                  obs_i.append(self_vel)
  
                  # pos and vel of other_agents: 2*10 + 2*10 = 40 bit
                  neighboring_agents = np.where(agent_adj[i] > 0)[0]
                  # UAV within Comm range is in the obs space!
                  n_insert_agent = 0
                  for j in neighboring_agents:
                      if j != i and n_insert_agent < n_max_neighbor_agent:
                          other_pos_diff = agent_diff[i, j].reshape([-1])  # 2 bit
                          other_vel = agent_vel[j].reshape([-1])
                          obs_i.append(other_pos_diff)
                          obs_i.append(other_vel)
                          n_insert_agent += 1
                  n_agent_padding = max(0, n_max_neighbor_agent - n_insert_agent)
                  obs_i.append(np.ones([n_agent_padding * 2, ]) * dummy_value)  # padding agent pos
                  obs_i.append(np.ones([n_agent_padding * 2, ]) * 0)  # padding agent vel
  
                  # pos of pois: 2*30 = 60 bit
                  neighboring_pois = np.where(agent_poi_adj[i] > 0)[0]
                  # POI within OBS range is in the obs space!
                  n_insert_poi = 0
                  for j in neighboring_pois:
                      if n_insert_poi < n_max_neighbor_poi - 1:
                          other_pos_diff = agent_poi_diff[i, j].reshape([-1])  # 2 bit
                          obs_i.append(other_pos_diff)
                          n_insert_poi += 1
                  n_poi_padding = max(0, n_max_neighbor_poi - n_insert_poi)
                  obs_i.append(np.ones([n_poi_padding * 2, ]) * dummy_value)  # padding agent pos
  
                  obs_i = np.concatenate(obs_i)  # [104, ]
                  obs.append(obs_i)
              obs = np.array(obs)  # [n_agent, 104]
          return obs.reshape([self.n_agent, self.obs_dim])
  
      def get_adj(self):
          pos_agents = np.concatenate([self.x_agents.reshape([-1, 1]),
                                       self.y_agents.reshape([-1, 1])], axis=-1)  # [n_agent, 2]
          pos_agents_i = pos_agents.reshape([self.n_agent, 1, 2]).repeat(self.n_agent, axis=1)
          pos_agents_j = pos_agents.reshape([1, self.n_agent, 2]).repeat(self.n_agent, axis=0)
          distance_mat = ((pos_agents_i - pos_agents_j) ** 2).sum(axis=-1) ** 0.5  # [n_agent, n_agent]
          adj = (distance_mat <= COMM_RANGE).astype(np.float32)
          collide_adj = (distance_mat <= COVER_RANGE * 0.5).astype(np.float32)
          self.collide_agent_num = collide_adj.sum(axis=1).reshape([-1]) - 1
          self.adj_t = adj
          mean_num_neighbor = (adj.sum() - self.n_agent) / self.n_agent
          self.num_neighbor_history.append(mean_num_neighbor)
          return adj.reshape([self.n_agent, self.n_agent])
  
      def step(self, actions):
          collision_penalty, self.batteries = self.agents.deploy_action(actions, self.obstacle_map, self.batteries)
          self.energy_consumption_history.append(self.agents.energy_consumption)
          if self.pixel_obs:
              self.refresh_agent_map()
          adj = self.get_adj()
  
          adj_graph = nx.from_numpy_matrix(adj)
          sub_graphs = tuple(
              adj_graph.subgraph(c).nodes() for c in nx.connected_components(adj_graph))  # ((0,1,2),(3,4,),...)
          reward = []
          id2group_dic = {}
          for i_group, group in enumerate(sub_graphs):
              for id in group:
                  id2group_dic[id] = i_group
          poi_pos = np.concatenate([self.x_pois.reshape([self.num_poi, 1]),
                                    self.y_pois.reshape([self.num_poi, 1])], axis=-1)
          agent_pos = np.concatenate([self.x_agents.reshape([-1, 1]),
                                      self.y_agents.reshape([-1, 1])], axis=-1)  # [n_agent, 2]
  
          self.poi_is_covered = cal_poi_is_covered(np.zeros([self.num_poi]), poi_pos, agent_pos)
          self.poi_coverage_history.append(self.poi_is_covered)
  
          if self.heuristic_reward:
              for i in range(self.n_agent):
                  i_group = id2group_dic[i]
                  group = sub_graphs[i_group]  # [1,4,6,...,55,100]
  
                  individual_covered, total_covered = cal_cover_num(poi_pos, tuple(group), agent_pos, i)
                  individual_rew = -1 if individual_covered == 0 else 1 * individual_covered
                  group_scale = len(group)
                  group_rew = 0. if group_scale == 1 else 0.1 * (total_covered - individual_covered) / (group_scale - 1)
  
                  # Calculate Minimum Charging Distance
                  # self.charger_pos/agents.pos is random(0,1)
                  charge_distance = []
                  for j in range(self.num_charger):
                      distance_euclidean = math.sqrt(pow(self.charger_pos[j][0] - self.agents.pos[i][0], 2) +
                                                     (pow(self.charger_pos[j][1] - self.agents.pos[i][1], 2)))
                      charge_distance.append(distance_euclidean)
                  closest_charger_euclidean = min(charge_distance)
  
                  charge_rew = 1.0 - closest_charger_euclidean
                  if closest_charger_euclidean <= 0.02:
                      charge_rew += 2.0
                      self.batteries[i] = BATTERY_CAP
                      self.is_charging[i] = 1
                  else:
                      self.is_charging[i] = 0
  
                  # Reward Shaping: UAV Battery Usage
                  if self.batteries[i] > 20:
                      rew = group_rew + individual_rew + charge_rew
                      # rew = group_rew + individual_rew
                  else:
                      rew = charge_rew
  
                  # penalty to agents near the boundary
                  collide_percent = COLLIDE_RANGE / NUM_GRIDS
                  if not in_rectangle(self.agents.pos[i][0], self.agents.pos[i][1], collide_percent, collide_percent,
                                      1 - collide_percent, 1 - collide_percent):
                      rew -= 1
                  '''
                  # rew -= self.collide_agent_num[i]
                  # rew /= self.agents.energy_consumption[i]
                  '''
                  reward.append(rew)
          else:
              num_covered = self.poi_is_covered.sum()
              current_c = num_covered / self.num_poi
              f, _, _, _ = self.cal_episodic_coverage_and_fairness()
              global_rew = f * current_c / self.agents.energy_consumption.mean()
              reward = [global_rew] * self.n_agent
  
          rew = np.array(reward)
          rew = rew + collision_penalty
          rew = rew.reshape([self.n_agent, ])
          done = np.zeros([self.n_agent, ])
  
          if not self.static_poi:
              self.poi_pos = np.clip(self.poi_pos + 0.001 * (np.random.randn(self.num_poi, 2)), a_min=0, a_max=1)
          obs = self.get_obs()
          self.timeslot += 1
          return rew, obs, adj, done
  
      def cal_episodic_coverage_and_fairness(self):
          try:
              w_t_k = np.stack(self.poi_coverage_history)  # [T, n_poi]
          except:
              # First time
              return -1, -1, -1, -1
          for i in range(1, len(w_t_k)):
              w_t_k[i] = w_t_k[i] + w_t_k[i - 1]
          max_time = w_t_k.shape[0]
          time_matrix = np.arange(1, max_time + 1).reshape([max_time, 1]).repeat(self.num_poi, axis=-1)  # [T, n_poi]
          c_t_k = w_t_k / time_matrix
          c_t = np.mean(c_t_k, axis=-1)
          final_averaged_coverage_score = c_t[-1]
          f_t = (np.sum(c_t_k, axis=-1) ** 2) / (self.num_poi * np.sum(c_t_k ** 2, axis=-1) + 1e-10)
          final_achieved_fairness_index = f_t[-1]
          return final_averaged_coverage_score, final_achieved_fairness_index, c_t, f_t
  
      def cal_episodic_mean_energy_consumption(self):
          return np.array(self.energy_consumption_history).mean()
  
      def cal_episodic_mean_num_neighbors(self):
          return np.array(self.num_neighbor_history).mean()
  
      def from_copy(self, other):
          for k, v in self.__dict__.items():
              self.__dict__[k] = copy.deepcopy(other.__dict__[k])
  
      def get_log_vars(self):
          """
          get vars to be logged in a episode
          """
          coverage_index, fairness_index, _, _ = self.cal_episodic_coverage_and_fairness()
          energy_index = self.cal_episodic_mean_energy_consumption()
          return {'coverage_index': coverage_index, 'fairness_index': fairness_index, 'energy_index': energy_index}
  
  
  def MB_greedy_policy(env):
      # It is a model-based individual greedy policy.
      # We don't choose actions jointly since it need to numerate joint action space with a size of 5^100
      # For each agent choosing action, we assume other agent choose stay there, then find the action with highest reward
      actions = []
      backup_env = copy.deepcopy(env)
      for agent_i in range(env.n_agent):
          rewards_for_each_actions = np.array([-999] * env.act_dims)
          imagine_env = copy.deepcopy(backup_env)
          for action_i in range(env.act_dims):
              dummy_actions = [env.act_dims - 1] * env.n_agent  # means stop
              dummy_actions[agent_i] = action_i
              _, _, rews, _ = imagine_env.step(np.array(dummy_actions))
              imagine_env.from_copy(backup_env)
              reward = rews[agent_i]
              rewards_for_each_actions[action_i] = reward
          best_action_agent_i = rewards_for_each_actions.argmax()
          actions.append(best_action_agent_i)
      return np.array(actions)
  
  
  def MB_GA_policy(env, max_iter=20):
      # For pop=500, iter=10,
      backup_env = copy.deepcopy(env)
      imagine_env = copy.deepcopy(env)
  
      def loss_func(actions):
          _, _, rewards, _ = imagine_env.step(actions)
          imagine_env.from_copy(backup_env)
          return - sum(rewards)
  
      ga = GA(func=loss_func, n_dim=env.n_agent, size_pop=50, max_iter=max_iter, lb=[0] * env.n_agent,
              ub=[env.act_dims - 1] * env.n_agent, precision=1) \
          # .to('cuda:0')
      joint_actions, joint_rewards = ga.run()
      return np.array(joint_actions)
  