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
        

[((10, 29, 10, 45), 'pybullet.connect', 'p.connect', ({(10, 39, 10, 44): 'p.GUI'}, {}), '(p.GUI)', True, 'import pybullet as p\n'), ((11, 8, 11, 31), 'pybullet.setGravity', 'p.setGravity', ({(11, 21, 11, 22): '(0)', (11, 23, 11, 24): '(0)', (11, 25, 11, 30): '(-9.81)'}, {}), '(0, 0, -9.81)', True, 'import pybullet as p\n'), ((17, 8, 17, 59), 'pybullet.setPhysicsEngineParameter', 'p.setPhysicsEngineParameter', (), '', True, 'import pybullet as p\n'), ((20, 23, 20, 57), 'pybullet.loadURDF', 'p.loadURDF', ({(20, 34, 20, 56): '"""../models/plane.urdf"""'}, {}), "('../models/plane.urdf')", True, 'import pybullet as p\n'), ((21, 8, 21, 81), 'pybullet.changeDynamics', 'p.changeDynamics', (), '', True, 'import pybullet as p\n'), ((23, 22, 23, 55), 'pybullet.loadURDF', 'p.loadURDF', ({(23, 33, 23, 54): '"""../models/goal.urdf"""'}, {}), "('../models/goal.urdf')", True, 'import pybullet as p\n'), ((24, 22, 24, 56), 'pybullet.loadURDF', 'p.loadURDF', ({(24, 33, 24, 55): '"""../models/goal2.urdf"""'}, {}), "('../models/goal2.urdf')", True, 'import pybullet as p\n'), ((27, 21, 27, 55), 'pybullet.loadURDF', 'p.loadURDF', ({(27, 32, 27, 54): '"""../models/ball1.urdf"""'}, {}), "('../models/ball1.urdf')", True, 'import pybullet as p\n'), ((28, 8, 28, 88), 'pybullet.resetBasePositionAndOrientation', 'p.resetBasePositionAndOrientation', ({(28, 42, 28, 52): 'self.ball1', (28, 54, 28, 67): '[2.0, 4.0, 0.5]', (28, 69, 28, 87): '(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))', True, 'import pybullet as p\n'), ((29, 21, 29, 55), 'pybullet.loadURDF', 'p.loadURDF', ({(29, 32, 29, 54): '"""../models/ball2.urdf"""'}, {}), "('../models/ball2.urdf')", True, 'import pybullet as p\n'), ((30, 8, 30, 88), 'pybullet.resetBasePositionAndOrientation', 'p.resetBasePositionAndOrientation', ({(30, 42, 30, 52): 'self.ball2', (30, 54, 30, 67): '[4.0, 2.0, 0.5]', (30, 69, 30, 87): '(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))', True, 'import pybullet as p\n'), ((32, 8, 32, 68), 'pybullet.resetDebugVisualizerCamera', 'p.resetDebugVisualizerCamera', ({(32, 37, 32, 40): '(7.0)', (32, 41, 32, 45): '(90.0)', (32, 47, 32, 52): '(-43.0)', (32, 54, 32, 67): '(1.0, 1.0, 0.0)'}, {}), '(7.0, 90.0, -43.0, (1.0, 1.0, 0.0))', True, 'import pybullet as p\n'), ((36, 8, 36, 83), 'pybullet.resetBasePositionAndOrientation', 'p.resetBasePositionAndOrientation', ({(36, 42, 36, 48): 'wallId', (36, 50, 36, 62): '[0.0, -1.0, 0]', (36, 64, 36, 82): '(0.0, 0.0, 0.5, 0.5)'}, {}), '(wallId, [0.0, -1.0, 0], (0.0, 0.0, 0.5, 0.5))', True, 'import pybullet as p\n'), ((38, 8, 38, 82), 'pybullet.resetBasePositionAndOrientation', 'p.resetBasePositionAndOrientation', ({(38, 42, 38, 48): 'wallId', (38, 50, 38, 61): '[0.0, 1.0, 0]', (38, 63, 38, 81): '(0.0, 0.0, 0.5, 0.5)'}, {}), '(wallId, [0.0, 1.0, 0], (0.0, 0.0, 0.5, 0.5))', True, 'import pybullet as p\n'), ((40, 8, 40, 83), 'pybullet.resetBasePositionAndOrientation', 'p.resetBasePositionAndOrientation', ({(40, 42, 40, 48): 'wallId', (40, 50, 40, 62): '[3.0, -1.0, 0]', (40, 64, 40, 82): '(0.0, 0.0, 0.5, 0.5)'}, {}), '(wallId, [3.0, -1.0, 0], (0.0, 0.0, 0.5, 0.5))', True, 'import pybullet as p\n'), ((42, 8, 42, 82), 'pybullet.resetBasePositionAndOrientation', 'p.resetBasePositionAndOrientation', ({(42, 42, 42, 48): 'wallId', (42, 50, 42, 61): '[3.0, 1.0, 0]', (42, 63, 42, 81): '(0.0, 0.0, 0.5, 0.5)'}, {}), '(wallId, [3.0, 1.0, 0], (0.0, 0.0, 0.5, 0.5))', True, 'import pybullet as p\n'), ((44, 8, 44, 80), 'pybullet.resetBasePositionAndOrientation', 'p.resetBasePositionAndOrientation', ({(44, 42, 44, 48): 'wallId', (44, 50, 44, 61): '[1.0, 2.0, 0]', (44, 63, 44, 79): '(0.0, 0.0, 0.0, 1.0)'}, {}), '(wallId, [1.0, 2.0, 0], (0.0, 0.0, 0.0, 1.0))', True, 'import pybullet as p\n'), ((46, 8, 46, 81), 'pybullet.resetBasePositionAndOrientation', 'p.resetBasePositionAndOrientation', ({(46, 42, 46, 48): 'wallId', (46, 50, 46, 62): '[2.0, -2.0, 0]', (46, 64, 46, 80): '(0.0, 0.0, 0.0, 1.0)'}, {}), '(wallId, [2.0, -2.0, 0], (0.0, 0.0, 0.0, 1.0))', True, 'import pybullet as p\n'), ((100, 8, 100, 26), 'pybullet.stepSimulation', 'p.stepSimulation', ({}, {}), '()', True, 'import pybullet as p\n'), ((131, 8, 131, 26), 'pybullet.stepSimulation', 'p.stepSimulation', ({}, {}), '()', True, 'import pybullet as p\n'), ((35, 17, 35, 49), 'pybullet.loadSDF', 'p.loadSDF', ({(35, 27, 35, 48): '"""../models/walls.sdf"""'}, {}), "('../models/walls.sdf')", True, 'import pybullet as p\n'), ((37, 17, 37, 49), 'pybullet.loadSDF', 'p.loadSDF', ({(37, 27, 37, 48): '"""../models/walls.sdf"""'}, {}), "('../models/walls.sdf')", True, 'import pybullet as p\n'), ((39, 17, 39, 49), 'pybullet.loadSDF', 'p.loadSDF', ({(39, 27, 39, 48): '"""../models/walls.sdf"""'}, {}), "('../models/walls.sdf')", True, 'import pybullet as p\n'), ((41, 17, 41, 49), 'pybullet.loadSDF', 'p.loadSDF', ({(41, 27, 41, 48): '"""../models/walls.sdf"""'}, {}), "('../models/walls.sdf')", True, 'import pybullet as p\n'), ((43, 17, 43, 49), 'pybullet.loadSDF', 'p.loadSDF', ({(43, 27, 43, 48): '"""../models/walls.sdf"""'}, {}), "('../models/walls.sdf')", True, 'import pybullet as p\n'), ((45, 17, 45, 49), 'pybullet.loadSDF', 'p.loadSDF', ({(45, 27, 45, 48): '"""../models/walls.sdf"""'}, {}), "('../models/walls.sdf')", True, 'import pybullet as p\n'), ((87, 12, 87, 30), 'pybullet.stepSimulation', 'p.stepSimulation', ({}, {}), '()', True, 'import pybullet as p\n'), ((86, 31, 86, 87), 'robot.Robot', 'Robot', ({(86, 37, 86, 70): '[1.0 * i + 0.5, 1.0 * j - 0.5, 0.3]', (86, 72, 86, 77): '(2 * i + j)', (86, 79, 86, 86): 'self.dt'}, {}), '([1.0 * i + 0.5, 1.0 * j - 0.5, 0.3], 2 * i + j, self.dt)', False, 'from robot import Robot\n'), ((115, 23, 115, 48), 'numpy.linalg.norm', 'np.linalg.norm', ({(115, 38, 115, 47): '(pos1 - pos2)'}, {}), '(pos1 - pos2)', True, 'import numpy as np\n')]