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

conn = sqlite3.connect('database.db', check_same_thread=False)
c = conn.cursor()
create(c)

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 init():
    source_imgs = c.execute("SELECT image FROM students").fetchall()
    cwd = os.getcwd()
    source_imgs = [os.path.join(cwd, "images", s[0]) for s in source_imgs]
    source_faces = [Image.open(s) for s in source_imgs]
    # source_embeddings = get_embeddings(model, source_faces)
    source_embeddings = c.execute("SELECT embeddings FROM students").fetchall()
    source_embeddings = [os.path.join(cwd, "embeds", s[0]) for s in source_embeddings]
    source_embeddings = [np.load(s) for s in source_embeddings]
    names = c.execute("SELECT name FROM students").fetchall()
    names = [n[0] for n in names]
    return names, source_faces, source_embeddings


def init():
    source_imgs = c.execute("SELECT image FROM students").fetchall()
    cwd = os.getcwd()
    source_imgs = [os.path.join(cwd, "images", s[0]) for s in source_imgs]
    source_faces = [Image.open(s) for s in source_imgs]
    # source_embeddings = get_embeddings(model, source_faces)
    source_embeddings = c.execute("SELECT embeddings FROM students").fetchall()
    source_embeddings = [os.path.join(cwd, "embeds", s[0]) for s in source_embeddings]
    source_embeddings = [np.load(s) for s in source_embeddings]
    names = c.execute("SELECT name FROM students").fetchall()
    names = [n[0] for n in names]
    return names, source_faces, source_embeddings

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.squeeze(), 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"

source_imgs, source_faces, source_embeddings = init()

detect = gr.Interface(
    find_names,
    [gr.Image(type="filepath", label="Class Photo"), gr.Number(label = "Minimum Size"), gr.Number(label = "Minimum Confidence")],
    [gr.Textbox(label = "Roll No") ,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):
    global i, imgs
    if not os.path.exists(os.path.join(os.getcwd(), "temp")):
        os.mkdir(os.path.join(os.getcwd(), "temp"))
    for file in files:
        # faces, id = extract_faces(detector, os.path.join(os.getcwd(), "temp", file))
        # imgs.append(faces)
        shutil.move(file.name, os.path.join(os.getcwd(), "temp", file.name.split('\\')[-1]))

    return None, "Uploaded!"


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)
    i+=1
    shutil.rmtree(os.path.join(os.getcwd(), "temp"))
    return imgs[i], "Loaded!", None, None

def save_img(label, email, roll, selected):
    global i, imgs, source_imgs, source_faces, source_embeddings
    if label:
        imgs[i][int(selected)].save(os.path.join(os.getcwd(), "images", f"{roll}.jpg"))
        np.save(os.path.join(os.getcwd(), "embeds", f"{roll}.npy"),get_embeddings(model, [imgs[i][int(selected)]]))
        insert(c, label, email, roll, f"{roll}.jpg", f"{roll}.npy")
        conn.commit()
        source_imgs, source_faces, source_embeddings = init()

    imgs[i].pop(int(selected))
    if(len(imgs[i]) == 0):
        i+=1
    if i < len(imgs):
        if not label:
            return imgs[i], None, None, None
        return imgs[i], "Saved!", None, None
    else:
        clear()
        return None, "Finished!", None, None

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

i = -1
imgs = None

with gr.Blocks() as upload:
    gr.Markdown("# Select Images to Upload and click Upload")
    with gr.Row():
        files = gr.Files(file_types=[".jpg", ".jpeg", ".png"], label="Upload images")
        # input = gr.Image(type="filepath")
        alert = gr.Textbox()
    upload_btn = gr.Button(value="Upload")
    upload_btn.click(upload_files, inputs=[files], outputs=[files, alert])
    with gr.Accordion("Annotate", open=False):
        with gr.Row():
            with gr.Column():
                input = gr.Textbox(label = "Name")
                email = gr.Textbox(label = "Email")
                roll = gr.Textbox(label = "Roll No")
            selected = gr.Number(visible=False)
            def get_select_index(evt: gr.SelectData):
                return evt.index
            output = gr.Gallery(label="Found Faces", height = 400)
            output.select(get_select_index, None, selected)
        with gr.Row():
            next_btn = gr.Button(value="Next")
            next_btn.click(load_image, inputs=[], outputs=[output, input, email, roll])
            save_btn = gr.Button(value="Save")
            save_btn.click(save_img, inputs=[input, email, roll, selected], outputs=[output, input, email, roll])
            clear_btn = gr.Button(value="Clear")
            clear_btn.click(clear, inputs=[], outputs=[output, input, email, roll])
# with gr.Blocks() as annotate:

def match(name, roll):
    if roll:
        results = c.execute(f"SELECT name, roll_no, image from students WHERE roll_no = {roll}").fetchall()
    elif name:
        results = c.execute(f'SELECT name, roll_no, image from students WHERE name = "{name}"').fetchall()
    else:
        results = c.execute(f"SELECT name, roll_no, image from students").fetchall()
    names = []
    rolls = []
    images = []
    cwd = os.getcwd()
    for r in results:
        names.append(r[0])
        rolls.append(r[1])
        images.append(os.path.join(cwd, "images", r[2]))
    return names, rolls, images, images

def update(name, roll, rolls, img, selected):
    global source_imgs, source_faces, source_embeddings
    c.execute(f'UPDATE students SET name = "{name}", roll_no = {roll} where roll_no = {rolls[int(selected)]}')
    conn.commit()
    source_imgs, source_faces, source_embeddings = init()
    result = c.execute("SELECT name, roll_no from students").fetchall()
    names = [r[0] for r in result]
    rolls = [r[1] for r in result]
    return None, None, img, img, names, rolls

with gr.Blocks() as find:
    with gr.Row():
        with gr.Column():
            name = gr.Textbox(label = "Name")
            roll = gr.Textbox(label = "Roll No")
        image = gr.Gallery(label = "Matches", height = 400)
    with gr.Row():
        names = gr.State()
        rolls = gr.State()
        images = gr.State()
        selected = gr.Number(visible=False)
        find_btn = gr.Button(value = "Find")
        find_btn.click(match, inputs = [name, roll], outputs = [names, rolls, image, images])
        update_btn = gr.Button(value = "Update")
        update_btn.click(update, inputs = [name, roll, rolls, images, selected], outputs = [name, roll, image, images, names, rolls])
        clear_btn = gr.ClearButton([name, roll, image])
        def get_select_index(evt: gr.SelectData, names, rolls):
            return names[evt.index], rolls[evt.index], evt.index
        image.select(get_select_index, [names, rolls], [name, roll, selected])

tabbed_interface = gr.TabbedInterface(
    [detect, upload, find],
    ["Attendance", "Register", "Find"],
)
if __name__ == "__main__":
    tabbed_interface.launch()