import argparse
import functools
import os
from PIL import Image
from paddle.fluid import core
import paddle.fluid as fluid
import paddle
import numpy as np
from scipy.misc import imsave
from model import *
import glob
from utility import add_arguments, print_arguments

parser = argparse.ArgumentParser(description=__doc__)
add_arg = functools.partial(add_arguments, argparser=parser)
# yapf: disable
add_arg('input',             str,   None, "The images to be infered.")
add_arg('output',            str,   "./infer_result", "The directory the infer result to be saved to.")
add_arg('init_model',        str,   None,       "The init model file of directory.")
add_arg('input_style',        str,  "A",       "The style of the input, A or B")
add_arg('use_gpu',           bool,  True,       "Whether to use GPU to train.")
# yapf: enable


def infer(args):
    data_shape = [-1, 3, 256, 256]
    input = fluid.layers.data(name='input', shape=data_shape, dtype='float32')
    if args.input_style == "A":
        model_name = 'g_a'
        fake = build_generator_resnet_9blocks(input, name="g_A")
    elif args.input_style == "B":
        model_name = 'g_b'
        fake = build_generator_resnet_9blocks(input, name="g_B")
    else:
        raise "Input with style [%s] is not supported." % args.input_style
    # prepare environment
    place = fluid.CPUPlace()
    if args.use_gpu:
        place = fluid.CUDAPlace(0)
    exe = fluid.Executor(place)
    exe.run(fluid.default_startup_program())
    fluid.io.load_persistables(exe, args.init_model + "/" + model_name)

    if not os.path.exists(args.output):
        os.makedirs(args.output)
    for file in glob.glob(args.input):
        image_name = os.path.basename(file)
        image = Image.open(file)
        image = image.resize((256, 256))
        image = np.array(image) / 127.5 - 1
        if len(image.shape) != 3:
            continue
        data = image.transpose([2, 0, 1])[np.newaxis, :].astype("float32")
        tensor = core.LoDTensor()
        tensor.set(data, place)

        fake_temp = exe.run(fetch_list=[fake.name], feed={"input": tensor})
        fake_temp = np.squeeze(fake_temp[0]).transpose([1, 2, 0])
        input_temp = np.squeeze(data).transpose([1, 2, 0])

        imsave(args.output + "/fake_" + image_name, (
            (fake_temp + 1) * 127.5).astype(np.uint8))


if __name__ == "__main__":
    args = parser.parse_args()
    print_arguments(args)
    infer(args)