import numpy as np
import pybullet as p
import itertools

from robot import Robot
    
class World():
    def __init__(self):
        # create the physics simulator
        self.physicsClient = p.connect(p.GUI)
        p.setGravity(0,0,-9.81)
        
        self.max_communication_distance = 2.0

        # We will integrate every 4ms (250Hz update)
        self.dt = 1./250.
        p.setPhysicsEngineParameter(self.dt, numSubSteps=1)

        # Create the plane.
        self.planeId = p.loadURDF("../models/plane.urdf")
        p.changeDynamics(self.planeId, -1, lateralFriction=5., rollingFriction=0)

        self.goalId = p.loadURDF("../models/goal.urdf")
        self.goalId = p.loadURDF("../models/goal2.urdf")
        
        # the balls
        self.ball1 = p.loadURDF("../models/ball1.urdf")
        p.resetBasePositionAndOrientation(self.ball1, [2., 4., 0.5], (0., 0., 0.5, 0.5))
        self.ball2 = p.loadURDF("../models/ball2.urdf")
        p.resetBasePositionAndOrientation(self.ball2, [4., 2., 0.5], (0., 0., 0.5, 0.5))

        p.resetDebugVisualizerCamera(7.0,90.0, -43.0, (1., 1., 0.0))
        
        # Add objects
        wallId = p.loadSDF("../models/walls.sdf")[0]
        p.resetBasePositionAndOrientation(wallId, [0., -1., 0], (0., 0., 0.5, 0.5))
        wallId = p.loadSDF("../models/walls.sdf")[0]
        p.resetBasePositionAndOrientation(wallId, [0., 1., 0], (0., 0., 0.5, 0.5))
        wallId = p.loadSDF("../models/walls.sdf")[0]
        p.resetBasePositionAndOrientation(wallId, [3., -1., 0], (0., 0., 0.5, 0.5))
        wallId = p.loadSDF("../models/walls.sdf")[0]
        p.resetBasePositionAndOrientation(wallId, [3., 1., 0], (0., 0., 0.5, 0.5))
        wallId = p.loadSDF("../models/walls.sdf")[0]
        p.resetBasePositionAndOrientation(wallId, [1., 2., 0], (0., 0., 0., 1.))
        wallId = p.loadSDF("../models/walls.sdf")[0]
        p.resetBasePositionAndOrientation(wallId, [2., -2., 0], (0., 0., 0., 1.))

        # tube
        # wallId = p.loadSDF("../models/walls.sdf")[0]
        # p.resetBasePositionAndOrientation(wallId, [-1., 5., 0], (0., 0., 0., 1.))
        # wallId = p.loadSDF("../models/walls.sdf")[0]
        # p.resetBasePositionAndOrientation(wallId, [-1., 6., 0], (0., 0., 0., 1.))
        
        # #arena
        # wallId = p.loadSDF("../models/walls.sdf")[0]
        # p.resetBasePositionAndOrientation(wallId, [-2, 4., 0], (0., 0., 0.5, 0.5))
        # wallId = p.loadSDF("../models/walls.sdf")[0]
        # p.resetBasePositionAndOrientation(wallId, [-2., 7., 0], (0., 0., 0.5, 0.5))
        # wallId = p.loadSDF("../models/walls.sdf")[0]
        # p.resetBasePositionAndOrientation(wallId, [-2., 9., 0], (0., 0., 0.5, 0.5))
        # wallId = p.loadSDF("../models/walls.sdf")[0]
        # p.resetBasePositionAndOrientation(wallId, [-2., 11., 0], (0., 0., 0.5, 0.5))
        # wallId = p.loadSDF("../models/walls.sdf")[0]
        # p.resetBasePositionAndOrientation(wallId, [-2., 13., 0], (0., 0., 0.5, 0.5))
        # wallId = p.loadSDF("../models/walls.sdf")[0]
        # p.resetBasePositionAndOrientation(wallId, [-3., 3., 0], (0., 0., 0., 1.))
        # wallId = p.loadSDF("../models/walls.sdf")[0]
        # p.resetBasePositionAndOrientation(wallId, [-5., 3., 0], (0., 0., 0., 1.))
        # wallId = p.loadSDF("../models/walls.sdf")[0]
        # p.resetBasePositionAndOrientation(wallId, [-7., 3., 0], (0., 0., 0., 1.))
        # wallId = p.loadSDF("../models/walls.sdf")[0]
        # p.resetBasePositionAndOrientation(wallId, [-8, 4., 0], (0., 0., 0.5, 0.5))
        # wallId = p.loadSDF("../models/walls.sdf")[0]
        # p.resetBasePositionAndOrientation(wallId, [-8., 6., 0], (0., 0., 0.5, 0.5))
        # wallId = p.loadSDF("../models/walls.sdf")[0]
        # p.resetBasePositionAndOrientation(wallId, [-8., 8., 0], (0., 0., 0.5, 0.5))
        # wallId = p.loadSDF("../models/walls.sdf")[0]
        # p.resetBasePositionAndOrientation(wallId, [-8., 10., 0], (0., 0., 0.5, 0.5))
        # wallId = p.loadSDF("../models/walls.sdf")[0]
        # p.resetBasePositionAndOrientation(wallId, [-8., 12., 0], (0., 0., 0.5, 0.5))

        
        # create 6 robots
        self.robots = []
        for (i,j) in itertools.product(range(3), range(2)):
            self.robots.append(Robot([1. * i + 0.5, 1. * j - 0.5, 0.3], 2*i+j, self.dt))
            p.stepSimulation()
        
        self.time = 0.0
        
        self.stepSimulation()
        self.stepSimulation()

    def reset(self):
        """
        Resets the position of all the robots
        """
        for r in self.robots:
            r.reset()
        p.stepSimulation()
        
    def stepSimulation(self):
        """
        Simulates one step simulation
        """
        
        # for each robot construct list of neighbors
        for r in self.robots:
            r.neighbors = [] #reset neighbors
            r.messages_received = [] #reset message received
            pos1, or1 = r.get_pos_and_orientation()
            for j,r2 in enumerate(self.robots):
                if(r.id != r2.id):
                    pos2, or2 = r2.get_pos_and_orientation()
                    if(np.linalg.norm(pos1-pos2) < self.max_communication_distance):
                        r.neighbors.append(j)
        
        # for each robot send and receive messages
        for i,r in enumerate(self.robots):
            for msg in r.messages_to_send:
                if msg[0] in r.neighbors: #then we can send the message
                    self.robots[msg[0]].messages_received.append([i, msg[1]]) #add the sender id
            r.messages_to_send = []
        
        # update the controllers
        if self.time > 1.0:
            for r in self.robots:
                r.compute_controller()
        
        # do one simulation step
        p.stepSimulation()
        self.time += self.dt
        

[((194, 210), 'pybullet.connect', 'p.connect', (['p.GUI'], {}), '(p.GUI)\n', (203, 210), True, 'import pybullet as p\n'), ((219, 244), 'pybullet.setGravity', 'p.setGravity', (['(0)', '(0)', '(-9.81)'], {}), '(0, 0, -9.81)\n', (231, 244), True, 'import pybullet as p\n'), ((386, 437), 'pybullet.setPhysicsEngineParameter', 'p.setPhysicsEngineParameter', (['self.dt'], {'numSubSteps': '(1)'}), '(self.dt, numSubSteps=1)\n', (413, 437), True, 'import pybullet as p\n'), ((490, 524), 'pybullet.loadURDF', 'p.loadURDF', (['"""../models/plane.urdf"""'], {}), "('../models/plane.urdf')\n", (500, 524), True, 'import pybullet as p\n'), ((533, 607), 'pybullet.changeDynamics', 'p.changeDynamics', (['self.planeId', '(-1)'], {'lateralFriction': '(5.0)', 'rollingFriction': '(0)'}), '(self.planeId, -1, lateralFriction=5.0, rollingFriction=0)\n', (549, 607), True, 'import pybullet as p\n'), ((630, 663), 'pybullet.loadURDF', 'p.loadURDF', (['"""../models/goal.urdf"""'], {}), "('../models/goal.urdf')\n", (640, 663), True, 'import pybullet as p\n'), ((686, 720), 'pybullet.loadURDF', 'p.loadURDF', (['"""../models/goal2.urdf"""'], {}), "('../models/goal2.urdf')\n", (696, 720), True, 'import pybullet as p\n'), ((771, 805), 'pybullet.loadURDF', 'p.loadURDF', (['"""../models/ball1.urdf"""'], {}), "('../models/ball1.urdf')\n", (781, 805), True, 'import pybullet as p\n'), ((814, 903), 'pybullet.resetBasePositionAndOrientation', 'p.resetBasePositionAndOrientation', (['self.ball1', '[2.0, 4.0, 0.5]', '(0.0, 0.0, 0.5, 0.5)'], {}), '(self.ball1, [2.0, 4.0, 0.5], (0.0, 0.0, \n    0.5, 0.5))\n', (847, 903), True, 'import pybullet as p\n'), ((916, 950), 'pybullet.loadURDF', 'p.loadURDF', (['"""../models/ball2.urdf"""'], {}), "('../models/ball2.urdf')\n", (926, 950), True, 'import pybullet as p\n'), ((959, 1048), 'pybullet.resetBasePositionAndOrientation', 'p.resetBasePositionAndOrientation', (['self.ball2', '[4.0, 2.0, 0.5]', '(0.0, 0.0, 0.5, 0.5)'], {}), '(self.ball2, [4.0, 2.0, 0.5], (0.0, 0.0, \n    0.5, 0.5))\n', (992, 1048), True, 'import pybullet as p\n'), ((1049, 1112), 'pybullet.resetDebugVisualizerCamera', 'p.resetDebugVisualizerCamera', (['(7.0)', '(90.0)', '(-43.0)', '(1.0, 1.0, 0.0)'], {}), '(7.0, 90.0, -43.0, (1.0, 1.0, 0.0))\n', (1077, 1112), True, 'import pybullet as p\n'), ((1202, 1281), 'pybullet.resetBasePositionAndOrientation', 'p.resetBasePositionAndOrientation', (['wallId', '[0.0, -1.0, 0]', '(0.0, 0.0, 0.5, 0.5)'], {}), '(wallId, [0.0, -1.0, 0], (0.0, 0.0, 0.5, 0.5))\n', (1235, 1281), True, 'import pybullet as p\n'), ((1339, 1417), 'pybullet.resetBasePositionAndOrientation', 'p.resetBasePositionAndOrientation', (['wallId', '[0.0, 1.0, 0]', '(0.0, 0.0, 0.5, 0.5)'], {}), '(wallId, [0.0, 1.0, 0], (0.0, 0.0, 0.5, 0.5))\n', (1372, 1417), True, 'import pybullet as p\n'), ((1475, 1554), 'pybullet.resetBasePositionAndOrientation', 'p.resetBasePositionAndOrientation', (['wallId', '[3.0, -1.0, 0]', '(0.0, 0.0, 0.5, 0.5)'], {}), '(wallId, [3.0, -1.0, 0], (0.0, 0.0, 0.5, 0.5))\n', (1508, 1554), True, 'import pybullet as p\n'), ((1612, 1690), 'pybullet.resetBasePositionAndOrientation', 'p.resetBasePositionAndOrientation', (['wallId', '[3.0, 1.0, 0]', '(0.0, 0.0, 0.5, 0.5)'], {}), '(wallId, [3.0, 1.0, 0], (0.0, 0.0, 0.5, 0.5))\n', (1645, 1690), True, 'import pybullet as p\n'), ((1748, 1826), 'pybullet.resetBasePositionAndOrientation', 'p.resetBasePositionAndOrientation', (['wallId', '[1.0, 2.0, 0]', '(0.0, 0.0, 0.0, 1.0)'], {}), '(wallId, [1.0, 2.0, 0], (0.0, 0.0, 0.0, 1.0))\n', (1781, 1826), True, 'import pybullet as p\n'), ((1882, 1961), 'pybullet.resetBasePositionAndOrientation', 'p.resetBasePositionAndOrientation', (['wallId', '[2.0, -2.0, 0]', '(0.0, 0.0, 0.0, 1.0)'], {}), '(wallId, [2.0, -2.0, 0], (0.0, 0.0, 0.0, 1.0))\n', (1915, 1961), True, 'import pybullet as p\n'), ((4600, 4618), 'pybullet.stepSimulation', 'p.stepSimulation', ([], {}), '()\n', (4616, 4618), True, 'import pybullet as p\n'), ((5790, 5808), 'pybullet.stepSimulation', 'p.stepSimulation', ([], {}), '()\n', (5806, 5808), True, 'import pybullet as p\n'), ((1158, 1190), 'pybullet.loadSDF', 'p.loadSDF', (['"""../models/walls.sdf"""'], {}), "('../models/walls.sdf')\n", (1167, 1190), True, 'import pybullet as p\n'), ((1295, 1327), 'pybullet.loadSDF', 'p.loadSDF', (['"""../models/walls.sdf"""'], {}), "('../models/walls.sdf')\n", (1304, 1327), True, 'import pybullet as p\n'), ((1431, 1463), 'pybullet.loadSDF', 'p.loadSDF', (['"""../models/walls.sdf"""'], {}), "('../models/walls.sdf')\n", (1440, 1463), True, 'import pybullet as p\n'), ((1568, 1600), 'pybullet.loadSDF', 'p.loadSDF', (['"""../models/walls.sdf"""'], {}), "('../models/walls.sdf')\n", (1577, 1600), True, 'import pybullet as p\n'), ((1704, 1736), 'pybullet.loadSDF', 'p.loadSDF', (['"""../models/walls.sdf"""'], {}), "('../models/walls.sdf')\n", (1713, 1736), True, 'import pybullet as p\n'), ((1838, 1870), 'pybullet.loadSDF', 'p.loadSDF', (['"""../models/walls.sdf"""'], {}), "('../models/walls.sdf')\n", (1847, 1870), True, 'import pybullet as p\n'), ((4327, 4345), 'pybullet.stepSimulation', 'p.stepSimulation', ([], {}), '()\n', (4343, 4345), True, 'import pybullet as p\n'), ((4257, 4319), 'robot.Robot', 'Robot', (['[1.0 * i + 0.5, 1.0 * j - 0.5, 0.3]', '(2 * i + j)', 'self.dt'], {}), '([1.0 * i + 0.5, 1.0 * j - 0.5, 0.3], 2 * i + j, self.dt)\n', (4262, 4319), False, 'from robot import Robot\n'), ((5138, 5165), 'numpy.linalg.norm', 'np.linalg.norm', (['(pos1 - pos2)'], {}), '(pos1 - pos2)\n', (5152, 5165), True, 'import numpy as np\n')]