import os
import cv2
import glob
import numpy as np
from utils import landmark_processor

from face_enhance.face_gan_pt import FaceGAN
from time import time

class FaceEnhancement(object):
    def __init__(self,  size=512, gpu_id=0):
        self.facegan = FaceGAN(size, gpu_id)
        self.size = size
        self.threshold = 0.9

        # the mask for pasting restored faces back
        self.mask = np.zeros((512, 512), np.float32)
        cv2.rectangle(self.mask, (26, 26), (486, 486), (1, 1, 1), -1, cv2.LINE_AA)
        self.mask = cv2.GaussianBlur(self.mask, (101, 101), 11)
        self.mask = cv2.GaussianBlur(self.mask, (101, 101), 11)

        self.kernel = np.array((
                [0.0625, 0.125, 0.0625],
                [0.125, 0.25, 0.125],
                [0.0625, 0.125, 0.0625]), dtype="float32")

    def process(self, img, landmarks1k):

        assert len(landmarks1k) == 1000

        image_to_face_mat = landmark_processor.get_transform_mat_face_restore(landmarks1k, self.size)
        tfm_inv = cv2.invertAffineTransform(image_to_face_mat)

        height, width = img.shape[:2]
        full_mask = np.zeros((height, width), dtype=np.float32)
        full_img = np.zeros(img.shape, dtype=np.uint8)

        fh, fw = (landmarks1k[0][1]-landmarks1k[154][1]), (landmarks1k[95][0]-landmarks1k[215][0])

        of = cv2.warpAffine(img, image_to_face_mat, (self.size, self.size), flags=cv2.INTER_LANCZOS4, borderMode=cv2.BORDER_CONSTANT, borderValue=[0, 0, 0])

        # enhance the face
        ef = self.facegan.process(of)
        tmp_mask = self.mask
        tmp_mask = cv2.resize(tmp_mask, ef.shape[:2])
        tmp_mask = cv2.warpAffine(tmp_mask, tfm_inv, (width, height), flags=3)

        if min(fh, fw)<100: # gaussian filter for small faces
            ef = cv2.filter2D(ef, -1, self.kernel)

        # tmp_img = cv2.warpAffine(ef, tfm_inv, (width, height), flags=3)
        tmp_img = cv2.warpAffine(ef, tfm_inv, (width, height), dst=img.copy(), borderMode=cv2.BORDER_TRANSPARENT)

        # cv2.imshow("tmp_img: ", tmp_img)

        mask = tmp_mask - full_mask
        full_mask[np.where(mask>0)] = tmp_mask[np.where(mask>0)]
        full_img[np.where(mask>0)] = tmp_img[np.where(mask>0)]

        full_mask = full_mask[:, :, np.newaxis]
        img = cv2.convertScaleAbs(img*(1-full_mask) + full_img*full_mask)

        # cv2.imshow("img: ", img)

        # cv2.waitKey()

        return img

if __name__=='__main__':
    
    indir = '/media/DATA_4T/zao_data/tencent_con_0701/ref_c/res_contrast_blur13_notps_yunfu_res'
    outdir = '/media/DATA_4T/zao_data/tencent_con_0701/ref_c/res_contrast_blur13_notps_yunfu_res_outs2'
    os.makedirs(outdir, exist_ok=True)

    faceenhancer = FaceEnhancement(base_dir="./", size=512, model="GPEN-512", channel_multiplier=2)

    files = sorted(glob.glob(os.path.join(indir, '*.*g')))
    for n, file in enumerate(files[:]):
        filename = os.path.basename(file)
        txtname = file.replace(".jpg", "_landmark1k.txt")

        im = cv2.imread(file, cv2.IMREAD_COLOR) # BGR
        print(txtname)
        landmark = np.loadtxt(txtname)
        if not isinstance(im, np.ndarray): print(filename, 'error'); continue

        start = time()

        img = faceenhancer.process(im, landmark)

        end = time()

        print("Time cost: {:.4f}".format(end - start))

        cv2.imwrite(os.path.join(outdir, '.'.join(filename.split('.')[:-1])+'_2.jpg'), img)