import gdown
import gradio as gr

import logging
import os

import cv2
import numpy as np
import tensorflow as tf

from ai.detection import detect
from laeo_per_frame.interaction_per_frame_uncertainty import LAEO_computation
from utils.hpe import hpe, project_ypr_in2d
from utils.img_util import resize_preserving_ar, percentage_to_pixel, draw_key_points_pose, \
    visualize_vector, draw_axis, draw_axis_3d, draw_cones

# <a href="https://malga.unige.it/" target="_blank"><nobr>Lab MaLGa UniGe</nobr></a>
WEBSITE = """
<div class="embed_hidden">
<h1 style='text-align: center'>Head Pose Estimation and LAEO computation </h1>
<h2 style='text-align: center'>
<a target="_blank" href="https://github.com/Malga-Vision/LAEO_demo"> <nobr> Code for LAEO </nobr></a>
<br>
<a target="_blank" href="https://github.com/Malga-Vision/HHP-Net/tree/master"> <nobr> Code for HPE </nobr></a>

</h2>
<h2 style='text-align: center'>
<nobr><a href="https://github.com/Malga-Vision" target="_blank"><nobr>MaLGa Vision GitHub</nobr></a> &emsp;</nobr>
</h2>

<h3 style="text-align:center;">
<a href="https://fede1995.github.io/" target="_blank"><nobr>Federico FT</nobr></a> &emsp;
</h3>

<h2> Description </h2>
<p>
This space illustrates a method for Head Pose Estimation and also LAEO detection. The code is based on experiments and research carried out at the UNiversity of Genoa (Italy) in the MaLGa Laboratory. 
This demo has been set up by Federico Figari Tomenotti.
DISCLAIMER: does not work properly on smartphones and sometimes on Safari web browser.
</p>
<h2> Usage </h2>
<p>
The flags allow the user to choose what to display on the result image, and to change the sensitivity for the person detection algorithm.
The Head Pose orientation can be described only as one vector (arrow) or a triplet of angles: yaw, pitch and roll projected on the image plane.
The uncertainty result is the mean of the uncertainty compute on the three angles.
The run botton is needed to run the demo on an image after changing flag settings.
For every detailed explanation on the algorithms refer to the paper which will be out soon.
</p>

</div>
"""

WEBSITE_citation = """
<h2 style='text-align: center'>
Citation
</h2>

If you find this code useful for your research, please use the following BibTeX entry.

``` 
@inproceedings{cantarini2022hhp,
  title={HHP-Net: A light Heteroscedastic neural network for Head Pose estimation with uncertainty},
  author={Cantarini, Giorgio and Tomenotti, Federico Figari and Noceti, Nicoletta and Odone, Francesca},
  booktitle={Proceedings of the IEEE/CVF Winter Conference on applications of computer vision},
  pages={3521--3530},
  year={2022}
}
```"""


def load_image(camera, ):
    # Capture the video frame by frame
    try:
        ret, frame = camera.read()
        return True, frame
    except:
        logging.Logger('Error reading frame')
        return False, None


def demo_play(img, laeo=True, rgb=False, show_keypoints=True, only_face=False, Head_Pose_representation='Vector', detection_threshold=0.45):
    # webcam in use

    # gpus = tf.config.list_physical_devices('GPU')

    # img = np.array(frame)


    img_resized, new_old_shape = resize_preserving_ar(img, input_shape_od_model)

    if not rgb:
        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) # covert at grey scale
        img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR) # it is still grey scale but with 3 channels to add the colours of the points and lines
        # img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
    else: # if RGB
        # img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
        pass

    print('inference centernet')
    detections, elapsed_time = detect(model, img_resized, detection_threshold,
                                      new_old_shape)  # detection classes boxes scores
    # probably to draw on resized
    # img_with_detections = draw_detections(img_resized, detections, max_boxes_to_draw, None, None, None)
    # cv2.imshow("aa", img_with_detections)

    det, kpt = percentage_to_pixel(img.shape, detections['detection_boxes'], detections['detection_scores'],
                                   detections['detection_keypoints'], detections['detection_keypoint_scores'])

    # center_xy, yaw, pitch, roll = head_pose_estimation(kpt, 'centernet', gaze_model=gaze_model)

    # _________ extract hpe and print to img
    people_list = []

    print('inferece hpe')

    for j, kpt_person in enumerate(kpt):
        yaw, pitch, roll, tdx, tdy = hpe(gaze_model, kpt_person, detector='centernet')

        # img = draw_axis_3d(yaw[0].numpy()[0], pitch[0].numpy()[0], roll[0].numpy()[0], image=img, tdx=tdx, tdy=tdy,
        #                    size=50)

        people_list.append({'yaw'      : yaw[0].numpy()[0],
                            'yaw_u'    : yaw[0].numpy()[1],
                            'pitch'    : pitch[0].numpy()[0],
                            'pitch_u'  : pitch[0].numpy()[1],
                            'roll'     : roll[0].numpy()[0],
                            'roll_u'   : roll[0].numpy()[1],
                            'center_xy': [tdx, tdy]
                            })


    if show_keypoints:
        # img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        for i in range(len(det)):
            # img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            img = draw_key_points_pose(img, kpt[i], only_face=only_face)

    # call LAEO
    clip_uncertainty = 0.5
    binarize_uncertainty = False
    if laeo:
        interaction_matrix = LAEO_computation(people_list, clipping_value=clip_uncertainty,
                                              clip=binarize_uncertainty)
    else:
        interaction_matrix = np.zeros((len(people_list), len(people_list)))
    # coloured arrow print per person

    print(f'Head pose representation: {Head_Pose_representation}')
    def visualise_hpe(yaw, pitch, roll, image=None, tdx=None, tdy=None, size=50, yaw_uncertainty=-1, pitch_uncertainty=-1, roll_uncertainty=-1, openpose=False, title="", color=(255, 0, 0)):
        if str(Head_Pose_representation).lower() == 'vector':
            vector = project_ypr_in2d(person['yaw'], person['pitch'], person['roll'])
            image = visualize_vector(image, [tdx, tdy], vector, title=title, color=color)
            return image
        elif str(Head_Pose_representation).lower() == 'axis':
           image = draw_axis_3d(yaw, pitch, roll, image=image, tdx=tdx, tdy=tdy, size=size)
           return image
        elif str(Head_Pose_representation).lower() == 'cone':
            _, image = draw_cones(yaw, pitch, roll, unc_yaw=yaw_uncertainty, unc_pitch=pitch_uncertainty, unc_roll=roll_uncertainty, image=image, tdx=tdx, tdy=tdy, size=size)
            return image
        else:
            return image


    for index, person in enumerate(people_list):
        green = round((max(interaction_matrix[index, :])) * 255)
        colour = (0, green, 0)
        if green < 40:
            colour = (255, 0, 0)
        img = visualise_hpe(person['yaw'], person['pitch'], person['roll'], image=img, tdx=person['center_xy'][0], tdy=person['center_xy'][1], size=50, yaw_uncertainty=person['yaw_u'], pitch_uncertainty=person['pitch_u'], roll_uncertainty=person['roll_u'], title="", color=colour)
        # vector = project_ypr_in2d(person['yaw'], person['pitch'], person['roll'])
        # img = visualize_vector(img, person['center_xy'], vector, title="",
        #                        color=colour)
    uncertainty_mean = [i['yaw_u'] + i['pitch_u'] + i['roll_u'] for i in people_list]
    uncertainty_mean_str =  ''.join([str(round(i, 2)) + ' ' for i in uncertainty_mean])
    return img, uncertainty_mean_str


if __name__=='__main__':
    if not os.path.exists("LAEO_demo_data"):
        gdown.download_folder("https://drive.google.com/drive/folders/1nQ1Cb_tBEhWxy183t-mIcVH7AhAfa6NO?usp=drive_link",
                              use_cookies=False)

    # Get the list of all files and directories
    path = "LAEO_demo_data/examples"
    dir_list = os.listdir(path)
    print("Files and directories in '", path, "' :")

    # prints all files
    print(dir_list)

    gaze_model_path = 'LAEO_demo_data/head_pose_estimation'
    gaze_model = tf.keras.models.load_model(gaze_model_path, custom_objects={"tf": tf})
    path_to_model = 'LAEO_demo_data/keypoint_detector/centernet_hg104_512x512_kpts_coco17_tpu-32'
    model = tf.saved_model.load(os.path.join(path_to_model, 'saved_model'))

    input_shape_od_model = (512, 512)
    # params
    min_score_thresh, max_boxes_to_draw, min_distance = .25, 50, 1.5

    print("Num GPUs Available: ", len(tf.config.list_physical_devices('GPU')))

    function_to_call = demo_play
    # outputs = gr.Image(shape=(512, 512))
    live = True
    title = "Head Pose Estimation and LAEO"

    print(os.getcwd())

    with gr.Blocks() as demo:
        gr.Markdown(WEBSITE)
        with gr.Tab("demo_webcam"):
            with gr.Row():
                with gr.Column():
                    input_img = gr.Image(label="Input Image", source="webcam")
                    laeo = gr.Checkbox(value=True, label="LAEO", info="Compute and display LAEO")
                    rgb = gr.Checkbox(value=False, label="rgb", info="Display output on W/B image")
                    show_keypoints = gr.Checkbox(value=True, label="show_keypoints", info="Display keypoints on image")
                    show_keypoints_only_face = gr.Checkbox(value=True, label="show_keypoints_only_face",
                                                              info="Display only face keypoints on image")
                    Head_Pose_representation = gr.Radio(["Vector", "Axis", "None"], label="Head_Pose_representation",
                                                        info="Which representation to show", value="Vector")
                    detection_threshold = gr.Slider(0.01, 1, value=0.45, step=0.01, interactive=True,
                                                    label="detection_threshold", info="Choose in [0, 1]")

                    button = gr.Button(label="RUN", type="default")
                with gr.Column():
                    outputs = gr.Image(label="Output Image", shape=(512, 512))
                    uncert = gr.Label(label="Uncertainty", value="0.0")

        input_img.change(function_to_call, inputs=[input_img, laeo, rgb, show_keypoints, show_keypoints_only_face,
                                                       Head_Pose_representation, detection_threshold], outputs=[outputs, uncert])
        button.click(function_to_call, inputs=[input_img, laeo, rgb, show_keypoints, show_keypoints_only_face,
                                                       Head_Pose_representation, detection_threshold], outputs=[outputs, uncert])

        with gr.Tab("demo_upload"):
            with gr.Row():
                with gr.Column():
                    input_img = gr.Image(label="Input Image", source="upload")
                    laeo = gr.Checkbox(value=True, label="LAEO", info="Compute and display LAEO")
                    rgb = gr.Checkbox(value=False, label="rgb", info="Display output on W/B image")
                    show_keypoints = gr.Checkbox(value=True, label="show_keypoints", info="Display keypoints on image")
                    show_keypoints_only_face = gr.Checkbox(value=True, label="show_keypoints_only_face",
                                                           info="Display only face keypoints on image")
                    Head_Pose_representation = gr.Radio(["Vector", "Axis", "None"],
                                                        label="Head_Pose_representation",
                                                        info="Which representation to show", value="Vector")
                    detection_threshold = gr.Slider(0.01, 1, value=0.45, step=0.01, interactive=True,
                                                    label="detection_threshold", info="Choose in [0, 1]")

                    button = gr.Button(label="RUN", type="default")
                with gr.Column():
                    outputs = gr.Image(height=238, width=585, label="Output Image")
                    uncert = gr.Label(label="Uncertainty", value="0.0")
                    examples_text =gr.Markdown("## Image Examples")
                    examples = gr.Examples([["LAEO_demo_data/examples/1.jpg"], ["LAEO_demo_data/examples/300wlp_0.png"],
                                ["LAEO_demo_data/examples/AWFL_2.jpg"],
                                ["LAEO_demo_data/examples/BIWI_3.png"]], inputs=input_img,)  # add all other flags

        input_img.change(function_to_call, inputs=[input_img, laeo, rgb, show_keypoints, show_keypoints_only_face,
                                                   Head_Pose_representation, detection_threshold],
                         outputs=[outputs, uncert])
        button.click(function_to_call, inputs=[input_img, laeo, rgb, show_keypoints, show_keypoints_only_face,
                                               Head_Pose_representation, detection_threshold],
                     outputs=[outputs, uncert])

        gr.Markdown(WEBSITE_citation)

    demo.launch()