# tk_face_recog.py

# tk_face_recog.py ZZZ

import os
import random
import string
from tkinter import Tk, simpledialog, messagebox, filedialog, Button, Frame, Label, Scrollbar, Listbox, Toplevel, VERTICAL
from PIL import Image, ImageOps, ImageTk
import math

def browse_for_directory():
    return filedialog.askdirectory(initialdir=current_directory, title='Select Directory')

def browse_for_file():
    return filedialog.askopenfilename(initialdir=current_directory, title='Select Training Data File')

def get_images_and_labels(image_directory, subdirs, label_mapping):
    image_paths = []
    for subdir in subdirs:
        path = os.path.join(image_directory, subdir)
        for filename in os.listdir(path):
            image_paths.append(os.path.join(path, filename))

    images = []
    labels = []
    for image_path in image_paths:
        img = Image.open(image_path).convert('L')
        img = ImageOps.fit(img, (100, 100), Image.ANTIALIAS)
        image = list(img.getdata())
        subdir = os.path.basename(os.path.dirname(image_path))
        label = label_mapping[subdir]
        images.append(image)
        labels.append(label)
    return images, labels

def train_recognizer(image_directory, subdirs, label_mapping):
    images, labels = get_images_and_labels(image_directory, subdirs, label_mapping)
    recognizer = {}
    for i, label in enumerate(labels):
        recognizer[label] = images[i]
    return recognizer

def euclidean_distance(img1, img2):
    return math.sqrt(sum((img1[i] - img2[i]) ** 2 for i in range(len(img1))))

def show_image(image_path, image_listbox, listbox_side):
    global top
    try:
        top.destroy()
    except:
        pass
    top = Toplevel(root)
    img = Image.open(image_path)
    width, height = img.size
    ratio = height / width
    new_height = int(500 * ratio)
    img = img.resize((500, new_height), Image.ANTIALIAS)
    img = ImageTk.PhotoImage(img)
    panel = Label(top, image=img)
    panel.image = img
    panel.place(x=0, y=0)

    screen_width = root.winfo_screenwidth()
    screen_height = root.winfo_screenheight()
    window_width = 500
    window_height = new_height

    if listbox_side == 'left':
        x = 640
    else:
        x = 40

    y = -30
    top.geometry(f"{window_width}x{window_height}+{x}+{y}")
    image_listbox.bind("<Leave>", lambda event: top.destroy())
    top.bind("<Button-1>", lambda event: top.destroy())

def browse_images():
    global image_directory
    image_directory = browse_for_directory()
    update_status("Selected image directory: " + image_directory)
    load_image_list()

def train_data():
    global recognizer, image_directory
    if not image_directory:
        update_status("Please select an image directory first.")
        return

    training_data_path = os.path.join(current_directory, "training_data", "training_data.txt")
    if not os.path.exists(training_data_path):
        create_training_data = messagebox.askyesno("Create Training Data", "Training data not found. Would you like to create a new training_data folder and training_data.txt file?")
        if create_training_data:
            os.makedirs(os.path.join(current_directory, "training_data"), exist_ok=True)
            with open(training_data_path, 'w') as f:
                f.write("")  # Create an empty training_data.txt file
            update_status("Created training_data folder and training_data.txt file.")
        else:
            update_status("Training data not found. Training aborted.")
            return

    recognizer = train_recognizer(image_directory, subdirs, label_mapping)
    load_nearest_images()

def update_status(message):
    status_label.config(text=message)

def load_image_list():
    global image_directory
    image_listbox.delete(0, 'end')
    for filename in os.listdir(image_directory):
        if '.' in filename:
            image_listbox.insert('end', filename)

def load_nearest_images():
    global recognizer, image_directory
    if not recognizer:
        update_status("Please load the training data first.")
        return
    image_files = [f for f in os.listdir(image_directory) if f.endswith('.jpg') or f.endswith('.png')]
    distances = []
    for image_file in image_files:
        image_path = os.path.join(image_directory, image_file)
        img = Image.open(image_path).convert('L')
        img = ImageOps.fit(img, (100, 100), Image.ANTIALIAS)
        img_data = list(img.getdata())
        min_diff = float('inf')
        for image in recognizer.values():
            diff = euclidean_distance(img_data, image)
            if diff < min_diff:
                min_diff = diff
        distances.append((min_diff, image_file))
    distances.sort()
    nearest_images = [image_file for _, image_file in distances[:2500]]
    nearest_listbox.delete(0, 'end')
    for filename in nearest_images:
        nearest_listbox.insert('end', filename)

def on_image_select(event):
    selected_image = image_listbox.get(image_listbox.curselection())
    image_path = os.path.join(image_directory, selected_image)
    show_image(image_path, image_listbox, 'left')

def on_nearest_select(event):
    selected_image = nearest_listbox.get(nearest_listbox.curselection())
    image_path = os.path.join(image_directory, selected_image)
    show_image(image_path, nearest_listbox, 'right')

root = Tk()
root.title("Face Recognition")
root.geometry("1200x640+10+10")

frame = Frame(root)
frame.pack(side='top', fill='x', padx=10, pady=10)

browse_button = Button(frame, text="Select Image Directory", command=browse_images)
browse_button.pack(side='left', padx=5)

load_button = Button(frame, text="Load Training Data", command=browse_for_file)
load_button.pack(side='left', padx=5)

train_button = Button(frame, text="Train Data", command=train_data)
train_button.pack(side='left', padx=5)

status_label = Label(frame, text="Status: Waiting for input")
status_label.pack(side='left', padx=5)

image_frame = Frame(root)
image_frame.pack(side='left', fill='y', padx=10, pady=10)

image_listbox = Listbox(image_frame, width=80)
image_listbox.pack(side='left', fill='y')
image_listbox_scrollbar = Scrollbar(image_frame, orient=VERTICAL, command=image_listbox.yview)
image_listbox_scrollbar.pack(side='right', fill='y')
image_listbox.config(yscrollcommand=image_listbox_scrollbar.set)

nearest_frame = Frame(root)
nearest_frame.pack(side='right', fill='y', padx=10, pady=10)

nearest_listbox = Listbox(nearest_frame, width=80)
nearest_listbox.pack(side='left', fill='y')
nearest_listbox_scrollbar = Scrollbar(nearest_frame, orient=VERTICAL, command=nearest_listbox.yview)
nearest_listbox_scrollbar.pack(side='right', fill='y')
nearest_listbox.config(yscrollcommand=nearest_listbox_scrollbar.set)

image_listbox.bind('<<ListboxSelect>>', on_image_select)
nearest_listbox.bind('<<ListboxSelect>>', on_nearest_select)

current_directory = os.path.dirname(os.path.abspath(__file__))
image_directory = ""
training_data_file = ""
recognizer = None
subdirs = []
label_mapping = {}
top = None

root.mainloop()