text: from envs.highway import Env
  from vehicle_handler import VehicleHandlerMerge, VehicleHandlerMergeMC
  
  class EnvMC(Env):
      def __init__(self, config):
          super().__init__(config)
          self.real_vehicles = []
          self.ima_vehicles = []
          self.real_mc_log = {}
          self.ima_mc_log = {}
          self.collision_detected = False
          self.debugging_mode = False
          self.metric_collection_mode = False
  
      def prohibit_lane_change(self, vehicle):
          """
          For cars to be modelled for MC
          """
          def _act(self, reservations):
              act_long = self.idm_action(self, self.neighbours['att'])
              return [max(-3, min(act_long, 3)), 0]
          vehicle.act = types.MethodType(_act, vehicle)
  
      def idm_to_neural_vehicle(self, vehicle):
          neural_vehicle = copy.deepcopy(self.neural_vehicle)
          for attrname, attrvalue in list(vehicle.__dict__.items()):
              if attrname != 'act':
                  setattr(neural_vehicle, attrname, copy.copy(attrvalue))
          return neural_vehicle
  
      def add_new_vehicles(self):
          new_entries = self.handler.handle_vehicle_entries(self.queuing_entries,
                                                            self.last_entries)
          for vehicle in new_entries:
              imagined_vehicle = copy.deepcopy(vehicle)
              imagined_vehicle.collision_detected = False
              self.ima_vehicles.append(imagined_vehicle)
              self.real_vehicles.append(vehicle)
  
      def set_ima_veh_decision(self, veh_real, veh_ima):
          for attrname in ['lane_decision', 'lane_y', 'target_lane']:
              attrvalue = getattr(veh_real, attrname)
              setattr(veh_ima, attrname, attrvalue)
  
      def set_ima_veh_neighbours(self, veh_real, veh_ima):
          for key, neighbour in veh_real.neighbours.items():
              if neighbour:
                  for veh in self.ima_vehicles:
                      if veh.id == neighbour.id:
                          veh_ima.neighbours[key] = veh
              else:
                  veh_ima.neighbours[key] = None
  
      def get_joint_action(self):
          """
          Returns the joint action of all vehicles on the road
          """
          acts_real = []
          acts_ima = []
          for veh_real, veh_ima in zip(self.real_vehicles, self.ima_vehicles):
              # real vehicles
              veh_real.neighbours = veh_real.my_neighbours(self.real_vehicles)
              act_long, act_lat = veh_real.act(self.handler.reservations)
              acts_real.append([act_long, act_lat])
              veh_real.act_long = act_long
              self.handler.update_reservations(veh_real)
  
              # imagined vehicles
              self.set_ima_veh_neighbours(veh_real, veh_ima)
              if veh_ima.vehicle_type == 'neural':
                  obs = veh_ima.neur_observe(veh_ima, veh_ima.neighbours['f'], \
                                                          veh_ima.neighbours['m'])
                  if not veh_ima.collision_detected:
                      veh_ima.update_obs_history(obs[0])
                      if veh_ima.time_lapse > 35 and veh_ima.control_type != 'neural':
                          veh_ima.control_type = 'neural'
  
                      if veh_ima.control_type == 'neural':
                          _act_long = veh_ima.act(obs)
                          # act_long = veh_ima.act(obs)
                          # _ = veh_ima.act(obs)
                          if self.metric_collection_mode:
                              veh_ima.act_long = act_long
                              self.mc_log_info(veh_real, veh_ima)
                      else:
                          act_long = veh_ima.idm_action(veh_ima, veh_ima.neighbours['att'])
                  else:
                      act_long = 0
  
              elif veh_ima.vehicle_type == 'idmmobil' and not veh_ima.collision_detected:
                  try:
                      act_long = veh_ima.idm_action(veh_ima, veh_ima.neighbours['att'])
                  except:
                      veh_ima.collision_detected = True
              else:
                  act_long = 0
  
  
              acts_ima.append([act_long, act_lat]) # lateral action is from veh_real
  
              if self.debugging_mode:
                  # veh_ima.act_long = _act_long
                  if veh_ima.vehicle_type == 'neural':
                      if veh_ima.control_type == 'neural':
                          veh_ima.act_long = _act_long
  
                      else:
                          veh_ima.act_long = act_long
                  else:
                      veh_ima.act_long = act_long
  
                  self.vis_log_info(veh_real, veh_ima)
  
          return acts_real, acts_ima
  
      def step(self, actions=None):
          """ steps the environment forward in time.
          """
          # self.remove_vehicles_outside_bound()
          acts_real, acts_ima = self.get_joint_action()
          for veh_real, veh_ima, act_real, act_ima in zip(
                                                      self.real_vehicles,
                                                      self.ima_vehicles,
                                                      acts_real,
                                                      acts_ima):
  
              veh_real.step(act_real)
              veh_ima.step(act_ima)
              veh_real.time_lapse += 1
              veh_ima.time_lapse += 1
  
          if self.time_step == 400:
          # if self.time_step > 600:
              ima_vehicles = []
              for vehicle in self.ima_vehicles:
                  if 100 < vehicle.glob_x < 400 and vehicle.vehicle_type != 'neural':
                  # if 50 < vehicle.time_lapse < 200 and vehicle.vehicle_type != 'neural'  \
                  #     and vehicle.id == 64:
                      neural_vehicle = self.idm_to_neural_vehicle(vehicle)
                      # neural_vehicle.id = 'neur_'+str(vehicle.id)
                      imagined_vehicle = neural_vehicle
                      imagined_vehicle.vehicle_type = 'neural'
                      imagined_vehicle.time_lapse = 0
                      ima_vehicles.append(imagined_vehicle)
                  else:
                      ima_vehicles.append(vehicle)
              self.ima_vehicles = ima_vehicles
          elif self.time_step <= 400:
              self.add_new_vehicles()
          self.time_step += 1
  
      def vis_log_info(self, veh_real, veh_ima):
          """
          For one off visualisation and debugging.
          """
          if veh_ima.vehicle_type == 'neural':
              veh_id =  veh_real.id
              if veh_real.neighbours['m'] and\
                                  veh_real.neighbours['att'] == veh_real.neighbours['m']:
                  att_real = 1
              else:
                  att_real = 0
  
              if veh_id not in self.real_mc_log:
                  self.real_mc_log[veh_id] = {}
                  self.ima_mc_log[veh_id] = {}
  
                  self.real_mc_log[veh_id]['act'] = []
                  self.real_mc_log[veh_id]['speed'] = []
                  self.real_mc_log[veh_id]['att'] = []
                  for key in veh_ima.driver_params.keys():
                      self.real_mc_log[veh_id][key] = []
  
                  self.ima_mc_log[veh_id]['act'] = []
                  self.ima_mc_log[veh_id]['speed'] = []
                  self.ima_mc_log[veh_id]['att'] = []
                  self.ima_mc_log[veh_id]['m_veh_exists'] = []
                  for key in veh_ima.driver_params.keys():
                      self.ima_mc_log[veh_id][key] = []
  
              self.real_mc_log[veh_id]['act'].append(veh_real.act_long)
              self.real_mc_log[veh_id]['speed'].append(veh_real.speed)
              self.real_mc_log[veh_id]['att'].append(att_real)
              for key in veh_ima.driver_params.keys():
                  self.real_mc_log[veh_id][key].append(veh_real.driver_params[key])
  
              # Imagined vehicle
              self.ima_mc_log[veh_id]['act'].append(veh_ima.act_long)
              self.ima_mc_log[veh_id]['speed'].append(veh_ima.speed)
  
              if not hasattr(veh_ima, 'att'):
                  self.ima_mc_log[veh_id]['att'].append(att_real)
              else:
                  self.ima_mc_log[veh_id]['att'].append(veh_ima.att)
  
  
              for key in veh_ima.driver_params.keys():
                  self.ima_mc_log[veh_id][key].append(veh_ima.driver_params[key])
  
              self.ima_mc_log[veh_id]['m_veh_exists'].append(\
                                                  veh_ima.m_veh_exists)
  
  
      def mc_log_info(self, veh_real, veh_ima):
          """
          Informatin to be logged:
          - ego (real and imagined) global_x for rwse and collision detection
          - ego (real and imagined) speed for rwse
          - ego (real and imagined) action for comparing action distributions
          """
          veh_id =  veh_real.id
          if veh_id not in self.real_mc_log:
              self.real_mc_log[veh_id] = {}
              self.ima_mc_log[veh_id] = {}
  
              self.real_mc_log[veh_id] = []
              self.ima_mc_log[veh_id] = []
  
          if veh_real.neighbours['m'] and\
                              veh_real.neighbours['att'] == veh_real.neighbours['m']:
              att_real = 1
          else:
              att_real = 0
  
          if veh_real.neighbours['att']:
              min_delta_x = veh_real.neighbours['att'].glob_x - veh_real.glob_x
          else:
              min_delta_x = 100
          real_mc_log = [self.time_step, veh_real.glob_x, \
                              veh_real.speed, veh_real.act_long, min_delta_x, att_real]
          self.real_mc_log[veh_id].append(real_mc_log)
  
          # Imagined vehicle
          if veh_ima.neighbours['att']:
              min_delta_x = veh_ima.neighbours['att'].glob_x - veh_ima.glob_x
          else:
              min_delta_x = 100
  
          # ima_mc_log = [self.time_step, veh_ima.glob_x, \
          #                         veh_ima.speed, veh_ima.act_long, min_delta_x, veh_ima.att]
          ima_mc_log = [self.time_step, veh_ima.glob_x, \
                                  veh_ima.speed, veh_ima.act_long, min_delta_x]
          self.ima_mc_log[veh_id].append(ima_mc_log)
  