import numpy as np
import cv2
import os
import subprocess
import glob
from options.test_options import TestOptions
from model.net import InpaintingModel_GMCNN
from util.utils import generate_rect_mask, generate_stroke_mask, getLatest

os.environ['CUDA_VISIBLE_DEVICES']=str(np.argmax([int(x.split()[2]) for x in subprocess.Popen(
        "nvidia-smi -q -d Memory | grep -A4 GPU | grep Free", shell=True, stdout=subprocess.PIPE).stdout.readlines()]
        ))

config = TestOptions().parse()

if os.path.isfile(config.dataset_path):
    pathfile = open(config.dataset_path, 'rt').read().splitlines()
elif os.path.isdir(config.dataset_path):
    pathfile = glob.glob(os.path.join(config.dataset_path, '*.png'))
else:
    print('Invalid testing data file/folder path.')
    exit(1)
total_number = len(pathfile)
test_num = total_number if config.test_num == -1 else min(total_number, config.test_num)
print('The total number of testing images is {}, and we take {} for test.'.format(total_number, test_num))

print('configuring model..')
ourModel = InpaintingModel_GMCNN(in_channels=4, opt=config)
ourModel.print_networks()
if config.load_model_dir != '':
    print('Loading pretrained model from {}'.format(config.load_model_dir))
    ourModel.load_networks(getLatest(os.path.join(config.load_model_dir, '*.pth')))
    print('Loading done.')

if config.random_mask:
    np.random.seed(config.seed)

for i in range(test_num):
    if config.mask_type == 'rect':
        mask, _ = generate_rect_mask(config.img_shapes, config.mask_shapes, config.random_mask)
    else:
        mask = generate_stroke_mask(im_size=(config.img_shapes[0], config.img_shapes[1]),
                                    parts=8, maxBrushWidth=20, maxLength=100, maxVertex=20)
    image = cv2.imread(pathfile[i])
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    h, w = image.shape[:2]

    if h >= config.img_shapes[0] and w >= config.img_shapes[1]:
        h_start = (h-config.img_shapes[0]) // 2
        w_start = (w-config.img_shapes[1]) // 2
        image = image[h_start: h_start+config.img_shapes[0], w_start: w_start+config.img_shapes[1], :]
    else:
        t = min(h, w)
        image = image[(h-t)//2:(h-t)//2+t, (w-t)//2:(w-t)//2+t, :]
        image = cv2.resize(image, (config.img_shapes[1], config.img_shapes[0]))

    image = np.transpose(image, [2, 0, 1])
    image = np.expand_dims(image, axis=0)
    image_vis = image * (1-mask) + 255 * mask
    image_vis = np.transpose(image_vis[0][::-1,:,:], [1, 2, 0])
    cv2.imwrite(os.path.join(config.saving_path, 'input_{:03d}.png'.format(i)), image_vis.astype(np.uint8))

    h, w = image.shape[2:]
    grid = 4
    image = image[:, :, :h // grid * grid, :w // grid * grid]
    mask = mask[:, :, :h // grid * grid, :w // grid * grid]
    result = ourModel.evaluate(image, mask)
    result = np.transpose(result[0][::-1,:,:], [1, 2, 0])
    cv2.imwrite(os.path.join(config.saving_path, '{:03d}.png'.format(i)), result)
    print(' > {} / {}'.format(i+1, test_num))
print('done.')