import os
import PIL
import cv2
import pickle
import argparse
import numpy as np
import face_alignment
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from matplotlib.path import Path


def parse_args():
    parser = argparse.ArgumentParser(description="Plot facial landmarks from an image.")
    parser.add_argument(
        "--image_path", 
        type=str, 
        default=None, 
        help="Path to the image file."
    )
    parser.add_argument("--size", type=int, default=512)
    parser.add_argument("--crop", action="store_true", help="Crop around the face image.")
    parser.add_argument(
        "--output_dir",
        type=str,
        default="output/landmarks/",
        help="Folder to save landmark images."
    )
    args = parser.parse_args()

    return args

def get_patch(landmarks, color='lime', closed=False):
    contour = landmarks
    ops = [Path.MOVETO] + [Path.LINETO]*(len(contour)-1)
    facecolor = (0, 0, 0, 0)      # Transparent fill color, if open
    if closed:
        contour.append(contour[0])
        ops.append(Path.CLOSEPOLY)
        facecolor = color
    path = Path(contour, ops)
    return patches.PathPatch(path, facecolor=facecolor, edgecolor=color, lw=4)

def bbox_from_landmarks(landmarks):
    landmarks_x, landmarks_y = zip(*landmarks)

    x_min, x_max = min(landmarks_x), max(landmarks_x)
    y_min, y_max = min(landmarks_y), max(landmarks_y)
    width = x_max - x_min
    height = y_max - y_min

    # Give it a little room; I think it works anyway
    x_min -= 25
    y_min -= 25
    width += 50
    height += 50
    bbox = (x_min, y_min, width, height)
    return bbox

def plot_landmarks(landmarks, crop=False, size=512):
    if crop:
        (x_min, y_min, width, height) = bbox_from_landmarks(landmarks)
        # print(x_min, y_min, width, height)
        landmarks_np = np.array(landmarks)
        landmarks_np[:, 0] = (landmarks_np[:, 0] - x_min) * size / width
        landmarks_np[:, 1] = (landmarks_np[:, 1] - y_min) * size / height
        landmarks = landmarks_np.tolist()
    # Precisely control output image size
    dpi = 72
    fig, ax = plt.subplots(1, figsize=[size/dpi, size/dpi], tight_layout={'pad':0})
    fig.set_dpi(dpi)

    black = np.zeros((size, size, 3))
    ax.imshow(black)

    face_patch = get_patch(landmarks[0:17])
    l_eyebrow = get_patch(landmarks[17:22], color='yellow')
    r_eyebrow = get_patch(landmarks[22:27], color='yellow')
    nose_v = get_patch(landmarks[27:31], color='orange')
    nose_h = get_patch(landmarks[31:36], color='orange')
    l_eye = get_patch(landmarks[36:42], color='magenta', closed=True)
    r_eye = get_patch(landmarks[42:48], color='magenta', closed=True)
    outer_lips = get_patch(landmarks[48:60], color='cyan', closed=True)
    inner_lips = get_patch(landmarks[60:68], color='blue', closed=True)

    ax.add_patch(face_patch)
    ax.add_patch(l_eyebrow)
    ax.add_patch(r_eyebrow)
    ax.add_patch(nose_v)
    ax.add_patch(nose_h)
    ax.add_patch(l_eye)
    ax.add_patch(r_eye)
    ax.add_patch(outer_lips)
    ax.add_patch(inner_lips)

    plt.axis('off')

    fig.canvas.draw()
    buffer, (width, height) = fig.canvas.print_to_buffer()
    assert width == height
    assert width == size

    buffer = np.frombuffer(buffer, np.uint8).reshape((height, width, 4))
    buffer = buffer[:, :, 0:3]
    plt.close(fig)
    return PIL.Image.fromarray(buffer)

def get_landmarks(image):
    fa = face_alignment.FaceAlignment(face_alignment.LandmarksType.TWO_D, flip_input=False, face_detector='sfd')
    faces = fa.get_landmarks_from_image(image)
    if faces is None or len(faces) == 0:
        return None
    landmarks = faces[0]
    return landmarks

def save_landmarks(args):
    os.makedirs(args.output_dir, exist_ok=True)
    
    image_name = os.path.basename(args.image_path)
    image = cv2.imread(args.image_path)
    image = cv2.resize(image, (args.size, args.size))
    landmarks = get_landmarks(image)
    if landmarks is None:
        print(f'No faces found in {image_name}')
        return
    
    filename = f'{args.output_dir}/{image_name}'
    if args.crop:
        landmarks_cropped_image = plot_landmarks(landmarks.tolist(), crop=True, size=args.size)
        landmarks_cropped_image.save(filename)
    else:
        landmarks_image = plot_landmarks(landmarks.tolist(), size=args.size)
        landmarks_image.save(filename)
    print(f'Landmark saved in {filename}')

if __name__ == '__main__':
    args = parse_args()
    save_landmarks(args)