import gradio as gr
import tensorflow as tf
import numpy as np
import cv2
from mtcnn.mtcnn import MTCNN
from numpy import asarray
from annoy import AnnoyIndex
import json
from facerecognition import facenet, facenet_pytorch
from dotenv import load_dotenv
import os
from util import file_helper, ml_math


load_dotenv()
# model = tf.lite.Interpreter('model/face_net.tflite')
# model.allocate_tensors()


tree_path = file_helper.download_gdrive(os.getenv('TREE_ID'), dir='tmp')
dataset_path = file_helper.download_gdrive(os.getenv('DATASET_ID'), dir="tmp")    

model_facenet = facenet.Facenet(tree_path)
model_facenet_pytorch = facenet_pytorch.FacenetPytorch(tree_path)

#
data = np.load(dataset_path)
# url_ids = data.f.url_ids
# image_urls = data.f.image_urls
face_url = data.f.face_url


def find_images(image, model):
    print('using model: ', model)    
    print((image).shape)

    face_ai = model_facenet_pytorch if model == 1 else model_facenet
    t = AnnoyIndex(512, 'angular')
    t.load(face_ai.face_tree_path())

    # image =  cv2.resize(image, (512,512))
    faces = face_ai.extract_faces(pixels = image)    
    print(f'{len(faces)} face detected...')
    if len(faces) == 0:
      return 'Foto Wajah Tidak Terdeteksi, pastikan kamu upload foto yang ada wajahnya ya...'
    if len(faces) > 1:
        return f'{len(faces)} wajah terdeteksi! UPLOAD FOTO HANYA DENGAN 1 WAJAH'
    
    result = []
    face_emb = face_ai.get_embedding(faces[0])
    nns = t.get_nns_by_vector(face_emb, 20)
    for n in nns:
        # result.append(drive_ids[n])
        face_origin = t.get_item_vector(n)
        distance = ml_math.cosine(face_emb, face_origin)
        print(f'{n} --> {distance}')
        if distance < 0.7: continue
        # link = image_urls[url_ids[n]] # f'https://drive.google.com/uc?export=view&id={drive_ids[n]}'
        link = face_url[n]
        result.append(link)


    if len(result) == 0:
        return '''
        Foto Serupa Tidak Ditemukan, Tips: 
         - Gunakan foto terupdate,
         - Jika kamu punya foto saat lari, lebih baik gunakan foto itu,
         - Jangan menggunakan foto blur,
         - Jangan upload foto yang menggunakan aksesoris yang tidak kamu gunakan saat lari (semisal kacamata)
        '''
    

    # return json.dumps(result)

    # link_format = [f'<a href="{link}" target="_blank">{i+1}. {link}</a>' for i, link in enumerate(result)]
    link_format = [f'<a href="{link}" target="_blank"> <img src="{link}" height="150"></img></a>' for i, link in enumerate(result)]
    message = f'''
    ### LINK DOWNLOAD
    {len(result)} photo di temukan, klik untuk memperbesar:    
    {"<br/><br/>".join(link_format)}
    '''
    # return  f'### Download foto kamu dari {len(result)} link berikut: ' + ' \n'.join(result)
    return message


models = ['Facenet','Facenet Pytorch']
gr_inputs = [
  "image",
  gr.Dropdown(models, type='index', value=models[0], label='Pilih Salah Satu')
  ]

gr.Markdown('Upload Foto Wajah Kamu (Pastikan hanya terdapat SATU wajah pada)')
iface = gr.Interface(fn=find_images, inputs=gr_inputs, outputs="markdown")
iface.launch()