import numpy as np
import matplotlib.pyplot as plt
import os.path
import json
import argparse
import scipy
import math
import pylab
from sklearn.preprocessing import normalize
from mpl_toolkits.mplot3d import Axes3D
import cv2
import caffe
from Quaternion import Quat
import math

# Make sure that caffe is on the python path:
caffe_root = '/home/ubuntu/kagan-posenet/caffe-posenet/'  # Change to your directory to caffe-posenet
import sys
sys.path.insert(0, caffe_root + 'python')


def quaternion_to_euler_angle(w, x, y, z):
	ysqr = y * y
	
	t0 = +2.0 * (w * x + y * z)
	t1 = +1.0 - 2.0 * (x * x + ysqr)
	X = math.degrees(math.atan2(t0, t1))
	
	t2 = +2.0 * (w * y - z * x)
	t2 = +1.0 if t2 > +1.0 else t2
	t2 = -1.0 if t2 < -1.0 else t2
	Y = math.degrees(math.asin(t2))
	
	t3 = +2.0 * (w * z + x * y)
	t4 = +1.0 - 2.0 * (ysqr + z * z)
	Z = math.degrees(math.atan2(t3, t4))
	
	return X, Y, Z



parser = argparse.ArgumentParser()
parser.add_argument('--model', type=str, required=True)
parser.add_argument('--weights', type=str, required=True)
parser.add_argument('--iter', type=int, required=True)
args = parser.parse_args()

results = np.zeros((args.iter,2))

caffe.set_mode_gpu()

net = caffe.Net(args.model,
                args.weights,
                caffe.TEST)



for i in range(0, args.iter):


	img = cv2.imread('./dataset_test/photo_' + str(i) + '.png', 1)
	print 'photo_' + str(i) + '.png'
	
	#cv2.imshow('Tested_Image/photo_' + str(i) + '.png', img)
	#key = cv2.waitKey(0)

	mean_vec = np.average(np.average(img, axis=0), axis=0)
	# load input and configure preprocessing
	transformer = caffe.io.Transformer({'data': net.blobs['data'].data.shape})
	transformer.set_mean('data',mean_vec )
	transformer.set_transpose('data', (2,0,1))
	transformer.set_channel_swap('data', (2,1,0))
	transformer.set_raw_scale('data', 255.0)
	#note we can change the batch size on-the-fly
	#since we classify only one image, we change batch size from 10 to 1
	net.blobs['data'].reshape(1,3,224,224)

	#load the image in the data layer
	im = caffe.io.load_image('./dataset_test/photo_' + str(i) + '.png')
	net.blobs['data'].data[...] = transformer.preprocess('data', im)

	#compute
	net.forward()
	predicted_x = net.blobs['cls3_fc_xyz'].data 
	predicted_q = net.blobs['cls3_fc_wpqr'].data 
	predicted_q = np.squeeze(predicted_q)
	predicted_x = np.squeeze(predicted_x)

	#print 'predicted_x = ', predicted_x, 'predicted_angle = ', angle

	#print 'Iteration:  ', i, '  Error XYZ (m):  ', error_x, '  Error Q (degrees):  ', theta

	with open('./dataset_test/dataset_test_quat.txt', 'r') as f:
		for j, line in enumerate(f):			
			if j == i:
				word = line.split()
	        	#pose_q2 = float(word[3])
	        	pose_q = np.array([float(word[4]), float(word[5]), float(word[6]), float(word[7])])
	        	pose_x = np.array([float(word[1]), float(word[2]), float(word[3])])

	#pose_q= net.blobs['label_wpqr'].data
	#pose_x= net.blobs['label_xyz'].data

	print 'pose_x= ', pose_x , 'pose_q= ', pose_q
	print 'predicted_x = ', predicted_x , 'predicted_q = ' , predicted_q

	predicted_angle = predicted_q
	actual_angle = pose_q 
	#actual_angle = actual_angle * 180 / math.pi #in degrees
	
	actual_angle = quaternion_to_euler_angle(actual_angle[0], actual_angle[1], actual_angle[2], actual_angle[3]) #in degrees
	predicted_angle = quaternion_to_euler_angle(predicted_angle[0], predicted_angle[1], predicted_angle[2], predicted_angle[3]) #in degrees
	L = 1
	act_vec = [pose_x[0], pose_x[1], (pose_x[0] + L*math.cos(actual_angle[0]*math.pi/180)), (pose_x[1] + L*math.sin(actual_angle[0]*math.pi/180))]
	pred_vec = [predicted_x[0], predicted_x[1], (predicted_x[0] + L*math.cos(predicted_angle[0]*math.pi/180)), (predicted_x[1] + L*math.sin(predicted_angle[0]*math.pi/180))]

	q1 = pose_q / np.linalg.norm(pose_q)
	q2 = predicted_q / np.linalg.norm(predicted_q)
	d = abs(np.sum(np.multiply(q1,q2)))
	theta = 2 * np.arccos(d) * 180/math.pi
	error_x = np.linalg.norm(pose_x-predicted_x)

	results[i,:] = [error_x,theta]

	print 'Iteration:  ', i, '  Error XYZ (m):  ', error_x, '  Error Q (degrees):  ', theta

mean_result = np.mean(results, axis=0)
median_result = np.median(results,axis=0)
print 'Median error ', median_result[0], 'm  and ', median_result[1], 'degrees.', ' Mean error = ', mean_result[0], 'm and ', mean_result[1], 'degrees'

np.savetxt('results_talha.txt', results, delimiter=' ')

print 'Success!'



	#print 'act_vec = ', act_vec, 'pred_vec= ', pred_vec
	#print 'actual_angle= ', actual_angle, 'predicted_angle= ', predicted_angle[0]
'''
	ax = plt.axes()
	ax.arrow(pose_x[0], pose_x[1], (L*math.cos(actual_angle*math.pi/180)), (L*math.sin(actual_angle*math.pi/180)), head_width=0.07, head_length=0.07, fc='g', ec='g', lw=7)
	ax.arrow(predicted_x[0], predicted_x[1], (L*math.cos(predicted_angle[0]*math.pi/180)), (L*math.sin(predicted_angle[0]*math.pi/180)), head_width=0.07, head_length=0.07, fc='r', ec='r', lw=7)
	ax.set_xlim([-8, 0])
	ax.set_ylim([-3, -6])
	
	cv2.imshow('Tested_Image/photo_' + str(i) + '.png', img)
	plt.grid()
	plt.show()
	cv2.waitKey(0)

	#key = cv2.waitKey(0)
	#if key == 32:
	#	print 'Drawing vectors'
	#	plt.show()
	#if key == 27:
	#	print 'Skipping to next image'
	#	continue
	#if key == 27:
	#	print 'Exitting the program' 
	#	break
'''