#!/usr/bin/env python
# encoding: utf-8
"""Perform facial landmark localization using menpofit

The estimated landmarks will be saved as .pts files next to the input image,
with the same file stem as the input image e.g. image.jpg -> image.pts

If the output landmark file already exists and the --force/-f flag is not passed
the output file will not be written.

Usage:
  menpofit [-f | --force] <path>...
  menpofit (-h | --help)
  menpofit --version

Options:
  <path>       Perform landmark localization on all images found at path
  --force -f   Overwrite existing annotations
  -h --help    Show this screen.
  --version    Show version.
"""
from docopt import docopt

import menpo.io as mio
from menpo.visualize import print_progress, print_dynamic
from menpo.image.base import normalize_pixels_range
from menpodetect import load_dlib_frontal_face_detector
import menpofit  # needed for version
from menpofit.aam import load_balanced_frontal_face_fitter

from menpocli.io import (resolve_importable_paths, save_pointcloud_as_landmark,
                         build_landmark_output_path)


def prepare_img_for_fit(img):
    new_img = img.copy()
    new_img.pixels = normalize_pixels_range(new_img.pixels)
    if new_img.n_channels == 3:
        new_img = new_img.as_greyscale(mode='luminosity')
    return new_img


def detect_and_fit_images(detector, fitter, img_paths, overwrite):
    for img_path in print_progress(img_paths):
        landmark_out_path = build_landmark_output_path(img_path)
        if not overwrite and landmark_out_path.exists():
            print_dynamic(
                '{} already exists. Please set the --force/-f flag if '
                'you wish to overwrite the file.\n'.format(landmark_out_path)
            )
            continue

        try:
            img = mio.import_image(img_path, normalize=False)
            bboxes = detector(img)
            for i, bbox in enumerate(bboxes):
                fr = fitter.fit_from_bb(prepare_img_for_fit(img), bbox)
                save_pointcloud_as_landmark(img_path, i, fr.final_shape)
        except:
            msg = 'Unable to process {}'.format(img_path).ljust(100)
            print_dynamic(msg + '\n')


if __name__ == '__main__':
    a = docopt(__doc__,
               version='menpofit v{}'.format(menpofit.__version__))
    print('\nM E N P O F I T  ' + 'v' + menpofit.__version__ + '\n')
    detector = load_dlib_frontal_face_detector()
    fitter = load_balanced_frontal_face_fitter()
    image_paths = resolve_importable_paths(a['<path>'])
    detect_and_fit_images(detector, fitter, image_paths, a['--force'])
