text: import logging
  import math
  import time
  import copy
  
  import pygame
  from pygame.locals import (QUIT, KEYDOWN, K_ESCAPE)
  import numpy as np
  import pickle
  import imageio
  
  import Box2D  
  from Box2D.b2 import (world, polygonShape, circleShape, staticBody, dynamicBody, kinematicBody)
  
  import matplotlib.pyplot as plt
  import os
  from helper import *
  import json
  import random
  from policies import *
  from prune import *
  from keras.models import load_model
  from mpc import *
  
  PPM = 60.0  # pixels per meter
  TIME_STEP = 0.1
  SCREEN_WIDTH, SCREEN_HEIGHT = 720, 720
  GROUPS = [(1, 1, 15), (0.75, 1, 12), (0.5, 1, 9), (0.25, 1, 6)]
  num_classes = 8
  base = plt.cm.get_cmap('Accent')
  color_list = base(np.linspace(0, 1, num_classes))
  COLORS = [(int(255*c[0]), int(255*c[1]), int(255*c[2]), int(255*c[3])) for c in color_list]
  WHITE = (255, 255, 255, 255)
  FRICTION = 0.5
  DAMPING_FACTOR = 1 - ((1 - 0.5) / 3)
  
  
  class Action:
      def __init__(self, vector, point):
          """
          action that consists of:
          vector: (x, y) force vector
          point: (x, y) point of contact
          all relative to the local origin of polygon.
          """
          self.vector = normalize(vector)
          self.point = point
  
      def __eq__(self, other):
          """
          check if vector == vector, point == point
          """
          return self.vector == other.vector and self.point == self.point
  
  
  class Polygon:
      def __init__(self, body, fixtures, vertices, color=WHITE):
          """body: polygon shape (dynamicBody)
          fixture: fixture
          vertices: list of relative coordinates
          """
          self.body = body
          self.fixtures = fixtures
          self.vertices = vertices
          self.color = color
          self.original_pos = np.array(self.body.position)
          self.bounding_circle_radius = math.sqrt(max(self.vertices[:, 0]**2 + self.vertices[:, 1]**2))
  
      def test_overlap(self, other_polygon):
          if self.dist(other_polygon) > 0:
              return False
          return True
  
      def dist(self, other_polygon):
          """do not work on a polygon and a rod"""
          shape1 = self.fixtures[0].shape
          shape2 = other_polygon.fixtures[0].shape
          transform1 = Box2D.b2Transform()
          pos1 = self.body.position
          angle1 = self.body.angle
          transform1.Set(pos1, angle1)
          transform2 = Box2D.b2Transform()
          pos2 = other_polygon.body.position
          angle2 = other_polygon.body.angle
          transform2.Set(pos2, angle2)
          point_a, point_b, distance, iterations = Box2D.b2Distance(
              shapeA=shape1, shapeB=shape2, transformA=transform1, transformB=transform2)
          return distance
  
      def dist_rod(self, rod_fix, rod_body):
          shape1 = self.fixtures[0].shape
          shape2 = rod_fix.shape
          transform1 = Box2D.b2Transform()
          pos1 = self.body.position
          angle1 = self.body.angle
          transform1.Set(pos1, angle1)
          transform2 = Box2D.b2Transform()
          pos2 = rod_body.position
          angle2 = rod_body.angle
          transform2.Set(pos2, angle2)
          point_a, point_b, distance, iterations = Box2D.b2Distance(
              shapeA=shape1, shapeB=shape2, transformA=transform1, transformB=transform2)
          return distance
  
  
  class SingulationEnv:
      def __init__(self):
          self.world = world(gravity=(0, 0), doSleep=True)
          self.objs = []
          self.rod = None
          self.bounding_convex_hull = np.array([])
          self.centroid = (0, 0)
          self.bounding_circle_radius = 0
          self.screen = None
          polygonShape.draw = self.my_draw_polygon
          circleShape.draw = self.my_draw_circle
  
      def my_draw_polygon(self, polygon, body, color):
          vertices = [(body.transform * v) * PPM for v in polygon.vertices]
          vertices = [(v[0], SCREEN_HEIGHT - v[1]) for v in vertices]
          pygame.draw.polygon(self.screen, color, vertices)  # inside rectangle
          pygame.draw.polygon(self.screen, (0, 0, 0, 0), vertices, 5)  # boundary of rectangle
  
      def my_draw_circle(self, circle, body, color):
          position = body.transform * circle.pos * PPM
          position = (position[0], SCREEN_HEIGHT - position[1])
          pygame.draw.circle(self.screen, color, [int(
              x) for x in position], int(circle.radius * PPM))
  
      def create_random_env(self, num_objs=3, group=0):
          assert num_objs >= 1
          # clear environment
          if len(self.objs) > 0:
              for obj in self.objs:
                  for fix in obj.fixtures:
                      obj.body.DestroyFixture(fix)
                  self.world.DestroyBody(obj.body)
          self.objs = []
          for i in range(num_objs):
              # create shape
              vertices = generatePolygon(GROUPS[group][0], GROUPS[group][1], GROUPS[group][2])
              raw_com = np.array(compute_centroid(vertices))
              vertices = (vertices - raw_com)
              bounding_r = math.sqrt(max(vertices[:, 0]**2 + vertices[:, 1]**2))
              vertices = vertices / bounding_r
              # place object
              if len(self.objs) == 0:
                  original_pos = np.array([np.random.uniform(4, 8), np.random.uniform(4, 8)])
                  body = self.world.CreateDynamicBody(position=original_pos.tolist(), allowSleep=False)
                  fixture = body.CreatePolygonFixture(density=1, vertices=vertices.tolist(), friction=FRICTION)
                  self.objs.append(Polygon(body, [fixture], vertices, COLORS[i]))
              else:
                  max_iter = 10000
                  while True:
                      if max_iter <= 0:
                          raise RuntimeError("max iter reaches")
                      max_iter -= 1
                      # place object close to the last object
                      original_pos = np.array([np.random.uniform(-1.8, 1.8), np.random.uniform(-1.8, 1.8)]) + np.array(self.objs[-1].body.position)
                      original_pos = np.clip(original_pos, 1, 11)
                      body = self.world.CreateDynamicBody(position=original_pos.tolist(), allowSleep=False)
                      fixture = body.CreatePolygonFixture(density=1, vertices=vertices.tolist(), friction=FRICTION)
                      curr_polygon = Polygon(body, [fixture], vertices, COLORS[i % len(COLORS)])
                      overlap = False
                      for obj in self.objs:
                          if obj.test_overlap(curr_polygon):
                              overlap = True
                              break
                      if overlap:
                          body.DestroyFixture(fixture)
                          self.world.DestroyBody(body)
                      else:
                          self.objs.append(curr_polygon)
                          break
          self.bounding_convex_hull = create_convex_hull(np.concatenate([obj.vertices+obj.original_pos for obj in self.objs]))
          self.centroid = compute_centroid(self.bounding_convex_hull.tolist())
          self.bounding_circle_radius = math.sqrt(max(
              (self.bounding_convex_hull - np.array(self.centroid))[:, 0]**2
              + (self.bounding_convex_hull - np.array(self.centroid))[:, 1]**2))
  
      def create_random_env_wrapper(self, num_objs=3, group=0):
          while True:
              try:
                  self.create_random_env(num_objs,group)
                  return
              except RuntimeError:
                  pass
  
      def visualize(self, path):
          """Capture an image of the current state.
          path: .../something.png
          """
          self.screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
          pygame.display.iconify()
          self.screen.fill(WHITE)
          for obj in self.objs:
              for fix in obj.fixtures:
                  fix.shape.draw(fix.shape, obj.body, obj.color)
          pygame.image.save(self.screen, path)
          pygame.display.quit()
          pygame.quit()
  
      def step(self, start_pt, end_pt, path, display=False, check_reachable=True):
          if display:
              images = []
              if not os.path.exists(path):
                  os.makedirs(path)
              self.screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
              pygame.display.iconify()
              self.screen.fill(WHITE)
          start_pt = np.array(start_pt)
          end_pt = np.array(end_pt)
          self.rod = self.world.CreateKinematicBody(position=(start_pt[0], start_pt[1]), allowSleep=False)
          self.rodfix = self.rod.CreatePolygonFixture(vertices=[(0.1, 0.1), (-0.1, 0.1), (-0.1, -0.1), (0.1, -0.1)])
          vector = np.array(normalize(end_pt - np.array(start_pt)))
          # reachability check
          if check_reachable:
              vertices_lst = [(0.0, 0.1), (0.0, -0.1), (-0.3, -0.1), (-0.3, 0.1)]
              testrod = self.world.CreateKinematicBody(position=(start_pt[0], start_pt[1]), allowSleep=False)
              testrodfix = testrod.CreatePolygonFixture(vertices=[rotatePt(pt, vector) for pt in vertices_lst])
              while (np.count_nonzero(np.array([o.dist_rod(testrodfix, testrod) for o in self.objs]) <= 0) > 0):
                  start_pt -= 0.1 * vector
                  testrod.DestroyFixture(testrodfix)
                  self.world.DestroyBody(testrod)
                  testrod = self.world.CreateKinematicBody(position=(start_pt[0], start_pt[1]), allowSleep=False)
                  testrodfix = testrod.CreatePolygonFixture(vertices=[rotatePt(pt, vector) for pt in vertices_lst])
              testrod.DestroyFixture(testrodfix)
              self.world.DestroyBody(testrod)
          self.rod.linearVelocity[0] = vector[0]
          self.rod.linearVelocity[1] = vector[1]
          self.rod.angularVelocity = 0.0
          timestamp = 0
          # display
          if display:
              for obj in self.objs:
                  for fix in obj.fixtures:
                      fix.shape.draw(fix.shape, obj.body, obj.color)
              self.rodfix.shape.draw(self.rodfix.shape, self.rod, (0, 0, 0, 255))
              img_path = os.path.join(path, '0.png')
              pygame.image.save(self.screen, img_path)
              images.append(imageio.imread(img_path))
              os.remove(img_path)
          first_contact = -1
          while (timestamp < 100):
              if first_contact == -1:
                  for i in range(len(self.objs)):
                      # if object is moving, classify it as contacted
                      if (self.objs[i].body.linearVelocity[0] ** 2 + self.objs[i].body.linearVelocity[1] ** 2 > 0.001):
                          first_contact = i
              for obj in self.objs:
                  obj.body.linearVelocity[0] *= DAMPING_FACTOR
                  obj.body.linearVelocity[1] *= DAMPING_FACTOR
                  obj.body.angularVelocity *= DAMPING_FACTOR
              if (math.sqrt(np.sum((start_pt - np.array(self.rod.position)) ** 2)) < 4):
                  vector = normalize((end_pt + 1e-8) - (start_pt + 1e-8))
                  self.rod.linearVelocity[0] = vector[0]
                  self.rod.linearVelocity[1] = vector[1]
              else:
                  self.rod.linearVelocity[0] = 0
                  self.rod.linearVelocity[1] = 0
              self.world.Step(TIME_STEP, 10, 10)
              timestamp += 1
              # display
              if display:
                  self.screen.fill((255, 255, 255, 255))
                  for obj in self.objs:
                      for fix in obj.fixtures:
                          fix.shape.draw(fix.shape, obj.body, obj.color)
                  self.rodfix.shape.draw(self.rodfix.shape, self.rod, (0, 0, 0, 255))
                  img_path = os.path.join(path, str(timestamp) + '.png')
                  pygame.image.save(self.screen, img_path)
                  images.append(imageio.imread(img_path))
                  os.remove(img_path)
          # display
          if display:
              pygame.display.quit()
              pygame.quit()
              imageio.mimsave(os.path.join(path, 'push.gif'), images)
          return first_contact
  
      def count_threshold(self, threshold=0.3):
          count = 0
          for i in range(len(self.objs)):
              isolated = True
              for j in range(len(self.objs)):
                  if i != j:
                      shape1 = self.objs[i].fixtures[0].shape
                      shape2 = self.objs[j].fixtures[0].shape
                      transform1 = Box2D.b2Transform()
                      pos1 = self.objs[i].body.position
                      angle1 = self.objs[i].body.angle
                      transform1.Set(pos1, angle1)
                      transform2 = Box2D.b2Transform()
                      pos2 = self.objs[j].body.position
                      angle2 = self.objs[j].body.angle
                      transform2.Set(pos2, angle2)
                      pointA, pointB, distance, iterations = Box2D.b2Distance(shapeA=shape1, shapeB=shape2, transformA=transform1, transformB=transform2)
                      if distance < threshold:
                          isolated = False
              if isolated:
                  count += 1
          return count
  
      def count_soft_threshold(self):
          count = 0.0
          for i in range(len(self.objs)):
              min_dist = 1e2
              for j in range(i+1, len(self.objs)):
                  shape1 = self.objs[i].fixtures[0].shape
                  shape2 = self.objs[j].fixtures[0].shape
                  transform1 = Box2D.b2Transform()
                  pos1 = self.objs[i].body.position
                  angle1 = self.objs[i].body.angle
                  transform1.Set(pos1, angle1)
                  transform2 = Box2D.b2Transform()
                  pos2 = self.objs[j].body.position
                  angle2 = self.objs[j].body.angle
                  transform2.Set(pos2, angle2)
                  pointA, pointB, distance, iterations = Box2D.b2Distance(shapeA=shape1, shapeB=shape2, transformA=transform1, transformB=transform2)
                  if distance < min_dist:
                      min_dist = distance
              # assume threshold = 0.3. Want f(0) = 0, f(0.3)=1
              count += (sigmoid(min_dist*10) - 0.5) * 2
          return count
  
      def collect_data_summary(self, start_pt, end_pt, golden_summary=None, first_step=False, img_path=None, display=False, sum_path=None):
          summary = {}
          abs_start_pt = np.array(start_pt)
          abs_end_pt = np.array(end_pt)
          summary["start pt"] = abs_start_pt.tolist()
          summary["end pt"] = abs_end_pt.tolist()
          for i in range(len(self.objs)):
              summary[str(i) + " dist to pushing line"] = pointToLineDistance(abs_start_pt, abs_end_pt, self.objs[i].body.position)
              summary[str(i) + " original pos"] = np.array(self.objs[i].body.position).tolist()
              summary[str(i) + " project dist"] = projectedPtToStartDistance(abs_start_pt, abs_end_pt, self.objs[i].body.position)
              summary[str(i) + " vertices"] = np.array(self.objs[i].vertices).tolist()
          summary["count soft threshold before push"] = self.count_soft_threshold()
          first_contact = self.step(start_pt, end_pt, path=img_path, display=display)
          for i in range(len(self.objs)):
              summary[str(i) + " change of pos"] = euclidean_dist(self.objs[i].body.position, self.objs[i].original_pos)
          summary["count soft threshold after push"] = self.count_soft_threshold()
          summary["first contact object"] = first_contact
          if sum_path is not None:
              with open(os.path.join(sum_path, 'summary.json'), 'w') as f:
                  json.dump(summary, f)
          return summary
  
      def reset(self):
          if self.rod:
              self.rod.DestroyFixture(self.rodfix)
              self.world.DestroyBody(self.rod)
              self.rod = None
          for obj in self.objs:
              obj.body.position[0] = obj.original_pos[0]
              obj.body.position[1] = obj.original_pos[1]
              obj.body.angle = 0.0
              obj.body.linearVelocity[0] = 0.0
              obj.body.linearVelocity[1] = 0.0
              obj.body.angularVelocity = 0.0
  
      def save_env(self, sum_path=None):
          """Save information about current state in a dictionary in sum_path/env.json"""
          summary = {}
          summary["num_objects"] = len(self.objs)
          summary["count soft threshold"] = self.count_soft_threshold()
          for i in range(len(self.objs)):
              summary[str(i) + " original pos"] = np.array(self.objs[i].body.position).tolist()
              summary[str(i) + " angle"] = self.objs[i].body.angle
              summary[str(i) + " vertices"] = np.array(self.objs[i].vertices).tolist()
          if sum_path != None:
              with open(os.path.join(sum_path, 'env.json'), 'w') as f:
                  json.dump(summary, f)
          return summary
  
      def load_env(self, dic):
          for i in range(dic["num_objects"]):
              original_pos = np.array(dic[str(i)+" original pos"])
              vertices = np.array(dic[str(i)+" vertices"])
              body = self.world.CreateDynamicBody(position=original_pos.tolist(), angle=dic.get(str(i)+" angle", 0.0), allowSleep=False)
              fixture = body.CreatePolygonFixture(density=1, vertices=vertices.tolist(), friction=FRICTION)
              self.objs.append(Polygon(body, [fixture], vertices, COLORS[i % len(COLORS)]))
  
      def save_curr_position(self):
          position = {}
          for i in range(len(self.objs)):
              position[i] = [self.objs[i].body.position[0], self.objs[i].body.position[1], self.objs[i].body.angle]
          return position
  
      def load_position(self, position):
          if self.rod:
              self.rod.DestroyFixture(self.rodfix)
              self.world.DestroyBody(self.rod)
              self.rod = None
          for i in range(len(self.objs)):
              self.objs[i].body.position[0] = position[i][0]
              self.objs[i].body.position[1] = position[i][1]
              self.objs[i].body.angle = position[i][2]
  
      def best_sequential_sample(self, num_samples, prune_method, reuse, max_step, data_path, reuse_path=None, metric="count soft threshold", open_loop=False):
          best_result = 0
          best_action = None
          best_sample = None
          best_info = None
  # 		import pdb; pdb.set_trace()
          summary = self.save_env()
          first_step_pt_lst = prune_method(self)
          complete_pt_lst = list(first_step_pt_lst)
          for i in range(num_samples):
              sample_path = os.path.join(data_path, "sample"+str(i))
              if not os.path.exists(sample_path):
                  os.makedirs(sample_path)
              before = self.save_env()
              result, action, info = collect_sequential_sample(complete_pt_lst, summary, self, reuse, max_step, sample_path, metric, open_loop)
              if before != self.save_env():
                  print("before != self.save_env()")
              if result > best_result:
                  best_result = result
                  best_action = action
                  best_sample = i
                  best_info = info
          if reuse:
              with open(os.path.join(reuse_path, "best_sample.json"), 'w') as f:
                  json.dump({"best_sample": best_sample}, f)
          else:
              with open(os.path.join(data_path,"best_sample.json"), 'w') as f:
                  json.dump({"best_sample": best_sample}, f)
          return best_result, best_action, best_info
  
      def prune_best(self, prune_method, metric="count threshold", position=None):
          pt_lst = prune_method(self)
          best_pt = None
          best_sep = -1e2
          print("len(pt_lst)", len(pt_lst))
          info = {}
          self.visualize("./A.png")
          for pts in pt_lst[:1000]:
              if position is None:
                  self.reset()
              else:
                  self.load_position(position)
              summary = self.collect_data_summary(pts[0], pts[1], None)
              if summary[metric +" after push"] - summary[metric + " before push"] >= best_sep:
                  best_pt = pts
                  best_sep = summary[metric +" after push"] - summary[metric + " before push"]
                  info[metric + " before push"] = summary[metric + " before push"]
                  info[metric + " after push"] = summary[metric + " after push"]
                  info["env after push"] = self.save_env()
                  self.visualize("./B.png")
          if position is None:
              self.reset()
          else:
              self.load_position(position)
          return best_pt, info
  
      def select_random(self, prune_method):
          pt_lst = prune_method(self)
          return random.choice(pt_lst)
  
      # Forward dynamics model
      def generate_state_action_pairs(self,path=None,display=False):
          state_action = []
          for o in self.objs:
              state_action.append(o.body.position[0])
              state_action.append(o.body.position[1])
              state_action.append(o.body.angle)
              for v in o.vertices:
                  state_action.append(v[0])
                  state_action.append(v[1])
          best_pts = self.select_random(com_only)
          if best_pts is None:
              return [], []
          best_summary = self.collect_data_summary(best_pts[0], best_pts[1], img_path=path, display=display)
          vec = normalize(np.array(best_pts[1]) - np.array(best_pts[0]))
          state_action.append(best_pts[0][0])
          state_action.append(best_pts[0][1])
          state_action.append(vec[0])
          state_action.append(vec[1])
          next_state = []
          for o in self.objs:
              next_state.append(o.body.position[0])
              next_state.append(o.body.position[1])
              next_state.append(o.body.angle)
              for v in o.vertices:
                  next_state.append(v[0])
                  next_state.append(v[1])
          if best_summary["first contact object"] == -1:
              return [], []
          return state_action, next_state
  
      def reset_env_five_objects(self, v):
          if self.rod:
              self.rod.DestroyFixture(self.rodfix)
              self.world.DestroyBody(self.rod)
              self.rod = None
          for i in range(len(self.objs)):
              obj = self.objs[i]
              index = int(33*i)
              obj.body.position[0] = float(v[index])
              obj.body.position[1] = float(v[index+1])
              obj.body.angle = float(v[index+2])
              obj.body.linearVelocity[0] = 0.0
              obj.body.linearVelocity[1] = 0.0
              obj.body.angularVelocity = 0.0
  
  
  if __name__ == "__main__":
      env = SingulationEnv()
      env.create_random_env_wrapper(num_objs=5)
      print("before")
      print("count_soft_threshold:", env.count_soft_threshold())
      for i in range(len(env.objs)):
          print("position%d:" % i, env.objs[i].body.position)
      push = random.choice(no_prune(env))
      env.step(push[0], push[1], "./no_reach/push", display=True, check_reachable=False)
      print("after")
      print("count_soft_threshold:", env.count_soft_threshold())
      for i in range(len(env.objs)):
          print("position:", env.objs[i].body.position)
  