import json
import os
import gradio as gr
from numpy.linalg import norm
from gradio_client import Client

os.environ['CUDA_LAUNCH_BLOCKING'] = '0'
G = ['Female', 'Male']

token = os.environ['token']
client = Client("https://vtechai-face-demo.hf.space/", token)

def analyze(im, conf=0.65, area="Asia", f_age=False, f_gender=False, f_emotion=False):
    im_path, jj = client.predict(im, conf, area, f_age, f_gender, f_emotion, fn_index=1)
    with open(jj) as f:
        jss = json.load(f)
    return im_path, jss
    
def face_search(im1,im2, threshold):
    im3, im4, te = client.predict(im1, im2, threshold, fn_index=3)
    return im3, im4, te


with gr.Blocks() as face_compare:
    with gr.Row():
        im1 = gr.Image(label='Register', type='filepath',  height=300) #.style(full_width=True, height=300)
        im2 = gr.Image(label='Image for Search', type='filepath',  height=300) #.style(full_width=True, height=300)
    
    with gr.Row():
        im3 = gr.Image(label='Output', height=300) #.style(height=300, full_width=True, full_height=True)
        im4 = gr.Image(label='Output', height=300) #.style(height=300, full_width=True, full_height=True)
    
    sl = gr.Slider(0.3, 1, step=0.05, value=0.5, label='Face Matching Threshold')
    text = gr.Text(label="Output", interactive=False)
    with gr.Row():
        btn = gr.Button(value="Run")
        btn_clean = gr.ClearButton([im1, im2, im3, im4])
    btn.click(fn=face_search, inputs=[im1, im2, sl], outputs=[im3, im4, text])
    # btn2 = gr.Button(value="Check", link="https://manhduy160396.wixsite.com/vtech")

    # example
    gr.Examples(
        examples=[[
            os.path.join(os.path.dirname(__file__), "example/ronaldo.jpg"),
            os.path.join(os.path.dirname(__file__), "example/face2.jpg")
        ]],
        inputs=[im1, im2]

    )

with gr.Blocks() as face_analyze:
    with gr.Row():
        im1 = gr.Image(shape=(300, 300), type='filepath', height=300, container=True)
        im2 = gr.Image(shape=(300, 300), height=300, container=True)

    with gr.Row():
        with gr.Column():
            area = gr.Radio(["Asia", "Europe/America"], label="Area?", value="Asia")
            cb_age = gr.Checkbox(label="Age", value=True)
            cb_gender = gr.Checkbox(label="Gender", value=True)
            cb_emotion = gr.Checkbox(label="Emotion", value=True)
            sl = gr.Slider(0, 1, step=0.05, value=0.65, label='Confidence Threshold')
        
        with gr.Column():
            js = gr.JSON(label="json")
    
    with gr.Row():
        btn = gr.Button(value="Run")
        btn_clean = gr.ClearButton([im1, im2])

    btn.click(fn=analyze, inputs=[im1, sl, area, cb_age, cb_gender, cb_emotion], outputs=[im2, js])
    btn2 = gr.Button(value="Check", link="https://manhduy160396.wixsite.com/vtech")


with gr.Blocks() as demo:
    gr.Markdown('<h1 style="text-align: center;">V-FaceSDK</h1>')
    gr.Markdown("*Chọn chức năng bạn muốn trải nghiệm")
    gr.TabbedInterface([face_analyze, face_compare], ["Face Analyze", "Face Compare"])


if __name__ == "__main__":
    demo.launch()