import pybullet as p
import math
import numpy as np



p.connect(p.GUI)
plane = p.loadURDF("plane100.urdf")
cube = p.loadURDF("cube.urdf",[0,0,1])

def getRayFromTo(mouseX,mouseY):
        width, height, viewMat, projMat, cameraUp, camForward, horizon,vertical, _,_,dist, camTarget = p.getDebugVisualizerCamera()
        camPos = [camTarget[0] - dist*camForward[0],camTarget[1] - dist*camForward[1],camTarget[2] - dist*camForward[2]]
        farPlane = 10000
        rayForward = [(camTarget[0]-camPos[0]),(camTarget[1]-camPos[1]),(camTarget[2]-camPos[2])]
        lenFwd = math.sqrt(rayForward[0]*rayForward[0]+rayForward[1]*rayForward[1]+rayForward[2]*rayForward[2])
        invLen = farPlane*1./lenFwd
        rayForward = [invLen*rayForward[0],invLen*rayForward[1],invLen*rayForward[2]]
        rayFrom = camPos
        oneOverWidth = float(1)/float(width)
        oneOverHeight = float(1)/float(height)

        dHor = [horizon[0] * oneOverWidth,horizon[1] * oneOverWidth,horizon[2] * oneOverWidth]
        dVer = [vertical[0] * oneOverHeight,vertical[1] * oneOverHeight,vertical[2] * oneOverHeight]
        rayToCenter=[rayFrom[0]+rayForward[0],rayFrom[1]+rayForward[1],rayFrom[2]+rayForward[2]]
        ortho=[- 0.5 * horizon[0] + 0.5 * vertical[0]+float(mouseX)*dHor[0]-float(mouseY)*dVer[0],
             - 0.5 * horizon[1] + 0.5 * vertical[1]+float(mouseX)*dHor[1]-float(mouseY)*dVer[1],
             - 0.5 * horizon[2] + 0.5 * vertical[2]+float(mouseX)*dHor[2]-float(mouseY)*dVer[2]]

        rayTo = [rayFrom[0]+rayForward[0]  +ortho[0],
                                        rayFrom[1]+rayForward[1]  +ortho[1],
                                        rayFrom[2]+rayForward[2]  +ortho[2]]
        lenOrtho = math.sqrt(ortho[0]*ortho[0]+ortho[1]*ortho[1]+ortho[2]*ortho[2])
        alpha = math.atan(lenOrtho/farPlane)
        return rayFrom,rayTo, alpha


width, height, viewMat, projMat, cameraUp, camForward, horizon,vertical, _,_,dist, camTarget = p.getDebugVisualizerCamera()
camPos = [camTarget[0] - dist*camForward[0],camTarget[1] - dist*camForward[1],camTarget[2] - dist*camForward[2]]
farPlane = 10000
rayForward = [(camTarget[0]-camPos[0]),(camTarget[1]-camPos[1]),(camTarget[2]-camPos[2])]
lenFwd = math.sqrt(rayForward[0]*rayForward[0]+rayForward[1]*rayForward[1]+rayForward[2]*rayForward[2])
oneOverWidth = float(1)/float(width)
oneOverHeight = float(1)/float(height)
dHor = [horizon[0] * oneOverWidth,horizon[1] * oneOverWidth,horizon[2] * oneOverWidth]
dVer = [vertical[0] * oneOverHeight,vertical[1] * oneOverHeight,vertical[2] * oneOverHeight]
        
lendHor = math.sqrt(dHor[0]*dHor[0]+dHor[1]*dHor[1]+dHor[2]*dHor[2]) 
lendVer = math.sqrt(dVer[0]*dVer[0]+dVer[1]*dVer[1]+dVer[2]*dVer[2])

cornersX = [0,width,width,0]
cornersY = [0,0,height,height]
corners3D=[]

imgW = int(width/10)
imgH = int(height/10)

img = p.getCameraImage(imgW,imgH, renderer=p.ER_BULLET_HARDWARE_OPENGL)
rgbBuffer=img[2]
depthBuffer = img[3]
print("rgbBuffer.shape=",rgbBuffer.shape)
print("depthBuffer.shape=",depthBuffer.shape)

#disable rendering temporary makes adding objects faster
p.configureDebugVisualizer(p.COV_ENABLE_RENDERING,0)
p.configureDebugVisualizer(p.COV_ENABLE_GUI,0)
p.configureDebugVisualizer(p.COV_ENABLE_TINY_RENDERER,0)
visualShapeId = p.createVisualShape(shapeType=p.GEOM_SPHERE, rgbaColor=[1,1,1,1], radius=0.03 )
collisionShapeId = -1 #p.createCollisionShape(shapeType=p.GEOM_MESH, fileName="duck_vhacd.obj", collisionFramePosition=shift,meshScale=meshScale)

for i in range (4):
	w = cornersX[i]
	h = cornersY[i]
	rayFrom,rayTo, _= getRayFromTo(w,h)
	rf = np.array(rayFrom)
	rt = np.array(rayTo)
	vec = rt-rf
	l = np.sqrt(np.dot(vec,vec))
	newTo = (0.01/l)*vec+rf
	#print("len vec=",np.sqrt(np.dot(vec,vec)))
			
	p.addUserDebugLine(rayFrom,newTo,[1,0,0])
	corners3D.append(newTo)
count = 0

stepX=5
stepY=5
for w in range(0,imgW,stepX):
	for h in range (0,imgH,stepY):
		count+=1
		if ((count % 100)==0):
			print(count,"out of ", imgW*imgH/(stepX*stepY))
		rayFrom,rayTo, alpha = getRayFromTo(w*(width/imgW),h*(height/imgH))
		rf = np.array(rayFrom)
		rt = np.array(rayTo)
		vec = rt-rf
		l = np.sqrt(np.dot(vec,vec))
		depthImg = float(depthBuffer[h,w])
		far=1000.
		near=0.01
		depth = far * near / (far - (far - near) * depthImg)
		depth/=math.cos(alpha)
		newTo = (depth/l)*vec+rf
		p.addUserDebugLine(rayFrom,newTo,[1,0,0])
		mb = p.createMultiBody(baseMass=0,baseCollisionShapeIndex=collisionShapeId, baseVisualShapeIndex = visualShapeId, basePosition = newTo, useMaximalCoordinates=True)
		color = rgbBuffer[h,w]
		color=[color[0]/255.,color[1]/255.,color[2]/255.,1]
		p.changeVisualShape(mb,-1,rgbaColor=color)
p.addUserDebugLine(corners3D[0],corners3D[1],[1,0,0])
p.addUserDebugLine(corners3D[1],corners3D[2],[1,0,0])
p.addUserDebugLine(corners3D[2],corners3D[3],[1,0,0])
p.addUserDebugLine(corners3D[3],corners3D[0],[1,0,0])
p.configureDebugVisualizer(p.COV_ENABLE_RENDERING,1)
print("ready\n")
#p.removeBody(plane)
#p.removeBody(cube)
while (1):
	p.setGravity(0,0,-10)