import os
import pathlib
import glob

import random
from PIL import Image
import plotly.graph_objects as go

import torch
import numpy as np
from basicsr.utils import imwrite
from gfpgan import GFPGANer
import face_recognition

import streamlit as st
from streamlit_image_comparison import image_comparison


DIR = os.getcwd()
model_path = os.path.join(DIR, 'GFPGANCleanv1-NoCE-C2.pth')

bg_upsampler = None

restorer = GFPGANer(
    model_path='GFPGANCleanv1-NoCE-C2.pth',
    upscale=2,
    arch='clean',
    channel_multiplier=2,
    bg_upsampler=bg_upsampler)

def enhancer(image):
    image = np.array(image)
    cropped_faces, restored_faces, restored_img = restorer.enhance(
        image, has_aligned=False, only_center_face=False, paste_back=True)

    return Image.fromarray(restored_faces[0])

def inference(img1, img2, tolerance):

    img1 = face_recognition.face_encodings(img1)[0]
    img2 = face_recognition.face_encodings(img2)[0]

    confidence = face_recognition.face_distance([img1], img2)[0]
    
    if confidence <= tolerance:
        result = 'People in the two images are same!'
    else:
        result = 'People in the two images are different!'
        
    return 1-confidence, result

def main():
    
    st.set_page_config(page_title='QFace.ai', page_icon='🤖', layout='centered')
    st.title('Qface.ai')
    col1, col2 = st.columns(2)
    tolerance = 0.6
    
    with st.form('Input Form'):
        
        with col1:
            st.header('Image 1')
            img1 = st.file_uploader('Upload Image 1', type=['png', 'jpg', 'jpeg'], accept_multiple_files=False)

        with col2:
            st.header('Image 2')
            img2 = st.file_uploader('Upload Image 2', type=['png', 'jpg', 'jpeg'], accept_multiple_files=False)
                
        submitted = st.form_submit_button('Enhance and Compare')
        
        if submitted:
            if img1 is not None and img2 is not None:
                
                with col1:
                    img1 = Image.open(img1)
                    eimg1 = enhancer(img1)
                    image_comparison(img1, eimg1,
                                     width=350,
                                     label1='Before',
                                     label2='After',
                                     show_labels=True,
                                     make_responsive=False)
                with col2:
                    img2 = Image.open(img2)
                    eimg2 = enhancer(img2)
                    image_comparison(img2, eimg2,
                                     width=350,
                                     label1='Before',
                                     label2='After',
                                     show_labels=True,
                                     make_responsive=False)
                    
            confidence, result = inference(np.array(eimg1), np.array(eimg2), tolerance)
            st.subheader(result)
            
            fig = go.Figure(go.Indicator(domain = {'x': [0, 1], 'y': [0, 1]},
                                         value = round(confidence, 3),
                                         mode = 'gauge+number',
                                         title = {'text': 'Confidence'},
                                         delta = {'reference': 0.5},
                                         gauge = {'axis': {'range': [None, 1]}}))
            
            st.plotly_chart(fig, use_container_width=True)
    

if __name__ == '__main__':
    main()