#!/usr/bin/env python2.7

import os
from subprocess import call
import sys
import getopt
import numpy
import torch
import PIL
import PIL.Image
from datetime import datetime
from network import Network, process
from moviepy.video.io.ffmpeg_reader import FFMPEG_VideoReader
from moviepy.video.io.ffmpeg_writer import FFMPEG_VideoWriter

##########################################################

arguments_strModel = 'lf'
arguments_strFirst = './images/first.png'
arguments_strSecond = './images/second.png'
arguments_strOut = './result.png'
arguments_strVideo = False
arguments_strVideoOut = datetime.now().strftime("sepconv_%Y%m%d_%H%M.mp4")
arguments_steps = 0
arguments_dilate = 1
arguments_smooth = False
arguments_aOffset = 0
arguments_bOffset = 0
arguments_mixVideos = False
arguments_averageVideos = False
arguments_mixImages = False
arguments_endpoint = None
arguments_dataset = "dataset"

for strOption, strArgument in getopt.getopt(sys.argv[1:], '', [ strParameter[2:] + '=' for strParameter in sys.argv[1::2] ])[0]:
	print("{}: {}".format(strOption, strArgument))
	if strOption == '--model':
		arguments_strModel = strArgument # which model to use, l1 or lf, please see our paper for more details
	elif strOption == '--first':
		arguments_strFirst = strArgument # path to the first frame
	elif strOption == '--second':
		arguments_strSecond = strArgument # path to the second frame
	elif strOption == '--out':
		arguments_strOut = strArgument # path to where the output should be stored
	elif strOption == '--video':
		arguments_strVideoOut = strArgument # path to video
	elif strOption == '--video-out':
		arguments_strVideoOut = strArgument # path to where the video should be stored
	elif strOption == '--steps':
		arguments_steps = int(strArgument)
	elif strOption == '--dilate':
		arguments_dilate = int(strArgument)
	elif strOption == '--smooth':
		arguments_smooth = bool(strArgument)
	elif strOption == '--mix-videos':
		arguments_mixVideos = bool(strArgument)
	elif strOption == '--mix-images':
		arguments_mixImages = bool(strArgument)
	elif strOption == '--average-videos':
		arguments_averageVideos = bool(strArgument)
	elif strOption == '--a-offset':
		arguments_aOffset = int(strArgument)
	elif strOption == '--b-offset':
		arguments_bOffset = int(strArgument)
	elif strOption == '--endpoint':
		arguments_endpoint = strArgument
	elif strOption == '--dataset':
		arguments_dataset = strArgument

if not os.path.exists('./renders'):
	os.mkdir('renders')

moduleNetwork = Network(arguments_strModel).cuda()
tensorOutput = torch.FloatTensor()

index = 0

def recurse_two_frames(moduleNetwork, tensorOutput, a_np, b_np, frame_index, morph_index, step, depth=0):
	print("generate {} {} {}".format(frame_index, morph_index, step))

	tensorInputFirst = torch.FloatTensor(a_np)
	tensorInputSecond = torch.FloatTensor(b_np)
	process(moduleNetwork, tensorInputFirst, tensorInputSecond, tensorOutput)
	middle_np = tensorOutput.clamp(0.0, 1.0).numpy()

	if morph_index == frame_index:
		print("frame {}, depth {}".format(frame_index, depth))
		return middle_np
	if morph_index > frame_index:
		next_index = morph_index - step
		next_a_np = a_np
		next_b_np = middle_np
		# print("next index: {} - {}".format(next_index, step))
	else:
		next_index = morph_index + step
		next_a_np = middle_np
		next_b_np = b_np
		# print("next index: {} + {}".format(next_index, step))
	return recurse_two_frames(moduleNetwork, tensorOutput, next_a_np, next_b_np, frame_index, next_index, step/2, depth+1)

def recurse_videos(moduleNetwork, tensorOutput, a, b, a_offset, b_offset, count, step, frame_index, dilate):
	global index
	index += 1
	if (index % 10) == 0:
		print("{}...".format(index))
	step /= 2
	a_fn = os.path.join(a, "frame_{:0>5}.png".format(int(frame_index + a_offset)))
	b_fn = os.path.join(b, "frame_{:0>5}.png".format(int(frame_index + b_offset)))
	print("{} => {}".format(a_fn, b_fn))
	a_np = load_image(a_fn)
	b_np = load_image(b_fn)
	img_np = recurse_two_frames(moduleNetwork, tensorOutput, a_np, b_np, frame_index, count / 2, count / 4)
	if step < 2 * dilate:
		return [img_np]
	else:
		left = recurse_videos(moduleNetwork, tensorOutput, a, b, a_offset, b_offset, count, step, frame_index - (step/2), dilate)
		right = recurse_videos(moduleNetwork, tensorOutput, a, b, a_offset, b_offset, count, step, frame_index + (step/2), dilate)
		return left + [img_np] + right

def process_two_videos(moduleNetwork, tensorOutput, a, b, a_offset, b_offset, steps, dilate):
	steps *= dilate
	return recurse_videos(moduleNetwork, tensorOutput, a, b, a_offset, b_offset, steps, steps, steps/2, dilate)

def average_two_videos(moduleNetwork, tensorOutput, a, b, a_offset, b_offset, steps, dilate):
	global index
	index += 1
	if (index % 10) == 0:
		print("{}...".format(index))
	frames = []
	steps *= dilate
	for i in range(1, steps * dilate + 1, dilate):
		a_fn = os.path.join(a, "frame_{:0>5}.png".format(int(i + a_offset)))
		b_fn = os.path.join(b, "frame_{:0>5}.png".format(int(i + b_offset)))
		print("{} => {}".format(a_fn, b_fn))
		a_np = load_image(a_fn)
		b_np = load_image(b_fn)
		tensorInputFirst = torch.FloatTensor(a_np)
		tensorInputSecond = torch.FloatTensor(b_np)
		process(moduleNetwork, tensorInputFirst, tensorInputSecond, tensorOutput)
		middle_np = tensorOutput.clamp(0.0, 1.0).numpy()
		frames.append(middle_np)
	return frames

def process_tree(moduleNetwork, a, b, tensorOutput, steps, dilate):
	global index
	index += 1
	if (index % 10) == 0:
		print("{}...".format(index))
	tensorInputFirst = torch.FloatTensor(a)
	tensorInputSecond = torch.FloatTensor(b)
	process(moduleNetwork, tensorInputFirst, tensorInputSecond, tensorOutput)
	middle_np = tensorOutput.clamp(0.0, 1.0).numpy()
	if steps < 4 * dilate:
		return [middle_np]
	else:
		left = process_tree(moduleNetwork, a, middle_np, tensorOutput, steps / 2, dilate)
		right = process_tree(moduleNetwork, middle_np, b, tensorOutput, steps / 2, dilate)
		return left + [middle_np] + right

def smooth_frames(moduleNetwork, tensorOutput, frames, smooth):
	if not smooth:
		return frames
	print("smoothing every other frame")
	firstFrame = frames[0]
	nextFrame = None
	new_frames = [firstFrame]
	for i in range(1, len(frames)-2, 2):
		firstFrame = frames[i]
		nextFrame = frames[i+2]

		tensorInputFirst = torch.FloatTensor(firstFrame)
		tensorInputSecond = torch.FloatTensor(nextFrame)
		process(moduleNetwork, tensorInputFirst, tensorInputSecond, tensorOutput)
		middle_np = tensorOutput.clamp(0.0, 1.0).numpy()

		new_frames += [firstFrame]
		new_frames += [middle_np]
	if nextFrame is not None:
		new_frames += [nextFrame]
	new_frames += [frames[len(frames)-1]]
	return new_frames

def dilate_frames(moduleNetwork, tensorOutput, frames, dilate):
	if dilate < 2:
		return frames
	print("dilating by a factor of {}".format(dilate))
	new_frames = []
	nextFrame = frames[0]
	for i in range(1, len(frames)):
		firstFrame = nextFrame
		nextFrame = frames[i]
		new_frames += [firstFrame]
		new_frames += process_tree(moduleNetwork, firstFrame, nextFrame, tensorOutput, dilate, 1)
	new_frames += [nextFrame]
	return new_frames

def store_frames(frames, outputPath, inputFirst=None, inputSecond=None, endpoint=None, dataset="unknown"):
	print('writing {}'.format(outputPath))
	print('frames: {}'.format(len(frames)))
	writer = FFMPEG_VideoWriter(outputPath, (1024, 512), 25)
	if inputFirst is not None:
		writer.write_frame(inputFirst)
	for frame in frames:
		writer.write_frame(tensor_to_image(frame))
	if inputSecond is not None:
		writer.write_frame(inputSecond)
	writer.close()

	if endpoint is not None:
	  call(["curl",
	    "-X", "POST",
	    "-F", "module=morph",
	    "-F", "activity=morph",
	    "-F", "generated=true",
	    "-F", "dataset=" + dataset,
	    "-F", "datatype=video",
	    "-F", "should_relay=true",
	    "-F", "file=@" + outputPath,
	    endpoint
	  ])

def tensor_to_image(np_val):
	return (numpy.rollaxis(np_val, 0, 3)[:,:,::-1] * 255.0).astype(numpy.uint8)

def load_image(path):
	return numpy.rollaxis(numpy.asarray(PIL.Image.open(path))[:,:,::-1], 2, 0).astype(numpy.float32) / 255.0

def load_image_tensor(path):
	return torch.FloatTensor(load_image(path))

if arguments_strVideo and arguments_strVideoOut:
	reader = FFMPEG_VideoReader(arguments_strVideo, False)
	writer = FFMPEG_VideoWriter(arguments_strVideoOut, reader.size, reader.fps*2)
	reader.initialize()
	nextFrame = reader.read_frame()
	for x in range(0, reader.nframes):
		firstFrame = nextFrame
		nextFrame = reader.read_frame()
		tensorInputFirst = torch.FloatTensor(numpy.rollaxis(firstFrame[:,:,::-1], 2, 0) / 255.0)
		tensorInputSecond = torch.FloatTensor(numpy.rollaxis(nextFrame[:,:,::-1], 2, 0) / 255.0)
		process(moduleNetwork, tensorInputFirst, tensorInputSecond, tensorOutput)
		writer.write_frame(firstFrame)
		writer.write_frame((numpy.rollaxis(tensorOutput.clamp(0.0, 1.0).numpy(), 0, 3)[:,:,::-1] * 255.0).astype(numpy.uint8))
	writer.write_frame(nextFrame)
	writer.close()
elif arguments_mixVideos:
	# Morph two videos
	print("morph two videos...")
	print("{} => {}".format(arguments_strFirst, arguments_strSecond))
	outputPath = './renders/' + arguments_strVideoOut
	frames = process_two_videos(moduleNetwork, tensorOutput, arguments_strFirst, arguments_strSecond, arguments_aOffset, arguments_bOffset, arguments_steps, arguments_dilate)
	frames = smooth_frames(moduleNetwork, tensorOutput, frames, arguments_smooth)
	frames = dilate_frames(moduleNetwork, tensorOutput, frames, arguments_dilate)
	store_frames(frames, outputPath, endpoint=arguments_endpoint, dataset=arguments_dataset)
elif arguments_averageVideos:
	print("average two videos...")
	print("{} => {}".format(arguments_strFirst, arguments_strSecond))
	outputPath = './renders/' + arguments_strVideoOut
	frames = average_two_videos(moduleNetwork, tensorOutput, arguments_strFirst, arguments_strSecond, arguments_aOffset, arguments_bOffset, arguments_steps, arguments_dilate)
	frames = smooth_frames(moduleNetwork, tensorOutput, frames, arguments_smooth)
	frames = dilate_frames(moduleNetwork, tensorOutput, frames, arguments_dilate)
	store_frames(frames, outputPath, endpoint=arguments_endpoint, dataset=arguments_dataset)
elif arguments_steps == 0:
	# Process image
	tensorInputFirst = load_image_tensor(arguments_strFirst)
	tensorInputSecond = load_image_tensor(arguments_strSecond)
	process(moduleNetwork, tensorInputFirst, tensorInputSecond, tensorOutput)
	PIL.Image.fromarray((numpy.rollaxis(tensorOutput.clamp(0.0, 1.0).numpy(), 0, 3)[:,:,::-1] * 255.0).astype(numpy.uint8)).save(arguments_strOut)
elif arguments_mixImages:
	print("{} => {}".format(arguments_strFirst, arguments_strSecond))
	inputFirst = load_image(os.path.join(arguments_strFirst, "frame_{:0>5}.png".format(int(arguments_aOffset)+1)))
	inputSecond = load_image(os.path.join(arguments_strSecond, "frame_{:0>5}.png".format(int(arguments_bOffset)+1)))
	outputPath = './renders/' + arguments_strVideoOut
	frames = process_tree(moduleNetwork, inputFirst, inputSecond, tensorOutput, arguments_steps * arguments_dilate, arguments_dilate)
	frames = smooth_frames(moduleNetwork, tensorOutput, frames, arguments_smooth)
	print("dilate... {}".format(arguments_dilate))
	frames = dilate_frames(moduleNetwork, tensorOutput, frames, arguments_dilate)
	store_frames(frames, outputPath, inputFirst, inputSecond, endpoint=arguments_endpoint, dataset=arguments_dataset)

else:
	# Morph two images
	print("{} => {}".format(arguments_strFirst, arguments_strSecond))
	inputFirst = load_image(arguments_strFirst)
	inputSecond = load_image(arguments_strSecond)
	outputPath = './renders/' + arguments_strVideoOut
	frames = process_tree(moduleNetwork, inputFirst, inputSecond, tensorOutput, arguments_steps * arguments_dilate, arguments_dilate)
	frames = smooth_frames(moduleNetwork, tensorOutput, frames, arguments_smooth)
	print("dilate... {}".format(arguments_dilate))
	frames = dilate_frames(moduleNetwork, tensorOutput, frames, arguments_dilate)
	store_frames(frames, outputPath, inputFirst, inputSecond, endpoint=arguments_endpoint, dataset=arguments_dataset)