# standard library imports
import os

# third party imports
import numpy as np
from PIL import Image

import torch.nn as nn
from torchvision import transforms


# local imports
import config
from . import utils
from . import geometric_transformer



class GeoTransformationInfer(nn.Module):
	def __init__(self, output_dir="./output/results"):
		super(GeoTransformationInfer, self).__init__()
		self.output_dir = output_dir
		utils.ensure_folder(self.output_dir)

	def forward(self, model_apparel, warped_image, model_image, warped_model_image, random_product_image, random_product_image_warped, output_on_random_product, batch_index, epoch):
		batch_size = warped_image.shape[0]

		model_apparel = model_apparel.cpu().numpy()
		warped_image = warped_image.cpu().numpy()
		model_image = model_image.cpu().numpy()
		warped_model_image = warped_model_image.cpu().numpy()
		random_product_image = random_product_image.cpu().numpy()
		random_product_image_warped = random_product_image_warped.cpu().numpy()
		output_on_random_product = output_on_random_product.cpu().numpy()

		for i in range(batch_size):
			self._save_image_sheet(
				batch_index*config.PARAMS["batch_size"] + i, 
				model_apparel[i], 
				warped_image[i], 
				model_image[i], 
				warped_model_image[i], 
				random_product_image[i],
				random_product_image_warped[i], 
				output_on_random_product[i], 
				epoch)

	def _save_image_sheet(self, 
		idx, 
		model_apparel, 
		warped_image, 
		model_image, 
		warped_model_image, 
		random_product_image,
		random_product_image_warped, 
		output_on_random_product, 
		epoch):

		# inverse normalization of the images along with channel first to channel last steps and finally converting np array to pillow format for saving
		model_apparel = np.moveaxis(model_apparel, 0, 2) * [0.229, 0.224, 0.225] + [0.485, 0.456, 0.406]
		model_apparel = Image.fromarray(np.uint8(model_apparel * 255))
		warped_image = np.moveaxis(warped_image, 0, 2) * [0.229, 0.224, 0.225] + [0.485, 0.456, 0.406]
		warped_image = Image.fromarray(np.uint8(warped_image * 255))
		model_image = np.moveaxis(model_image, 0, 2) * [0.229, 0.224, 0.225] + [0.485, 0.456, 0.406]
		model_image = Image.fromarray(np.uint8(model_image * 255))
		warped_model_image = np.moveaxis(warped_model_image, 0, 2) * [0.229, 0.224, 0.225] + [0.485, 0.456, 0.406]
		warped_model_image = Image.fromarray(np.uint8(warped_model_image * 255))
		random_product_image = np.moveaxis(random_product_image, 0, 2) * [0.229, 0.224, 0.225] + [0.485, 0.456, 0.406]
		random_product_image = Image.fromarray(np.uint8(random_product_image * 255))
		random_product_image_warped = np.moveaxis(random_product_image_warped, 0, 2) * [0.229, 0.224, 0.225] + (0.485, 0.456, 0.406)
		random_product_image_warped = Image.fromarray(np.uint8(random_product_image_warped * 255))
		output_on_random_product = np.moveaxis(output_on_random_product, 0, 2) * [0.229, 0.224, 0.225] + [0.485, 0.456, 0.406]
		output_on_random_product = Image.fromarray(np.uint8(output_on_random_product * 255))

		sheet = Image.new('RGB', (1568, 224), 'white')
		sheet.paste(model_apparel, (0, 0))
		sheet.paste(warped_image, (224, 0))
		sheet.paste(model_image, (448, 0))
		sheet.paste(warped_model_image, (672, 0))
		sheet.paste(random_product_image, (896, 0))
		sheet.paste(random_product_image_warped, (1120, 0))
		sheet.paste(output_on_random_product, (1344, 0))
		sheet.save(os.path.join(self.output_dir, "image_sheet_{}-epoch{}".format(idx, str(epoch).zfill(3)) + ".jpg"))


[((75, 10, 75, 48), 'PIL.Image.new', 'Image.new', ({(75, 20, 75, 25): '"""RGB"""', (75, 27, 75, 38): '(1568, 224)', (75, 40, 75, 47): '"""white"""'}, {}), "('RGB', (1568, 224), 'white')", False, 'from PIL import Image\n'), ((61, 34, 61, 63), 'numpy.uint8', 'np.uint8', ({(61, 43, 61, 62): 'model_apparel * 255'}, {}), '(model_apparel * 255)', True, 'import numpy as np\n'), ((63, 33, 63, 61), 'numpy.uint8', 'np.uint8', ({(63, 42, 63, 60): 'warped_image * 255'}, {}), '(warped_image * 255)', True, 'import numpy as np\n'), ((65, 32, 65, 59), 'numpy.uint8', 'np.uint8', ({(65, 41, 65, 58): 'model_image * 255'}, {}), '(model_image * 255)', True, 'import numpy as np\n'), ((67, 39, 67, 73), 'numpy.uint8', 'np.uint8', ({(67, 48, 67, 72): 'warped_model_image * 255'}, {}), '(warped_model_image * 255)', True, 'import numpy as np\n'), ((69, 41, 69, 77), 'numpy.uint8', 'np.uint8', ({(69, 50, 69, 76): 'random_product_image * 255'}, {}), '(random_product_image * 255)', True, 'import numpy as np\n'), ((71, 48, 71, 91), 'numpy.uint8', 'np.uint8', ({(71, 57, 71, 90): 'random_product_image_warped * 255'}, {}), '(random_product_image_warped * 255)', True, 'import numpy as np\n'), ((73, 45, 73, 85), 'numpy.uint8', 'np.uint8', ({(73, 54, 73, 84): 'output_on_random_product * 255'}, {}), '(output_on_random_product * 255)', True, 'import numpy as np\n'), ((60, 18, 60, 50), 'numpy.moveaxis', 'np.moveaxis', ({(60, 30, 60, 43): 'model_apparel', (60, 45, 60, 46): '(0)', (60, 48, 60, 49): '(2)'}, {}), '(model_apparel, 0, 2)', True, 'import numpy as np\n'), ((62, 17, 62, 48), 'numpy.moveaxis', 'np.moveaxis', ({(62, 29, 62, 41): 'warped_image', (62, 43, 62, 44): '(0)', (62, 46, 62, 47): '(2)'}, {}), '(warped_image, 0, 2)', True, 'import numpy as np\n'), ((64, 16, 64, 46), 'numpy.moveaxis', 'np.moveaxis', ({(64, 28, 64, 39): 'model_image', (64, 41, 64, 42): '(0)', (64, 44, 64, 45): '(2)'}, {}), '(model_image, 0, 2)', True, 'import numpy as np\n'), ((66, 23, 66, 60), 'numpy.moveaxis', 'np.moveaxis', ({(66, 35, 66, 53): 'warped_model_image', (66, 55, 66, 56): '(0)', (66, 58, 66, 59): '(2)'}, {}), '(warped_model_image, 0, 2)', True, 'import numpy as np\n'), ((68, 25, 68, 64), 'numpy.moveaxis', 'np.moveaxis', ({(68, 37, 68, 57): 'random_product_image', (68, 59, 68, 60): '(0)', (68, 62, 68, 63): '(2)'}, {}), '(random_product_image, 0, 2)', True, 'import numpy as np\n'), ((70, 32, 70, 78), 'numpy.moveaxis', 'np.moveaxis', ({(70, 44, 70, 71): 'random_product_image_warped', (70, 73, 70, 74): '(0)', (70, 76, 70, 77): '(2)'}, {}), '(random_product_image_warped, 0, 2)', True, 'import numpy as np\n'), ((72, 29, 72, 72), 'numpy.moveaxis', 'np.moveaxis', ({(72, 41, 72, 65): 'output_on_random_product', (72, 67, 72, 68): '(0)', (72, 70, 72, 71): '(2)'}, {}), '(output_on_random_product, 0, 2)', True, 'import numpy as np\n')]