import cv2
import time
import imutils
import numpy as np
import tqdm
from imutils import face_utils
import numpy as np
import argparse
import imutils
import dlib
import cv2
from imutils import face_utils
import glob
import os
import sys
os.environ["CUDA_VISIBLE_DEVICES"]=sys.argv[1]

dlib.DLIB_USE_CUDA = True
detector = dlib.get_frontal_face_detector()
predictor = dlib.shape_predictor('shape_predictor_68_face_landmarks.dat')


def overlay_transparent(background, overlay, x, y):

    background_width = background.shape[1]
    background_height = background.shape[0]

    if x >= background_width or y >= background_height:
        return background

    h, w = overlay.shape[0], overlay.shape[1]

    if x + w > background_width:
        w = background_width - x
        overlay = overlay[:, :w]

    if y + h > background_height:
        h = background_height - y
        overlay = overlay[:h]

    if overlay.shape[2] < 4:
        overlay = np.concatenate(
            [
                overlay,
                np.ones((overlay.shape[0], overlay.shape[1], 1), dtype = overlay.dtype) * 255
            ],
            axis = 2,
        )

    overlay_image = overlay[..., :3]
    mask = overlay[..., 3:] / 255.0

    background[y:y+h, x:x+w] = (1.0 - mask) * background[y:y+h, x:x+w] + mask * overlay_image

    return background

def shapeconverter(shape):
    x0s = []
    y0s = []
    x1s = []
    y1s = []

    x0s.append(shape[48][0])
    x0s.append(shape[49][0])
    x0s.append(shape[50][0])
    x0s.append(shape[51][0])
    x0s.append(shape[52][0])
    x0s.append(shape[53][0])
    x0s.append(shape[48][0])
    x0s.append(shape[60][0])
    x0s.append(shape[61][0])
    x0s.append(shape[62][0])
    x0s.append(shape[63][0])
    x0s.append(shape[64][0])
    x0s.append(shape[48][0])
    x0s.append(shape[67][0])
    x0s.append(shape[66][0])
    x0s.append(shape[65][0])
    x0s.append(shape[64][0])
    x0s.append(shape[48][0])
    x0s.append(shape[59][0])
    x0s.append(shape[58][0])
    x0s.append(shape[57][0])
    x0s.append(shape[56][0])
    x0s.append(shape[55][0])

    y0s.append(shape[48][1])
    y0s.append(shape[49][1])
    y0s.append(shape[50][1])
    y0s.append(shape[51][1])
    y0s.append(shape[52][1])
    y0s.append(shape[53][1])
    y0s.append(shape[48][1])
    y0s.append(shape[60][1])
    y0s.append(shape[61][1])
    y0s.append(shape[62][1])
    y0s.append(shape[63][1])
    y0s.append(shape[64][1])
    y0s.append(shape[48][1])
    y0s.append(shape[67][1])
    y0s.append(shape[66][1])
    y0s.append(shape[65][1])
    y0s.append(shape[64][1])
    y0s.append(shape[48][1])
    y0s.append(shape[59][1])
    y0s.append(shape[58][1])
    y0s.append(shape[57][1])
    y0s.append(shape[56][1])
    y0s.append(shape[55][1])

    x1s.append(shape[49][0])
    x1s.append(shape[50][0])
    x1s.append(shape[51][0])
    x1s.append(shape[52][0])
    x1s.append(shape[53][0])
    x1s.append(shape[54][0])
    x1s.append(shape[61][0])
    x1s.append(shape[61][0])
    x1s.append(shape[62][0])
    x1s.append(shape[63][0])
    x1s.append(shape[64][0])
    x1s.append(shape[54][0])
    x1s.append(shape[67][0])
    x1s.append(shape[66][0])
    x1s.append(shape[65][0])
    x1s.append(shape[64][0])
    x1s.append(shape[54][0])
    x1s.append(shape[59][0])
    x1s.append(shape[58][0])
    x1s.append(shape[57][0])
    x1s.append(shape[56][0])
    x1s.append(shape[55][0])
    x1s.append(shape[54][0])

    y1s.append(shape[49][1])
    y1s.append(shape[50][1])
    y1s.append(shape[51][1])
    y1s.append(shape[52][1])
    y1s.append(shape[53][1])
    y1s.append(shape[54][1])
    y1s.append(shape[61][1])
    y1s.append(shape[61][1])
    y1s.append(shape[62][1])
    y1s.append(shape[63][1])
    y1s.append(shape[64][1])
    y1s.append(shape[54][1])
    y1s.append(shape[67][1])
    y1s.append(shape[66][1])
    y1s.append(shape[65][1])
    y1s.append(shape[64][1])
    y1s.append(shape[54][1])
    y1s.append(shape[59][1])
    y1s.append(shape[58][1])
    y1s.append(shape[57][1])
    y1s.append(shape[56][1])
    y1s.append(shape[55][1])
    y1s.append(shape[54][1])
    return x0s, y0s, x1s, y1s


def facecontour(image):
    for i in range(0, 16):
        cv2.line(image, (shape[i][0], shape[i][1]), (shape[i+1][0], shape[i+1][1]), color, thickness)
    return image



def split_list(seq, num):
    avg = len(seq) / float(num)
    out = []
    last = 0.0
    while last < len(seq):
        out.append(seq[int(last):int(last + avg)])
        last += avg
    return out

color = (255,0,0)
thickness = 1
addition = 3
images = glob.glob("fridman_images/*")
images = split_list(images, 16)

for a,frame in tqdm.tqdm(enumerate(images[int(sys.argv[2])])):
    frame = cv2.imread(frame)
    image = frame.copy()
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    rects = detector(gray, 1)
    for (i, rect) in enumerate(rects):
        shape = predictor(gray, rect)
        shape = face_utils.shape_to_np(shape)
        x0, y0, x1, y1 = (shape[48][0]-addition, shape[50][1]-addition, shape[54][0]+addition, shape[57][1]+addition)
        cv2.rectangle(image, (x0-20, y0-45), (x1+20, y1+40),(0,0,0), -1)
        x0s, y0s, x1s, y1s = shapeconverter(shape)
        for x in range(0,len(x0s)):
            cv2.line(image, (x0s[x], y0s[x]), (x1s[x], y1s[x]), color, thickness)
    cv2.imwrite(f"images/lips/{sys.argv[2]}_{a}.jpg", image[shape[30][1]-256:shape[30][1]+256, shape[30][0]-256:shape[30][0]+256])
    cv2.imwrite(f"images/originals/{sys.argv[2]}_{a}.jpg", frame[shape[30][1]-256:shape[30][1]+256, shape[30][0]-256:shape[30][0]+256])