matplotlib

import pybullet as p
import numpy as np
import pybullet_data
import time

import numpy as np
import matplotlib.pyplot as plt

plt.style.use('seaborn-poster')


# Initialize the physics simulation
physicsClient = p.connect(p.GUI)
p.setAdditionalSearchPath(pybullet_data.getDataPath())

# Load the objects into the simulation
planeId = p.loadURDF("plane.urdf")
sphereId = p.loadURDF("sphere_small.urdf", [0, 1, 0])
cylinderId = p.loadURDF("sphere_small.urdf", [0.5, 1, 0])
robotId = p.loadURDF("r2d2.urdf", [0.5,0.5,0.5],useFixedBase=True)

# Get the positions of the objects
spherePos, sphereOrn = p.getBasePositionAndOrientation(sphereId)
cylinderPos, cylinderOrn = p.getBasePositionAndOrientation(cylinderId)

# Calculate the distance between the two objects
distance = np.linalg.norm(np.array(spherePos) - np.array(cylinderPos))

print("The distance between the sphere and cylinder is: ", distance)


while True:
    # Get the image from the camera
    camera_pos = [0, 1, 2]
    viewMatrix = p.computeViewMatrix(
        cameraEyePosition=camera_pos,
        cameraTargetPosition=[0, 1, 0],
        cameraUpVector=[0, 1, 0])

    # Set the intrinsic parameters of the camera
    fov = 60
    aspect = 320.0 / 240.0
    near = 0.01
    far = 100

    # Get the image from the camera
    width, height, rgbImg, depthImg, segImg = p.getCameraImage(width=320, height=240, viewMatrix=viewMatrix,
                                                               projectionMatrix=p.computeProjectionMatrixFOV(fov,
                                                                                                             aspect,
                                                                                                             near, far))
    print("1")
    print(depthImg)
    print(len(depthImg))

    # Convert the depth image to a point cloud
    points = np.zeros((height * width, 3))
    for y in range(height):
        for x in range(width):
            depth = depthImg[y, x]
            if depth != 0:
                points[y * width + x, 0] = (x - 160) * depth / 160.0
                points[y * width + x, 1] = (y - 120) * depth / 120.0
                points[y * width + x, 2] = depth
    print("2")
    points.sort()
    print(points)
    print(len(points))
    x = np.zeros(len(points))
    y = np.zeros(len(points))
    z = np.zeros(len(points))
    for i in range(len(points)):
        x[i] = points[i,0]
        print(x)
        y[i] = points[i,1]
        z[i] = points[i,2]
        if z[i] > 0.00001:
            print([x[i],y[i],z[i]])
            print('ok')
        if i > 1 and z[i] - z[i - 1] > 0.2:
            z[i]

    fig = plt.figure(figsize=(10, 10))
    ax = plt.axes(projection='3d')
    ax.grid()

    ax.scatter(x, y, z, c='r', s=0.01)
    ax.set_title('3D Scatter Plot')

    # Set axes label
    ax.set_xlabel('x', labelpad=20)
    ax.set_ylabel('y', labelpad=20)
    ax.set_zlabel('z', labelpad=20)

    plt.show()
    # Find the points closest to the sphere and cylinder
    spherePoint = points[np.argmin(np.linalg.norm(points - np.array(spherePos), axis=1))]
    cylinderPoint = points[np.argmin(np.linalg.norm(points - np.array(cylinderPos), axis=1))]

    # Calculate the distance between the two objects based on the point cloud
    distanceFromCamera = np.linalg.norm(np.array(spherePoint)-np.array(cylinderPoint))
    print("The distance between the sphere and cylinder based on the camera is: ", distanceFromCamera)

    p.stepSimulation()
    time.sleep(1. / 240.)