import cv2
import numpy as np
import detectar_patron as dp
import commands
import sys
import optparse

import re
import rt_ref as rt
import mat_rot as mr
import ransac

from draw import *

PATRON=-1
CORNERS=0
MOSTRAR_0=0
S=5
W=7
H=4
path=""

#############################################################
#                       FUNCIONES                           #
#############################################################
#---------------- Cargar parametros instrinsecos ----------
def cargar_intrinsecos(Ncam):
	camera_matrix = [np.load(path+"/cam"+str(i)+"/mat_calib.npy") for i in range(Ncam)]
	dist_coeff = [np.load(path+"/cam"+str(i)+"/coef_dist.npy") for i in range(Ncam)]
	return camera_matrix,dist_coeff


#--------- Cantidad de camaras e imagenes ---------------------
# Aca se va a detectar la cantidad de camaras segun los nombres de las img guardadas
def leer_capturas():
	camList=commands.getoutput("ls -d "+path+"/cam*")
	camVecList=[str(x) for x in filter(None,re.split('[\n]',camList))]
	Ncam=len(camVecList)
	# Aca se va a leer la cantidad de img tomadas por camara
	imgList=commands.getoutput("ls "+path+"/cam0/img*")
	imgVecList=[str(x) for x in filter(None,re.split('[\n]',imgList))]
	Nimg=len(imgVecList)
	cam=np.zeros((Ncam),np.int)
	print "Camaras:",Ncam
	print "Capturas:",Nimg
	return cam,Ncam,Nimg

#-------- Coordenadas 3D del fiducial --------------------------
def coord_3D_fiducial(PATRON):
	if PATRON==1:
		p3D=np.float32([[-20.3,-12.05,0],[-20.3,11.6,0],[18.4,11.6,0],[18.4,-12.05,0],[13.1,-7.3,0],[13.1,6.8,0],[-15.5,6.8,0],[-15.5,0,0],[0,0,0],[0,-7.3,0]])
	elif PATRON==2:
		p3D=np.float32([[-29.75,-18,0],[-29.75,17.85,0],[26.9,17.85,0],[26.9,-18,0],[19.8,-10.8,0],[19.8,10.15,0],[-22.65,10.15,0],    [-22.65,0,0],[0,0,0],[0,-10.8,0]])
	elif PATRON==3:
		p3D = np.array([[S*int(i/W),S*int(i%W)*2+S*(int(i/W)%2),0.0] for i in range(W*H)])
	else:
		p3D = np.array([[S*int(i%W),S*int(i/W),0.0] for i in range(W*H)])		
	return p3D
#-------- Coordenadas 3D de los ejes ---------------------------
def coord_3D_ejes(PATRON):
	if PATRON==1 or PATRON==2:
		ejes=np.float32([[0,0,0],[20.0,0,0],[0,20.0,0],[0,0,-20.0]])
	elif PATRON==3:
		ejes=np.float32([[0,0,0],[40.0,0,0],[0,40.0,0],[0,0,40.0]])
	else:
		ejes=np.float32([[0,0,0],[50,0,0],[0,50,0],[0,0,-50]])
	return ejes

#-------- Mostrar puntos detectados ----------------------------	
def mostrar_esquinas(p2D):
	for k in range(len(p2D)):
		cv2.circle(img,tuple(p2D[k].ravel()),3,(0,0,255),-1)			
	   	cv2.putText(img,str(k),tuple(p2D[k].ravel()),1,1,(0,0,255),1)

#--------------- Dibujar ejes 3D --------------------------------
def dibujar_ejes(img,ejes,Rvec,Tvec,matriz_calib,coef_dist):			
	pd,jacob=cv2.projectPoints(ejes,Rvec,Tvec,matriz_calib,coef_dist)				
	project=np.int0(pd)
	cv2.line(img,tuple(project[0].ravel()),tuple(project[1].ravel()),(255,0,0),2)
	cv2.line(img,tuple(project[0].ravel()),tuple(project[2].ravel()),(0,255,0),2)
	cv2.line(img,tuple(project[0].ravel()),tuple(project[3].ravel()),(0,0,255),2)
	return img

#----------- Mostrar resultados T-R respecto al fiducial ---------
def mostrar_TR_0(found,Tvec,Rvec):
	if found==1:
		print "Detectado: SI"
		print "Traslacion: ("+str(Tvec[0,0])+", "+str(Tvec[1,0])+", "+str(Tvec[2,0])+")"		
		Rmat0=mr.MatRot(Rvec)
		Rvec0=mr.AngRot(Rmat0,1)		
		print "Rotacion: ("+str(Rvec0[0])+", "+str(Rvec0[1])+", "+str(Rvec0[2])+")"
	else:
		print "Detectado: NO"

#--------------------- Mostrar resultados -----------------
def mostrar_resultados(Ttotal,Rtotal,Htotal,T_error,R_error):
	print "\n\n================================================================"
	print "\t PARAMETROS EXTRINSECOS:\n"
	for i in range(Ncam):
		print "\t Camara "+str(i)+":\n"
		print "_Traslacion: T("+str(Ttotal[i,0])+", "+str(Ttotal[i,1])+", "+str(Ttotal[i,2])+") --> error = "+str(T_error[i])+"\n"
		print "_Rotacion:   R("+str(Rtotal[i,0])+", "+str(Rtotal[i,1])+", "+str(Rtotal[i,2])+") --> error = "+str(R_error[i])+"\n"
		print "_H: "
		print Htotal[i],"\n"
		print "\n-------------------------------------------------------------\n"
#------------- Graficar Camaras --------------------------
def graficar_camaras(Htotal):
	fig = plt.figure('Calibracion Extrinseca')
	ax = fig.add_subplot(111, projection='3d')
	for j in range(0,Ncam):
		ax = cam_draw(ax,Htotal[j],"cam"+str(j))

	xlim = ax.get_xlim()
	ylim = ax.get_ylim()
	zlim = ax.get_zlim()
	ax = set_limits(np.array([xlim,ylim,zlim]),ax)
	plt.show()

# ---- Calculo de parametros extrinsecos entre 2 camaras ---------
def calc_stereo_extrinsic(Tvec_cam,Rvec_cam,Tvec_cam_ref,Rvec_cam_ref,found_cam,found_cam_ref):
	Tvec_ref = []
	Rvec_ref = []
	inliers = []
	for i in range(Nimg):

		if found_cam[i]==1 and found_cam_ref[i]==1:
			RmatRef,TvecRef=rt.rt_ref(Rvec_cam_ref[i],Tvec_cam_ref[i],Rvec_cam[i],Tvec_cam[i])                    
			RvecRef=mr.AngRot(RmatRef,1)
			Tvec_ref.append(TvecRef)
			Rvec_ref.append(RvecRef)
			inliers.append(1)			
			
	if len(inliers) == 0:
		vec = np.empty((3))
		vec[:] = np.NAN
		return vec,vec,vec,vec
	else:
		Tvec_ref=np.array(Tvec_ref)
		Rvec_ref=np.array(Rvec_ref)
		inliers=np.array(inliers)
		Ttotal=np.zeros((3),np.float32)
		Rtotal=np.zeros((3),np.float32)
		T_error=np.zeros((3),np.float32)
		R_error=np.zeros((3),np.float32)

		inliers=ransac.buscar(np.array(inliers),np.array(Tvec_ref),1)
		inliers=ransac.buscar(np.array(inliers),np.array(Rvec_ref),1)

		Ttotal,T_error=ransac.minimizar(np.array(Tvec_ref),np.array(inliers))
		Rtotal,R_error=ransac.minimizar(np.array(Rvec_ref),np.array(inliers))
		return Ttotal,Rtotal,T_error,R_error

#-------------- cambio de sistema de coordenadas respecto a la cam0 -----------------

def cambio_ref(cam,H_base,Htotal,Ttotal,Rtotal,nodo_h):
	#Htotal[cam] = np.dot(H_base,mr.get_homog(Ttotal[cam],Rtotal[cam]))
	Htotal[cam] = np.dot(H_base,mr.get_homog(Ttotal[cam],np.array([Rtotal[cam][0],-Rtotal[cam][1],-Rtotal[cam][2]])))
	for n in nodo_h[cam]:
		Htotal = cambio_ref(n,Htotal[cam],Htotal,Ttotal,Rtotal,nodo_h)	
	return Htotal

#------------- arbol de correspondencia entre camaras -----------------------
def crear_arbol(cam_p, nodo_p, nodo_h, dist, chk):
	if len(nodo_h[cam_p]) > 0:
		while 0 in [chk[p] for p in nodo_h[cam_p]]:
			id_h = [chk[p] for p in nodo_h[cam_p]].index(0)
			D = dist[cam_p,nodo_h[cam_p][id_h]]
			for i in range(len(nodo_h[cam_p])):
				if dist[cam_p,nodo_h[cam_p][i]]*1.2 < D and chk[nodo_h[cam_p][i]]==0:
					D = dist[cam_p,nodo_h[cam_p][i]]
					id_h = i
			cam_h = nodo_h[cam_p][id_h]
			chk[cam_h] = 1
			update_nodo_h = [cam_h]
			for i in range(len(nodo_h[cam_p])):
				if i!=id_h:
					if dist[cam_h,nodo_h[cam_p][i]]*1.2 < dist[cam_p,nodo_h[cam_p][i]] and chk[nodo_h[cam_p][i]]==0:
						nodo_p[nodo_h[cam_p][i]] = cam_h
						nodo_h[cam_h].append(nodo_h[cam_p][i])
					else:
						update_nodo_h.append(nodo_h[cam_p][i])
			nodo_h[cam_p] = update_nodo_h
			nodo_p,nodo_h,chk = crear_arbol(cam_h,nodo_p,nodo_h,dist,chk)
	return nodo_p, nodo_h, chk

# ----------- calculo del error acumulado -----------------------------------
def acumular_error(cam,Te_prev,Re_prev,T_error,R_error,nodo_h):
	T_error[cam] += Te_prev
	R_error[cam] += Re_prev
	for i in nodo_h[cam]:
		T_error,R_error = acumular_error(i,T_error[cam],R_error[cam],T_error,R_error,nodo_h)
	return T_error,R_error


######################################################################
#                        PARAMETROS                                  #        
######################################################################


desc = """ Este script permite calcular los parametros extrinsecos de un sistema multicamara, utilizando la combinacion de pares estereo. Para mas informacion sobre como usar el script ingrese a: https://ciii.frc.utn.edu.ar/Vision/MultiCameraCalibration. 
"""
epi = """ Para eligir el tipo de fiducial utilice:
 -f CHESS		 si se usa un chessboard.
 -f CIRCLE		 si se usa un patron de circulos.
 -f L1			 si se usa un patron L con dimensiones A3.
 -f L2			 si se usa un patron L con dimensiones A2.
"""
parser = optparse.OptionParser(add_help_option=False,description=desc)
parser.add_option('-p', '--path', help='path del dataset',action='store',dest='path',type='string')
parser.add_option('-v', '--view', help='ver los resultados de la calibracion individual',default=False, action='store_true',dest='view')
parser.add_option('-f', '--fid', help='fiducial utilizado en el dataset',action='store',dest='fid',type='string')
parser.add_option('-W', '--width', help='cantidad de puntos horizontales en caso de usar un chessboard o un patron de circulos',action='store',dest='width',type='int')
parser.add_option('-H', '--height', help='cantidad de puntos verticales en caso de usar un chessboard o un patron de circulos',action='store',dest='height',type='int')
parser.add_option('-S', '--size', help='ancho de cada cuadrado en caso de usar un chessboard o distancia entre centros en un patron de circulos',action='store',dest='size',type='float')
parser.add_option('-h', '--help', dest='help',help='muestra este mensaje de ayuda',action='store_true')

(opts, args) = parser.parse_args()

if opts.help==True:
	parser.print_help()
	print epi
	exit(-1)

if opts.fid is None:
	print "Debe especificarse el tipo de fiducial presente en el dataset."
	parser.print_help()
	print epi
	exit(-1)
elif (opts.fid!='CHESS')and(opts.fid!='CIRCLE')and(opts.fid!='L1')and(opts.fid!='L2')and opts.fid is not None:
	print "El tipo de fiducial seleccionado no esta permitido."
	parser.print_help()
	print epi
	exit(-1)
else:	PATRON = opts.fid

if (opts.fid=='CHESS')and((opts.width is None)or(opts.height is None)or(opts.size is None)):
	print "Debe especificar las dimensiones del chessboard."
	parser.print_help()
	print epi
	exit(-1)

if (opts.fid=='CIRCLE')and((opts.width is None)or(opts.height is None)or(opts.size is None)):
	print "Debe especificar las dimensiones del patron de circulos."
	parser.print_help()
	print epi
	exit(-1)

if ((opts.fid=='CIRCLE')or(opts.fid=='CHESS')):
	W = opts.width
	H = opts.height
	S = opts.size

if opts.view==True:
	MOSTRAR_0 = 1

if opts.fid=='CIRCLE':
	PATRON = 3
elif opts.fid=='CHESS':
	PATRON = 0
elif opts.fid=='L1':
	PATRON = 1
else:
	PATRON = 2

if opts.path is None:
	print "Debe especificar el path del data set."
	parser.print_help()
	print epi
	exit(-1)
else:
	path = opts.path

######################################################################
#                        VARIABLES                                   #
######################################################################
cam_index,Ncam,Nimg=leer_capturas()
#camera_matrix=np.zeros((3,3,Ncam),np.float)
#dist_coeff=np.zeros((1,5,Ncam),np.float)
RvecRef=np.zeros((3),np.float32)
TvecRef=np.zeros((3),np.float32)   
Rvec=np.zeros((Nimg,Ncam,3,1),np.float32)
Tvec=np.zeros((Nimg,Ncam,3,1),np.float32)      
found=np.zeros((Nimg,Ncam,1),np.int)
T=np.zeros((Nimg,Ncam,1),np.float32)
Q=np.zeros((Nimg,Ncam,1),np.float32)
p2D=np.zeros((10,2),np.float32)          
p3D=coord_3D_fiducial(PATRON)
ejes=coord_3D_ejes(PATRON)
project=np.zeros((4,2),np.int)
criteria=(cv2.TERM_CRITERIA_EPS+cv2.TERM_CRITERIA_MAX_ITER,30,0.001)          
camera_matrix,dist_coeff=cargar_intrinsecos(Ncam)
u1=u2=128
zoom=1

#########################################################################################  
#------ Param. extrinsecos respecto a la camara que tomo la imagen ---------------------
if MOSTRAR_0==1:
	print "==================================================================="     
	print "1_ Parametros Extrinsecos: Respecto a la camara que tomo la imagen\n"
for i in range(Nimg):
     for j in range(Ncam):
		img=cv2.imread(path+"/cam"+str(j)+"/img"+str(i)+"_"+str(j)+".jpg",-1)
		if PATRON==0:			# patron chessboard
			gris=cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
			found[i,j],p2D=cv2.findChessboardCorners(gris,(W,H))
			if found[i,j]==1:
				cv2.cornerSubPix(gris,p2D,(11,11),(-1,-1),criteria)
		elif PATRON == 3:		# patron circulos
			gris = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
			found[i,j],p2D = cv2.findCirclesGrid(gris,(3,9),flags=cv2.CALIB_CB_ASYMMETRIC_GRID)
		elif (PATRON==1)or(PATRON==2):	# patron L	
			found[i,j],p2D,u1,u2=dp.detectarPatron(img,u1,u2,zoom)		
		if found[i,j]==1:
			if CORNERS==1:
				mostrar_esquinas(p2D)
			ret,Rvec[i,j,:],Tvec[i,j,:]=cv2.solvePnP(p3D,p2D,camera_matrix[j],dist_coeff[j])
			T[i,j]=np.linalg.norm(Tvec[i,j])
			img=dibujar_ejes(img,ejes,Rvec[i,j,:],Tvec[i,j,:],camera_matrix[j],dist_coeff[j])
		if MOSTRAR_0==1:
			cv2.imshow("img"+str(i)+"_"+str(j),img)
			print "\nImagen: img"+str(i)+"_"+str(j)
			mostrar_TR_0(found[i,j],Tvec[i,j],Rvec[i,j])
			cv2.waitKey()
			cv2.destroyAllWindows()

#----------------- Param. extrinsecos entre camaras ---------------------

nodo_p = [-2 for i in range(Ncam)]
nodo_h = [[] for i in range(Ncam)]

Htotal = [np.identity(4) for i in range(Ncam)]
Ttotal = np.zeros((Ncam,3),np.float32)
Rtotal = np.zeros((Ncam,3),np.float32)
T_error = np.zeros((Ncam,3),np.float32)
R_error = np.zeros((Ncam,3),np.float32)

#---------------------- 1ra estimacion -----------------------------------
# Se calcula la posicion de las camaras formando pares stereo. En cada par
# se indica una camara como referencia. La referencia global es la cam0.
nodo_p[0] = -1

for cam_ref in range(Ncam):
	for cam in range(Ncam):
		if cam_ref!=cam and nodo_p[cam]==-2:
			Ttotal[cam],Rtotal[cam],T_error[cam],R_error[cam] = calc_stereo_extrinsic(Tvec[:,cam],Rvec[:,cam],Tvec[:,cam_ref],Rvec[:,cam_ref],found[:,cam],found[:,cam_ref])
			if (str(Ttotal[cam][0])!='nan') and (str(Rtotal[cam][0])!='nan'):
				nodo_p[cam] = cam_ref
				nodo_h[cam_ref].append(cam)


H_base = np.identity(4)

#--------------- Matriz homografica de la camara 0 ---------------------------------------
# Debe indicarse la traslacion y rotacion de la camara 0, respecto del sistema global
# La matriz Ho mostrada a continuacion puede usarse teniendo como referencia la esquina 
# del salon.

Ho = np.identity(4)
#Ho = np.array([[-7.03203813e-01, -4.83828169e-01, 5.20974760e-01, -8.18166372e+00],
# [ 7.04680287e-01, -5.71676364e-01, 4.20252099e-01, 1.32402279e+01],
# [ 9.44991531e-02, 6.62643522e-01, 7.42949172e-01, -3.11059947e+02],
# [ 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 1.00000000e+00]])

Htotal = cambio_ref(0,H_base,Htotal,Ttotal,Rtotal,nodo_h)

for j in range(len(Htotal)):
	Rtotal[j] = mr.AngRot(Htotal[j][:3,:3],1)
	Ttotal[j] = Htotal[j][:3,3]

#--------------------- 2da estimacion ------------------------------------
# Una vez que se tiene una estimacion de la posicion de cada camara, se 
# cambian las referencias usadas, de modo de reducir las  distancias de 
# los pares stereo
 
dist = np.zeros((Ncam,Ncam),np.float32)
for i in range(Ncam):
	for j in range(Ncam):
		dist[i,j] = np.linalg.norm(Ttotal[i]-Ttotal[j])
chk = [0 for i in range(Ncam)]
chk[0] = 1
print 
nodo_p,nodo_h,chk = crear_arbol(0,nodo_p,nodo_h,dist,chk)

for cam in range(Ncam):
	cam_ref = nodo_p[cam]
	if cam_ref>=0:
		Ttotal[cam],Rtotal[cam],T_error[cam],R_error[cam] = calc_stereo_extrinsic(Tvec[:,cam],Rvec[:,cam],Tvec[:,cam_ref],Rvec[:,cam_ref],found[:,cam],found[:,cam_ref])
	

Htotal = cambio_ref(0,H_base,Htotal,Ttotal,Rtotal,nodo_h)

for j in range(len(Htotal)):
	Htotal[j] = np.dot(Ho,Htotal[j])
	H = Htotal[j]
	Rtotal[j] = mr.AngRot(Htotal[j][:3,:3],1)
	Ttotal[j] = H[:3,3]

T_error,R_error = acumular_error(0,[0,0,0],[0,0,0],T_error,R_error,nodo_h)

#------------------ resultados --------------------------------
mostrar_resultados(Ttotal,Rtotal,Htotal,T_error,R_error)
graficar_camaras(Htotal)
cv2.waitKey()
cv2.destroyAllWindows()





