import os
import gradio as gr
import yaml
from argparse import ArgumentParser
from tqdm import tqdm

import numpy as np
import imageio
from skimage.transform import resize
from skimage import img_as_ubyte
from scipy.spatial import ConvexHull
import torch
from sync_batchnorm import DataParallelWithCallback
import face_alignment

from modules.generator import OcclusionAwareGenerator_SPADE
from modules.keypoint_detector import KPDetector


def normalize_kp(kp_source, kp_driving, kp_driving_initial, adapt_movement_scale=False,
                 use_relative_movement=False, use_relative_jacobian=False):
    kp_new = {k: v for k, v in kp_driving.items()}

    if adapt_movement_scale:
        source_area = ConvexHull(kp_source['value'][0].data.cpu().numpy()).volume
        driving_area = ConvexHull(kp_driving_initial['value'][0].data.cpu().numpy()).volume
        adapt_movement_scale = np.sqrt(source_area) / np.sqrt(driving_area)
        kp_new['value'] = kp_driving['value'] * adapt_movement_scale    # for reenactment demo
    else:
        adapt_movement_scale = 1

    if use_relative_movement:
        kp_value_diff = (kp_driving['value'] - kp_driving_initial['value'])
        kp_value_diff *= adapt_movement_scale
        kp_new['value'] = kp_value_diff + kp_source['value']

        if use_relative_jacobian:
            jacobian_diff = torch.matmul(kp_driving['jacobian'], torch.inverse(kp_driving_initial['jacobian']))
            kp_new['jacobian'] = torch.matmul(jacobian_diff, kp_source['jacobian'])

    return kp_new


def load_checkpoints(config_path, checkpoint_path, cpu=False):
    with open(config_path) as f:
        # config = yaml.load(f)
        config = yaml.load(f, Loader=yaml.FullLoader)

    generator = OcclusionAwareGenerator_SPADE(**config['model_params']['generator_params'],
                                              **config['model_params']['common_params'])
    if not cpu:
        generator.cuda()

    kp_detector = KPDetector(**config['model_params']['kp_detector_params'],
                             **config['model_params']['common_params'])
    if not cpu:
        kp_detector.cuda()

    if cpu:
        checkpoint = torch.load(checkpoint_path, map_location=torch.device('cpu'))
    else:
        checkpoint = torch.load(checkpoint_path)

    generator.load_state_dict(checkpoint['generator'])
    kp_detector.load_state_dict(checkpoint['kp_detector'])

    if not cpu:
        generator = DataParallelWithCallback(generator)
        kp_detector = DataParallelWithCallback(kp_detector)

    generator.eval()
    kp_detector.eval()

    return generator, kp_detector


def make_animation(source_image, driving_video, generator, kp_detector, relative=True, adapt_movement_scale=True,
                   cpu=False):
    with torch.no_grad():
        predictions = []
        source = torch.tensor(source_image[np.newaxis].astype(np.float32)).permute(0, 3, 1, 2)
        if not cpu:
            source = source.cuda()
        driving = torch.tensor(np.array(driving_video)[np.newaxis].astype(np.float32)).permute(0, 4, 1, 2, 3)
        kp_source = kp_detector(source)
        kp_driving_initial = kp_detector(driving[:, :, 0])

        for frame_idx in tqdm(range(driving.shape[2])):
            driving_frame = driving[:, :, frame_idx]
            if not cpu:
                driving_frame = driving_frame.cuda()
            kp_driving = kp_detector(driving_frame)
            kp_norm = normalize_kp(kp_source=kp_source, kp_driving=kp_driving,
                                   kp_driving_initial=kp_driving_initial, use_relative_movement=relative,
                                   use_relative_jacobian=relative, adapt_movement_scale=adapt_movement_scale)
            out = generator(source, kp_source=kp_source, kp_driving=kp_norm)

            predictions.append(np.transpose(out['prediction'].data.cpu().numpy(), [0, 2, 3, 1])[0])
    return predictions


def find_best_frame_func(source, driving, cpu=False):
    def normalize_kp_infunc(kp):
        kp = kp - kp.mean(axis=0, keepdims=True)
        area = ConvexHull(kp[:, :2]).volume
        area = np.sqrt(area)
        kp[:, :2] = kp[:, :2] / area
        return kp

    fa = face_alignment.FaceAlignment(face_alignment.LandmarksType._2D, flip_input=True,
                                      device='cpu' if cpu else 'cuda')
    kp_source = fa.get_landmarks(255 * source)[0]
    kp_source = normalize_kp_infunc(kp_source)
    norm = float('inf')
    frame_num = 0
    for i, image in tqdm(enumerate(driving)):
        kp_driving = fa.get_landmarks(255 * image)[0]
        kp_driving = normalize_kp_infunc(kp_driving)
        new_norm = (np.abs(kp_source - kp_driving) ** 2).sum()
        if new_norm < norm:
            norm = new_norm
            frame_num = i
    return frame_num


def drive_im(source_image, driving_image, adapt_scale):
    source_image = resize(source_image, (256, 256))[..., :3]
    driving_image = [resize(driving_image, (256, 256))[..., :3]]

    prediction = make_animation(source_image, driving_image, generator, kp_detector, relative=False,
                                adapt_movement_scale=adapt_scale, cpu=cpu)
    return img_as_ubyte(prediction[0])


def drive_vi(source_image, driving_video, mode, find_best_frame, best_frame, relative, adapt_scale):
    reader = imageio.get_reader(driving_video)
    fps = reader.get_meta_data()['fps']
    driving_video = []
    try:
        for im in reader:
            driving_video.append(im)
    except RuntimeError:
        pass
    reader.close()

    
    if mode == 'reconstruction':
        source_image = driving_video[0]

    source_image = resize(source_image, (256, 256))[..., :3]
    driving_video = [resize(frame, (256, 256))[..., :3] for frame in driving_video]

    i = 0
    if find_best_frame != "specific ref frame" or best_frame > 0:
        i = best_frame if find_best_frame == "specific ref frame" else find_best_frame_func(source_image, driving_video, cpu=cpu)
        print("Best frame: " + str(i))
        driving_forward = driving_video[i:]
        driving_backward = driving_video[:(i + 1)][::-1]
        predictions_forward = make_animation(source_image, driving_forward, generator, kp_detector,
                                             relative=relative, adapt_movement_scale=adapt_scale, cpu=cpu)
        predictions_backward = make_animation(source_image, driving_backward, generator, kp_detector,
                                              relative=relative, adapt_movement_scale=adapt_scale, cpu=cpu)
        predictions = predictions_backward[::-1] + predictions_forward[1:]
    else:
        predictions = make_animation(source_image, driving_video, generator, kp_detector, relative=relative,
                                     adapt_movement_scale=adapt_scale, cpu=cpu)
    result_video_path = "result_video.mp4"
    imageio.mimsave(result_video_path, [img_as_ubyte(frame) for frame in predictions], fps=fps)
    return result_video_path, i


config = "config/vox-256.yaml"
checkpoint = "00000099-checkpoint.pth.tar"
cpu = True # decided by the deploying environment

description = "We propose a Face Neural Volume Rendering (FNeVR) network for more realistic face animation, by taking the merits of 2D motion warping on facial expression transformation and 3D volume rendering on high-quality image synthesis in a unified framework.<br>[Paper](https://arxiv.org/abs/2209.10340) and [Code](https://github.com/zengbohan0217/FNeVR)"
im_description = "We can animate a face portrait by a single image in this tab.<br>Please input the origin face and the driving face which provides pose and expression information, then we can obtain the virtual generated face.<br>We can select \"adaptive scale\" parameter for better optic flow estimation using adaptive movement scale based on convex hull of keypoints."
vi_description = "We can animate a face portrait by a video in this tab.<br>Please input the origin face and the driving video which provides pose and expression information, then we can obtain the virtual generated video.<br>Please select inference mode (reenactment for different identities and reconstruction for the same identities).<br>We can select \"relative motion\" paramter to use relative keypoint coordinates for preserving global object geometry, select \"adaptive scale\" parameter for better optic flow estimation using adaptive movement scale based on convex hull of keypoints, and select \"find best ref frame\" parameter to generate video from the frame that is the most alligned with source image."
acknowledgements = "This work was supported by “the Fundamental Research Funds for the Central Universities”, and the National Natural Science Foundation of China under Grant 62076016, Beijing Natural Science Foundation-Xiaomi Innovation Joint Fund L223024. Besides, we gratefully acknowledge the support of [MindSpore](https://www.mindspore.cn), CANN (Compute Architecture for Neural Networks) and Ascend AI processor used for this research.<br>Our FNeVR implementation is inspired by [FOMM](https://github.com/AliaksandrSiarohin/first-order-model) and [DECA](https://github.com/YadiraF/DECA). We appreciate the authors of these papers for making their codes available to the public."

generator, kp_detector = load_checkpoints(config_path=config, checkpoint_path=checkpoint, cpu=cpu)

# iface = gr.Interface(fn=drive_im, 
#                     inputs=[gr.Image(label="Origin face"), 
#                             gr.Image(label="Driving face"), 
#                             gr.CheckboxGroup(label="adapt scale")], 
#                     outputs=gr.Image(label="Generated face"), examples=[["sup-mat/source.png"], ["sup-mat/driving.png"]], 
#                     title="Demostration of FNeVR", description=description)

with gr.Blocks(title="Demostration of FNeVR") as demo:
    gr.Markdown("# <center> Demostration of FNeVR")
    gr.Markdown(description)

    with gr.Tab("Driving by image"):
        gr.Markdown(im_description)

        with gr.Row():
            with gr.Column():
                gr.Markdown("#### Inputs")
                inp2 = gr.Image(label="Driving face")
                inp1 = gr.Image(label="Origin face")
                
                gr.Markdown("#### Parameter")
                inp3 = gr.Checkbox(value=True, label="adaptive scale")

                btn1 = gr.Button(value="Animate")
            with gr.Column():
                gr.Markdown("#### Output")
                outp = gr.Image(label="Generated face")

                with gr.Row():
                    with gr.Column():
                        btn2 = gr.Button(value="Reset")
                    with gr.Column():
                        btn3 = gr.Button(value="Cancel")

        gr.Examples([["sup-mat/driving.png", "sup-mat/source.png"]], [inp2, inp1])

        def reset_output():
            return outp.update(value=None)
        
        def reset_all():
            return inp1.update(value=None), inp2.update(value=None), inp3.update(value=True), outp.update(value=None)

        run = btn1.click(fn=drive_im, inputs=[inp1, inp2, inp3], outputs=outp)
        btn2.click(fn=reset_all, outputs=[inp1, inp2, inp3, outp])
        btn3.click(fn=reset_output, outputs=[outp], cancels=[run])
    with gr.Tab("Driving by video"):
        gr.Markdown(vi_description)

        with gr.Row():
            with gr.Column():
                gr.Markdown("#### Inputs")
                inp2 = gr.Video(label="Driving video")
                inp1 = gr.Image(label="Origin face")
                
                gr.Markdown("#### Parameters")
                inp3 = gr.Radio(choices=["reenactment", "reconstruction"], value="reenactment", label="mode (if \"reconstruction\" selected, origin face is the first frame of driving video)")
                inp6 = gr.Checkbox(value=True, label="relative motion")
                inp7 = gr.Checkbox(value=True, label="adaptive scale")
                inp4 = gr.Radio(choices=["find best ref frame (more time consumed)", "specific ref frame"], value="find best ref frame (more time consumed)", label="set ref frame (used by relative motion and adaptive scale)")
                inp5 = gr.Number(label="specific ref frame (default: 0)", value=0, precision=0, visible=False)

                def reset_ref(inp4):
                    return inp5.update(visible=True) if inp4 == "specific ref frame" else inp5.update(value=0, visible=False)

                inp4.change(fn=reset_ref, inputs=inp4, outputs=inp5)

                btn1 = gr.Button(value="Animate")
            with gr.Column():
                gr.Markdown("#### Output")
                outp1 = gr.Video(label="Generated video")
                outp2 = gr.Number(label="Ref frame", value=0, precision=0)

                # file = gr.File(value="result_video.mp4", visible=False)

                with gr.Row():
                    with gr.Column():
                        btn2 = gr.Button(value="Reset")
                    with gr.Column():
                        btn3 = gr.Button(value="Cancel")

        gr.Examples([["sup-mat/driving.mp4", "sup-mat/source_for_video.png", "specific ref frame", 53]], [inp2, inp1, inp4, inp5])

        def reset_output():
            return outp1.update(value=None), outp2.update(value=0)

        def reset_all():
            return inp1.update(value=None), inp2.update(value=None), inp3.update(value="reenactment"), inp4.update(value="find best ref frame (more time consumed)"), inp5.update(value=0), inp6.update(value=True), inp7.update(value=True), outp1.update(value=None), outp2.update(value=0)

        run = btn1.click(fn=drive_vi, inputs=[inp1, inp2, inp3, inp4, inp5, inp6, inp7], outputs=[outp1, outp2])
        btn2.click(fn=reset_all, outputs=[inp1, inp2, inp3, inp4, inp5, inp6, inp7, outp1, outp2])
        btn3.click(fn=reset_output, outputs=[outp1, outp2], cancels=[run])
    with gr.Tab("Real time animation"):
        gr.Markdown("Real time animation is coming soon.")

    gr.Markdown("## Acknowledgements")
    gr.Markdown(acknowledgements)

demo.queue()
demo.launch()