import os
import faiss
import gradio as gr
from helpers import *
import shutil
from PIL import Image

detector = load_detector()
model = load_model()

source_imgs = []
for r, _, f in os.walk(os.getcwd() + "/images"):
    for file in f:
        if (
            (".jpg" in file.lower())
            or (".jpeg" in file.lower())
            or (".png" in file.lower())
        ):
            exact_path = r + "/" + file
            source_imgs.append(exact_path)

source_faces = []
for img in source_imgs:
  try:
    faces, id = extract_faces(detector, img)
    source_faces.append(faces[id])
    # source_faces.append(Image.open(img))
  except Exception as e:
    print(f"Skipping {img}, {e}")

source_embeddings = get_embeddings(model, source_faces)

def find_names(image, minSize, minConf):
  imgs, _ = extract_faces(detector, image)
  ims = []
  for i, face in enumerate(imgs):
    if((face.size[0] * face.size[1]) > minSize):
      ims.append(face)
  imgs = ims

  embeds = get_embeddings(model, imgs)
  d = np.zeros((len(source_embeddings), len(embeds)))
  for i, s in enumerate(source_embeddings):
    for j, t in enumerate(embeds):
      d[i][j] = findCosineDistance(s, t)
  ids = np.argmin(d, axis = 0)
  names = []
  for j, i in enumerate(ids):
    if 1 - d[i][j] > minConf:
        names.append(source_imgs[i].split("/")[-1].split(".")[0])
    else:
        names.append("Unknown")
  recognition(imgs, ids, names, source_faces, d, source_imgs)
  return ",".join(names), "Recognition.jpg"

detect = gr.Interface(
    find_names,
    [gr.Image(type="filepath", label="Class Photo"), gr.Number(label = "Minimum Size"), gr.Number(label = "Minimum Confidence")],
    ["text" ,gr.Image(type = "filepath", label="Matching")],
    examples = [
        [os.path.join(os.path.dirname(__file__), "examples/group1.jpg"), 1000, 0.3],
        [os.path.join(os.path.dirname(__file__), "examples/group2.jpg"), 1000, 0.3]
    ]
)


def upload_files(files):
    if not os.path.exists(os.path.join(os.getcwd(), "temp")):
        os.mkdir(os.path.join(os.getcwd(), "temp"))
    for file in files:
        shutil.move(file.name, os.path.join(os.getcwd(), "temp", file.name.split('\\')[-1]))
    return None, "Uploaded!"

with gr.Blocks() as upload:
    gr.Markdown("# Select Images to Upload and click Upload")
    with gr.Row():
        input = gr.Files(file_types=[".jpg", ".jpeg", ".png"], label="Upload images")
        # input = gr.Image(type="filepath")
        output = gr.Textbox()
    upload_btn = gr.Button(value="Upload")
    upload_btn.click(upload_files, inputs=[input], outputs=[input, output])

def load_image():
    global i, imgs
    images = os.listdir(os.path.join(os.getcwd(), "temp"))
    imgs = []
    for image in images:
        faces, id = extract_faces(detector, os.path.join(os.getcwd(), "temp", image))
        # imgs.append(Image.open(os.path.join(os.getcwd(), "temp", image)))
        imgs.append(faces[id])
    return imgs[0], "Loaded!"

def save_img(label):
    global i, imgs
    imgs[i].save(os.path.join(os.getcwd(), "images", f"{label}.jpg"))
    i+=1
    if i < len(imgs):
        return imgs[i], "Saved!"
    else:
        clear()
        return None, "Finished!"

def clear():
    global i, imgs
    i = 0
    imgs = None
    shutil.rmtree(os.path.join(os.getcwd(), "temp"))
    return None, None

i = 0
imgs = None

with gr.Blocks() as annotate:
    with gr.Row():
        input = gr.Textbox(label = "Enter Label")
        output = gr.Image(type="pil", label="Image").style(height=400)
    with gr.Row():
        next_btn = gr.Button(value="Next")
        next_btn.click(load_image, inputs=[], outputs=[output, input])
        save_btn = gr.Button(value="Save")
        save_btn.click(save_img, inputs=[input], outputs=[output, input])
        clear_btn = gr.Button(value="Clear")
        clear_btn.click(clear, inputs=[], outputs=[output, input])

tabbed_interface = gr.TabbedInterface(
    [detect, upload, annotate],
    ["Attendance", "Upload", "Annotate"],
)
if __name__ == "__main__":
    tabbed_interface.launch()