# Learning the basic structure of an autoencoder.
import tensorflow as tf
print("Num GPUs Available: ", len(tf.config.list_physical_devices('GPU')))


from tensorflow.keras.layers import Input, Dense, Conv2D, MaxPooling2D, UpSampling2D
from tensorflow.keras.models import Sequential, Model
from keras.utils.image_utils import img_to_array, load_img
from matplotlib.pyplot import imshow
from PIL import Image, ImageDraw
import numpy as np
import cv2

np.random.seed(342)  # Seed for reproducibility, we want to see the same results every time we run the code,
# reproducibility.

# Image Preprocessing
# Convert to an image numpy array: (float32)
# Batches
# Normalise the image.
profile_img = cv2.imread('profile.PNG')
profile_img = cv2.cvtColor(profile_img, cv2.COLOR_BGR2RGB)  # convert to RGB.
print(profile_img.shape)  # (28, 28, 3)

reshaped = cv2.resize(profile_img, (256, 256))
print(reshaped.shape)  # (256, 256, 3)
img_array = img_to_array(reshaped).astype('float32') / 255.  # Normalise the image.
# Need 1 batch of the image.
img_array = np.expand_dims(img_array, axis=0)

# cv2.imshow('Image', reshaped)
# cv2.waitKey(0)

# Get the image dimensions at batch x height x width x channels
# Roughly 1x28x28x3 dims.
print(img_array.shape)  # (1, 256, 256, 3)

SIZE = 256
model = Sequential()
# Encoder Layer
model.add(Conv2D(32, (3, 3), activation='relu', padding='same', input_shape=(SIZE, SIZE, 3)))
model.add(MaxPooling2D((2, 2), padding='same')) # Shrinking
model.add(Conv2D(64, (3, 3), activation='relu', padding='same'))
model.add(MaxPooling2D((2, 2), padding='same')) # Shrinking
model.add(Conv2D(128, (3, 3), activation='relu', padding='same'))
model.add(MaxPooling2D((2, 2), padding='same')) # Bottleneck
model.add(Conv2D(256, (3, 3), activation='relu', padding='same'))
model.add(MaxPooling2D((2, 2), padding='same')) # Bottleneck
model.add(Conv2D(512, (3, 3), activation='relu', padding='same'))

# Middle Layer
model.add(MaxPooling2D((2, 2), padding='same')) # Shrinking

# Decoder Layer
model.add(Conv2D(512, (3, 3), activation='relu', padding='same'))
model.add(UpSampling2D((2, 2)))  # Expanding
model.add(Conv2D(256, (3, 3), activation='relu', padding='same'))
model.add(UpSampling2D((2, 2)))  # Expanding
model.add(Conv2D(128, (3, 3), activation='relu', padding='same'))
model.add(UpSampling2D((2, 2)))  # Expanding
model.add(Conv2D(64, (3, 3), activation='relu', padding='same'))
model.add(UpSampling2D((2, 2)))  # Expanding
model.add(Conv2D(32, (3, 3), activation='relu', padding='same'))
model.add(UpSampling2D((2, 2)))  # Expanding

# Output Layer
model.add(Conv2D(3, (3, 3), activation='relu', padding='same'))

model.compile(optimizer='adam', loss='mean_squared_error', metrics=['accuracy'])




if __name__ == '__main__':
	model.summary()

	# Image Gif Generator based on keras callbacks.
	model_outputs = []


	def on_epoch_end(epoch, logs):
		model_outputs.append(model.predict(img_array))


	# Callbacks
	predict_callback = tf.keras.callbacks.LambdaCallback(on_epoch_end=on_epoch_end)

	# I realised Keras is really easy to use, and I am loving it.
	# Training the model
	model.fit(x=img_array, y=img_array, epochs=1000, batch_size=20, shuffle=True, callbacks=[predict_callback])


	predict = model.predict(x=img_array)  # We hope to get similar to the original image.
	cv2.imshow("output", predict[0].reshape(SIZE, SIZE, 3))
	cv2.waitKey(0)

	print("Creating a gif/video from the model outputs.")
	# Make a gif creator from an array of given images.
	print(len(model_outputs))  # 50

	images = [(output[0] * 255).astype('uint8') for output in model_outputs]
	try:
		images = [(output[0] * 255).astype('uint8') for output in model_outputs]

		# Specify the video file name and codec
		video_filename = 'model_outputs_5.4.mp4'
		fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # Use appropriate codec for your system

		# Get the height and width from the first image
		height, width, _ = images[0].shape

		# Create a VideoWriter object
		fps = 30  # Adjust the frames per second (FPS) as needed
		video = cv2.VideoWriter(video_filename, fourcc, fps, (width, height))

		# Write each image to the video
		for image in images:
			# Convert BRG to RGB
			rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
			video.write(rgb_image)

		# Release the VideoWriter
		video.release()
	except Exception as e:
		print(e)
		images = [Image.fromarray((output[0] * 255).astype('uint8')) for output in model_outputs]
		images[0].save('model_outputs_5.4m_param.gif', save_all=True, append_images=images[1:], optimize=False, duration=20,
		               loop=0)
	#

	# # Save the images as a GIF
	#

	images = [Image.fromarray((output[0] * 255).astype('uint8')) for output in model_outputs]
	images[0].save('model_outputs_5.4m_param.gif', save_all=True, append_images=images[1:], optimize=False, duration=20,
		               loop=0)