import json
import zipfile
import numpy as np
import cv2
import os
import gradio as gr
from deepface import DeepFace
from ultralytics import YOLO
import urllib.request

with open("config.json", "r") as f:
    config = json.load(f)

FACE_DIST_TRESH = config['FACE_DIST_TRESH']
FACE_DET_TRESH = config['FACE_DET_TRESH']
YOLO_WEIGHTS_URL = config['YOLO_WEIGHTS_URL']

yolo_weights_filename = os.path.basename(YOLO_WEIGHTS_URL)

if not os.path.exists(yolo_weights_filename):
   urllib.request.urlretrieve(YOLO_WEIGHTS_URL, yolo_weights_filename)

model = YOLO(yolo_weights_filename)

def find_distance(base_face, check_face):
  result = DeepFace.verify(base_face, check_face, enforce_detection=False)
  return result['distance']

def find_faces(image):
  outputs = model(image)
  faces = []
  for box in outputs[0].boxes:
    if float(box.conf) >= FACE_DET_TRESH:
      x, y, w, h = [int(coord) for coord in box.xywh[0]]
      x_center, y_center = x + w / 2, y + h / 2
      x1 = int(x_center - w)
      y1 = int(y_center - h)
      crop_img = image[y1:y1+h, x1:x1+w]
      faces.append(crop_img)
  return faces
  
def load_images_from_zip(zip_path):
    images = []
    with zipfile.ZipFile(zip_path, 'r') as zip_file:
        for file_name in zip_file.namelist():
            with zip_file.open(file_name) as file:
                img_bytes = np.asarray(bytearray(file.read()), dtype=np.uint8)
                img = cv2.imdecode(img_bytes, cv2.IMREAD_COLOR)
                images.append(img)
    return [img for img in images if img is not None]

def create_image(images):
    if not images:
        return np.zeros((1, 1, 3), dtype=np.uint8)
    
    max_width = max([img.shape[1] for img in images])
    total_height = sum([img.shape[0] for img in images])

    new_image = np.zeros((total_height, max_width, 3), dtype=np.uint8)

    y_offset = 0
    for img in images:
        new_image[y_offset:y_offset+img.shape[0], :img.shape[1], :] = img
        y_offset += img.shape[0]

    new_image = cv2.resize(new_image, (0, 0), fx=0.5, fy=0.5)

    return new_image

def check(avatars_zip, photos_zip):
  avatars = load_images_from_zip(avatars_zip.name)
  photos = load_images_from_zip(photos_zip.name)
  not_found_faces = []
  for photo in photos:
    photo = cv2.cvtColor(photo, cv2.COLOR_RGB2BGR)
    input_faces = find_faces(photo)
    for input_face in input_faces:
        avatars_checked = 0
        for avatar in avatars:
            avatar = cv2.cvtColor(avatar, cv2.COLOR_RGB2BGR)
            distance = find_distance(avatar, input_face)
            if distance > FACE_DIST_TRESH:
                break
            else:
                avatars_checked+=1
            if avatars_checked == len(avatars):
                not_found_faces.append(input_face)

  return create_image(not_found_faces)

title = '<h1 style="text-align:center">FootballChecker</h1>'

with gr.Blocks(theme='soft', title='FootballChecker') as blocks:
    gr.HTML(title)
    gr.Markdown('Face checking toolkit')
    with gr.Row():
        avatars = gr.inputs.File(label="Avatar Images (zip)")
        photos = gr.inputs.File(label="Check Photos (zip)")
    inputs = [avatars, photos]
    process_button = gr.Button('Check images')
    outputs=gr.outputs.Image(type="numpy", label="Results")
    process_button.click(fn=check, inputs=inputs, outputs=outputs)

blocks.launch()