# Copyright 2020 Adobe. All rights reserved.
# This file is licensed to you under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License. You may obtain a copy
# of the License at http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software distributed under
# the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR REPRESENTATIONS
# OF ANY KIND, either express or implied. See the License for the specific language
# governing permissions and limitations under the License.
import os
import argparse
import bz2
import torch
from keras.utils import get_file

from face_alignment import image_align
from landmarks_detector import LandmarksDetector

from models.AutoToon import AutoToonModel
from test_utils import save_images


LANDMARKS_MODEL_URL = 'http://dlib.net/files/shape_predictor_68_face_landmarks.dat.bz2'

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


def unpack_bz2(src_path):
    data = bz2.BZ2File(src_path).read()
    dst_path = src_path[:-4]
    with open(dst_path, 'wb') as fp:
        fp.write(data)
    return dst_path


if __name__ == "__main__":
    """
    Runs inference on a directory of input images `in_dir`. Images can have more than one input face, and will
    be numbered starting from 01 through the n faces.

    Images are passed through the AutoToon model, and the following outputs are produced and written to
    the output directory `out_dir`:
    - IMG_NAME_orig.jpg, the original cropped and centered input image,
    - IMG_NAME_out.jpg, the output of the AutoToon model for the input image,
    - IMG_NAME_quiver.jpg, the warping field quiver plot generated by the model,
    - IMG_NAME_overlaid.jpg, the overlaid warping field quiver plot on top of the output image,
    - IMG_NAME_xflow.jpg, the visualized heatmap for the x-direction warping field with the overlaid quiver plot,
    - IMG_NAME_yflow.jpg, the visualized heatmap for the y-direction warping field with the overlaid quiver plot.

    Borrows some code from https://github.com/Puzer/stylegan-encoder (`align_images.py`) to extract and align faces
    from images using dlib, as well as the functions and files from the original FFHQ dataset preparation steps in
    face_alignment.py` and `landmarks_detector.py`.

    Usage: test.py --in_dir in/ --out_dir out/ --scale 1
    """

    landmarks_model_path = unpack_bz2(get_file('shape_predictor_68_face_landmarks.dat.bz2',
                                               LANDMARKS_MODEL_URL, cache_subdir='temp'))

    parser = argparse.ArgumentParser()
    parser.add_argument('-i', '--in_dir', default="in", help='raw input images directory')
    parser.add_argument('-o', '--out_dir', default="out", help='output images directory')
    parser.add_argument('-s', '--scale', type=float, default=1, help='cartoon scaling extent')

    args = parser.parse_args()
    IN_DIR = args.in_dir
    OUT_DIR = args.out_dir
    SCALE = args.scale

    # initialize model and load pretrained weights
    model = AutoToonModel()
    model.eval()
    model.load_model('autotoon', './models', device=device)
    print('successfully loaded AutoToon model weights')
    os.makedirs(OUT_DIR, exist_ok=True)

    landmarks_detector = LandmarksDetector(landmarks_model_path)
    for img_name in os.listdir(IN_DIR):
        raw_img_path = os.path.join(IN_DIR, img_name)
        #print(raw_img_path)
        try:
            if '.DS_Store' not in raw_img_path:
                for i, face_landmarks in enumerate(landmarks_detector.get_landmarks(raw_img_path), start=1):
                    face_img_name = '%s_%02d.png' % (os.path.splitext(img_name)[0], i)
                    aligned_face_path = os.path.join(OUT_DIR, face_img_name)

                    img = image_align(raw_img_path, aligned_face_path, face_landmarks)

                    img_out_name = img_name.split('.')[0] + '_%02d' % i + '.' + img_name.split('.')[1]

                    # see test_utils.py for details on how images are processed and saved
                    save_images(model, OUT_DIR, img=img, img_name=img_out_name, scale=SCALE, visualize=False)
                    print('finished processing image', img_name)
        except Exception as e:
            print(e)
