'''
it is an autoencoder that learns a latent variable model for its input data. 
So instead of letting your neural network learn an arbitrary function, 
you are learning the parameters of a probability distribution modeling your data. 
If you sample points from this distribution, you can generate new input data 
samples: a VAE is a "generative model".
'''
import numpy as np
import matplotlib.pyplot as plt
import os
from scipy.stats import norm
import glob
from scipy import misc
from PIL import Image
from keras.layers import Input, Dense, Lambda, Layer, Concatenate, concatenate, Reshape, Flatten
from keras import optimizers, losses
from keras.models import Model
from keras.layers.convolutional import Conv2D, Conv2DTranspose
from keras import backend as K
from keras import metrics
from keras.datasets import mnist


# root_x_train_lpath="/Users/anurag/Documents/Anurag/ASU-Courses/CSE591-DeepLearning/Project/KITTI_dataset/data_scene_flow_multiview/training/image_2_resized" # x_train
# root_x_test_lpath="/Users/anurag/Documents/Anurag/ASU-Courses/CSE591-DeepLearning/Project/KITTI_dataset/data_scene_flow_multiview/testing/image_2_resized" # x_test
# root_x_train_rpath="/Users/anurag/Documents/Anurag/ASU-Courses/CSE591-DeepLearning/Project/KITTI_dataset/data_scene_flow_multiview/training/image_3_resized" # y_train
# root_x_test_rpath="/Users/anurag/Documents/Anurag/ASU-Courses/CSE591-DeepLearning/Project/KITTI_dataset/data_scene_flow_multiview/testing/image_3_resized"	# y_test


cwd = os.getcwd()
root_x_train_lpath="/home/asolank1/deepProject/training/image_2_resized" # x_train
root_x_test_lpath="/home/asolank1/deepProject/testing/image_2_resized"   # x_test
root_x_train_rpath="/home/asolank1/deepProject/training/image_3_resized" # y_train
root_x_test_rpath="/home/asolank1/deepProject/testing/image_3_resized"  # y_test

num_images = 9

# x_train = np.ndarray(shape=(4200, 81000), dtype=float)
# # (4200, 1397250)
# x_test = np.ndarray(shape=(4189, 81000), dtype=float)
# # (4189 , 1397250)
# y_train = np.ndarray(shape=(4200, 81000), dtype=float)
# # (4200 , 1397250)
# y_test = np.ndarray(shape=(4189, 81000), dtype=float)
# # (4189 , 1397250)

x_train = [] #np.empty(shape=(0, 0), dtype=float)
# (4200, 1397250)
x_test = [] #np.empty(shape=(5, 81000), dtype=float)
# (4189 , 1397250)
y_train = [] #np.empty(shape=(5, 81000), dtype=float)
# (4200 , 1397250)
y_test = [] #np.empty(shape=(5, 81000), dtype=float)
# (4189 , 1397250)


############################
#      Image resizing code below
############################
"""
path = "/Users/anurag/Documents/Anurag/ASU-Courses/CSE591-DeepLearning/Project/KITTI_dataset/data_scene_flow_multiview/testing/image_2/"
dirs = os.listdir( path )
path2save="/Users/anurag/Documents/Anurag/ASU-Courses/CSE591-DeepLearning/Project/KITTI_dataset/data_scene_flow_multiview/testing/image_2_resized/"
#print dirs
basewidth = 300
def resize():
	for item in dirs:
		print path+item
		if os.path.isfile(path+item) and not item.startswith('.'):
			im = Image.open(path+item)
			wpercent = (basewidth/float(im.size[0]))
			hsize = int((float(im.size[1])*float(wpercent)))
			#f, e = os.path.splitext(path+item)
			im = im.resize((basewidth,hsize), Image.ANTIALIAS)
			im.save(path2save +item , 'JPEG', quality=90)

resize()
exit(0)
"""
# REDUCED image size = (90, 300)
#print os.listdir(root_x_train_lpath)
# reading train left images 
i = 0
for image_path in glob.glob(root_x_train_lpath+"/*.png"):
	 if i <= num_images:
		image = misc.imread(image_path)
		image = image.astype(np.float32)
		#image = image.reshape(image.shape[0]*image.shape[1]*image.shape[2], 1)
		x_train.append(image)
		i += 1
x_train = np.array(x_train)
# print x_train.shape
#x_train = np.reshape(x_train, (x_train.shape[0], x_train.shape[1]))
#print x_train.shape
x_train = x_train/255.0
print "x_train shape", x_train.shape

# reading train right images
i = 0
for image_path in glob.glob(root_x_train_rpath+"/*.png"):
	if i <= num_images:
		image = misc.imread(image_path)
		image = image.astype(np.float32)
		#image = image.reshape(image.shape[0]*image.shape[1]*image.shape[2], 1)
		y_train.append(image)
		i += 1
y_train = np.array(y_train)
#y_train = np.reshape(y_train, (y_train.shape[0], y_train.shape[1]))
y_train = y_train/255.0		
print "y_train shape", y_train.shape

# reading test Left images 
i = 0
for image_path in glob.glob(root_x_test_lpath+"/*.png"):
	if i <= num_images:
		# pth = os.path.abspath('.')+'/test.jpg'
		image = misc.imread(image_path)
		image = image.astype(np.float32)
		#image = image.reshape(image.shape[0]*image.shape[1]*image.shape[2], 1)
		x_test.append(image)
		i += 1
x_test = np.array(x_test)
# print "x_test shape", x_test.shape
#x_test = np.reshape(x_test, (x_test.shape[0], x_test.shape[1]))	
x_test = x_test/255.0
print "x_test shape", x_test.shape

# reading test right images 
i = 0
for image_path in glob.glob(root_x_test_rpath+"/*.png"):
	if i <= num_images:
		image = misc.imread(image_path)
		image = image.astype(np.float32)
		#image = image.reshape(image.shape[0]*image.shape[1]*image.shape[2], 1)
		y_test.append(image)
		i += 1
y_test = np.array(y_test)
#y_test = np.reshape(y_test, (y_test.shape[0], y_test.shape[1]))	
y_test = y_test/255.0
print "y_test shape", y_test.shape

# exit(0)
batch_size = num_images+1
#original_dim = 1397250
original_dim = 81000
latent_dim = 128																																																																																																																																																																																																																																																																																																																
intermediate_dim = 256
nb_epochs = 50																																																																																																																																																																																											
epsilon_std = 1.0
theta_dim = 20

############################
#	ENCODER
############################
x = Input(shape=(90,300,3))
# print x.shape
h = Conv2D(intermediate_dim, (5,5), strides=(2,2), padding='same', activation='relu')(x)
# print h.layer()
#z_int = Flatten()(h)
z_1 = Conv2D(latent_dim, (5,5), strides=(2,2), padding='same', activation='relu')(h)
# z_log_var = Dense(latent_dim, activation='relu')(h)
z_2 = Flatten()(z_1)
z = Dense(latent_dim, activation='relu')(z_2)
#y_train = Input(batch_shape=(batch_size, original_dim))

def sampling(args):
    z_mean, z_log_var = args
    epsilon = K.random_normal(shape=(latent_dim, ), mean=0., stddev=epsilon_std)
    return z_mean + K.exp(z_log_var / 2) * epsilon


# note that "output_shape" isn't necessary with the TensorFlow backend
#z = Lambda(sampling, output_shape=(latent_dim, ))([z_mean, z_log_var])
#z = Lambda(sampling)([z_mean, z_log_var])
#z_dense = Dense(latent_dim, activation='linear')(h)


# print "z, shape", z_mean.output_shape

theta = np.random.rand(x_train.shape[0], theta_dim)
# print "============",theta.shape
theta_h = Input(shape=(theta_dim, ))
#theta_h_reshaped = Reshape((latent_dim,))(theta_h)
theta_h_dense = Dense(theta_dim, activation='relu')(theta_h)
#flatten_h_dense = Flatten()(theta_h_dense)
#reshape_layer = Reshape((batch_size, theta_dim))(theta_h_dense)
enc = Model(inputs=x, outputs=z)
enc.summary()
dec = Model(inputs=theta_h, outputs=theta_h_dense)
dec.summary()

concat = Concatenate(axis=1)
merge_layer = concat([z, theta_h_dense])

############################
#	DECODER - 1
############################
decoder_h = Dense(intermediate_dim, activation='relu')
h_decoded_1 = decoder_h(z)

mer = Model(inputs=x, outputs=h_decoded_1)
mer.summary()

decoder_mean = Conv2DTranspose((90,300,3), (5,5), strides=(2,2), padding='same', activation='sigmoid')
x_decoded_mean_1 = decoder_mean(h_decoded_1)

############################
#	DECODER - 2
############################
decoder_h = Dense(intermediate_dim, activation='relu')
decoder_mean = Conv2DTranspose(original_dim, (5,5), strides=(2,2), padding='same', activation='sigmoid')
h_decoded_2 = decoder_h(merge_layer)
x_decoded_mean_2 = decoder_mean(h_decoded_2)

############################
#	LOSSES
############################
def vae_loss1(x, x_decoded_mean):
    xent_loss = original_dim * metrics.binary_crossentropy(x, x_decoded_mean)
    kl_loss = - 0.5 * K.sum(1 + z_log_var - K.square(z_mean) - K.exp(z_log_var), axis=-1)
    return K.mean(xent_loss + kl_loss)

def vae_loss2(y_train, x_decoded_mean):
    xent_loss = original_dim * metrics.binary_crossentropy(y_train, x_decoded_mean)
    kl_loss = - 0.5 * K.sum(1 + z_log_var - K.square(z_mean) - K.exp(z_log_var), axis=-1)
    return K.mean(xent_loss + kl_loss)

def vae_loss3(x, x_decoded_mean):
    xent_loss1 = original_dim * metrics.binary_crossentropy(x, x_decoded_mean_1)
    xent_loss2 = original_dim * metrics.binary_crossentropy(y_train, x_decoded_mean_2)
    kl_loss = - 0.5 * K.sum(1 + z_log_var - K.square(z_mean) - K.exp(z_log_var), axis=-1)
    return K.mean(xent_loss1 + xent_loss2 + kl_loss)

def vae_loss4(x, x_decoded_mean):
	# np.sqrt(np.mean((a - b) ** 2))
	loss1 = np.sqrt(K.mean((x - x_decoded_mean_1)**2))
	loss2 = np.sqrt(K.mean((y_train - x_decoded_mean_2)**2))
	kl_loss = - 0.5 * K.sum(1 + z_log_var - K.square(z_mean) - K.exp(z_log_var), axis=-1)
	return K.mean(loss1 + loss2 + kl_loss)

# loss = vae_loss(x, x_decoded_mean)
# vae = Model(x, x_decoded_mean)
adam = optimizers.Adam(lr=0.0005, beta_1=0.9, beta_2=0.999, epsilon=1e-08, decay=0.0)
vae = Model(inputs=[x, theta_h], outputs=[x_decoded_mean_1, x_decoded_mean_2])
vae.compile(optimizer=adam, loss=vae_loss4)
print "VAE: ",vae.summary()
#vae = Model(y_train, x_decoded_mean)
#vae.compile(optimizer='rmsprop', loss=vae_loss)

print "=======1=====",theta.shape

vae.fit([x_train, theta], [x_train,y_train],
        #shuffle=True,
        epochs=nb_epochs,
        batch_size=batch_size,)
        # validation_split=0.15)
        #validation_data=([x_train, theta], [x_test,y_test]))

print "======2======",theta.shape

encoded_data_1, encoded_data_2 = vae.predict([x_test, theta])

print encoded_data_1.shape
print encoded_data_2.shape
k = 0 
for num in range(x_test.shape[0]):
	_left = np.reshape(encoded_data_1[num], (90,300,3))
	_right = np.reshape(encoded_data_2[num], (90,300,3))
	left = np.reshape(x_train[num], (90, 300, 3))
	right = np.reshape(y_train[num], (90, 300, 3))
	plt.figure()
	# plt.imshow(left)
	# plt.savefig('left.png')
	# plt.imshow(right)
	# plt.savefig('right.png')
	plt.imshow(_left)
	# os.chdir(generated_left_image)
	plt.savefig("_left_{0}.png".format(num), bbox_inches='tight')
	# os.chdir(generated_right_image)
	plt.imshow(_right)
	plt.savefig("_right_{0}.png".format(num), bbox_inches='tight')
	# os.chdir(cwd)
	# plt.clr()
	plt.close()
	if k == 5:
		exit(0)
	k += 1
	# plt.imshow(_right)
	# plt.savefig(generated_right_image+"_right"+str(num)+".png")

'''exit(0)


_left = np.reshape(encoded_data_1[2], (90,300,3))
_right = np.reshape(encoded_data_2[2], (90,300,3))
left = np.reshape(x_train[1], (90, 300, 3))
right = np.reshape(y_train[1], (90, 300, 3))
plt.figure()
plt.imshow(left)
plt.savefig('left.png')
plt.imshow(right)
plt.savefig('right.png')
plt.imshow(_left)
plt.savefig('_left.png')
plt.imshow(_right)
plt.savefig('_right.png')


ax = plt.subplot(2,2,1)
plt.imshow(left)
ax = plt.subplot(2,2,3)
plt.imshow(right)
ax = plt.subplot(2,2,2)
plt.imshow(_left)
ax = plt.subplot(2,2,4)
plt.imshow(_right)
plt.show()


exit(0)


# build a model to project inputs on the latent space
encoder = Model(x, z_mean)

print "------------1-------------"
# build a digit generator that can sample from the learned distribution
decoder_input = Input(batch_shape=(batch_size, latent_dim))
decoder_h = Dense(intermediate_dim, activation='relu')(decoder_input)
_h_decoded = Dense(original_dim, activation='sigmoid')(decoder_h)
generator = Model(decoder_input, _h_decoded)
#generator.compile(loss='binary_crossentropy', optimizer='rmsprop')
'''