#! /usr/bin/env python
import os
import cv2
import argparse
import numpy as np

from face_detection import select_face, mask_mouth
from face_swap import face_swap


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='FaceSwapApp')
    parser.add_argument('--src', required=True, help='Path for source image')
    parser.add_argument('--dst', required=True, help='Path for target image')
    parser.add_argument('--out', required=True, help='Path for storing output images')
    parser.add_argument('--point_plot', default=False, help='facial landmark visualizations')
    parser.add_argument('--warp_2d', default=False, action='store_true', help='2d or 3d warp')
    parser.add_argument('--correct_color', default=False, action='store_true', help='Correct color')
    parser.add_argument('--no_debug_window', default=False, action='store_true', help='Don\'t show debug window')
    args = parser.parse_args()

    # Read images
    src_img = cv2.imread(args.src)
    src_img_mask = mask_mouth(args.src)
    dst_img = cv2.imread(args.dst)

    # Select src face
    src_points, src_shape, src_face = select_face(src_img_mask)
    # Select dst face
    dst_points, dst_shape, dst_face = select_face(dst_img)

    if src_points is None or dst_points is None:
        print('Detect 0 Face !!!')
        exit(-1)

    output = face_swap(src_face, dst_face, src_points, dst_points, dst_shape, dst_img, args)

    dir_path = os.path.dirname(args.out)
    if not os.path.isdir(dir_path):
        os.makedirs(dir_path)

    cv2.imwrite(args.out, output)

    ##For debug
    if not args.no_debug_window:
        if args.point_plot:
            tmp_img=src_img.copy()
            x, y, w, h = src_shape

            for point in src_points:
                x_, y_ = point.tolist()
                cv2.circle(tmp_img, [x_+x,y_+y], 6,  (0,0,255), -1)
                cv2.circle(tmp_img, [x_+x,y_+y], 3,  (255,255,255), -1)

            tmp_img2=dst_img.copy()
            x, y, w, h = dst_shape
            for point in dst_points:
                x_, y_ = point.tolist()
                cv2.circle(tmp_img2, [x_+x,y_+y], 6,  (0,0,255), -1)
                cv2.circle(tmp_img2, [x_+x,y_+y], 3,  (255,255,255), -1)

            result=np.hstack((tmp_img, tmp_img2, output))
        else:
            result=np.hstack((src_img, dst_img, output))
        cv2.namedWindow('result', 0)
        cv2.imshow("result", result)
        cv2.waitKey(0)
        
        cv2.destroyAllWindows()