import face_recognition
import numpy as np
import pickle
from mtcnn import MTCNN
from PIL import Image
import cv2
import faiss
import imgaug.augmenters as iaa
import os
import gradio as gr

def detect_and_align_face(image_path):
    detector = MTCNN()
    image = cv2.imread(image_path)
    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    detections = detector.detect_faces(image_rgb)
    
    if len(detections) == 0:
        raise ValueError("No face detected in the image.")
    
    detection = detections[0]
    x, y, width, height = detection['box']
    keypoints = detection['keypoints']
    face = image_rgb[y:y+height, x:x+width]
    
    left_eye = keypoints['left_eye']
    right_eye = keypoints['right_eye']
    delta_x = right_eye[0] - left_eye[0]
    delta_y = right_eye[1] - left_eye[1]
    angle = np.arctan2(delta_y, delta_x) * (180.0 / np.pi)
    
    center = ((x + x + width) // 2, (y + y + height) // 2)
    rot_matrix = cv2.getRotationMatrix2D(center, angle, scale=1.0)
    aligned_image = cv2.warpAffine(image_rgb, rot_matrix, (image_rgb.shape[1], image_rgb.shape[0]))
    aligned_face = aligned_image[y:y+height, x:x+width]
    
    return Image.fromarray(aligned_face)

def load_encodings(file_path):
    with open(file_path, "rb") as file:
        data = pickle.load(file)
    return np.array(data["encodings"]), data["labels"]

def save_encodings(encodings, labels, file_path):
    data = {"encodings": encodings, "labels": labels}
    with open(file_path, "wb") as file:
        pickle.dump(data, file)

def create_faiss_index(known_encodings):
    dimension = known_encodings.shape[1]
    index = faiss.IndexFlatL2(dimension)
    index.add(known_encodings)
    return index

def encode_face(image):
    img_array = np.array(image)
    encodings = face_recognition.face_encodings(img_array)
    return encodings[0] if encodings else None

def augment_image(image, num_augmented=5):
    image = np.array(image)
    aug = iaa.Sequential([
        iaa.Fliplr(0.5),  # horizontal flips
        iaa.Affine(rotate=(-25, 25)),  # rotation
        iaa.AdditiveGaussianNoise(scale=(0, 0.05*255)),  # noise
        iaa.Multiply((0.8, 1.2)),  # brightness
        iaa.GaussianBlur(sigma=(0.0, 1.0))  # blur
    ])
    augmented_images = [Image.fromarray(aug(image=image)) for _ in range(num_augmented)]
    return augmented_images

def update_dataset_with_verified_image(image, encodings_file, label, num_augmented=5):
    known_encodings, known_labels = load_encodings(encodings_file)
    augmented_images = augment_image(image, num_augmented=num_augmented)
    images_to_encode = [image] + augmented_images
    for img in images_to_encode:
        img_array = np.array(img)
        encoding = face_recognition.face_encodings(img_array)[0]
        known_encodings = np.append(known_encodings, [encoding], axis=0)
        known_labels.append(label)
    save_encodings(known_encodings, known_labels, encodings_file)

def verify_face_with_faiss(image, encodings_file, similarity_threshold=70, num_augmented=5):
    aligned_face = image.convert("RGB")
    target_encoding = face_recognition.face_encodings(np.array(aligned_face))[0].reshape(1, -1)
    
    known_encodings, known_labels = load_encodings(encodings_file)
    known_encodings = np.array(known_encodings)
    
    index = create_faiss_index(known_encodings)
    
    distances, indices = index.search(target_encoding, 1)
    
    best_match_index = indices[0][0]
    best_similarity_percentage = (1 - distances[0][0]) * 100
    
    is_match = best_similarity_percentage >= similarity_threshold
    
    if is_match:
        matched_label = known_labels[best_match_index]
        update_dataset_with_verified_image(image, encodings_file, matched_label, num_augmented=num_augmented)
        return True, f"Match found: {matched_label}, Similarity: {best_similarity_percentage:.2f}%"
    else:
        return False, "No match found."

# Define the Gradio interface
def gradio_interface(image, similarity_threshold=70):
    encodings_file = "face_encoding.pkl"
    result, message = verify_face_with_faiss(image, encodings_file, similarity_threshold=similarity_threshold)
    return message

# Launch the Gradio interface
iface = gr.Interface(fn=gradio_interface, 
                     inputs=[gr.Image(type="pil"), gr.Slider(0, 100, value=70, label="Similarity Threshold")], 
                     outputs="text", 
                     title="Face Recognition with MTCNN and FAISS",
                     description="Upload an image to see if it matches any face in the database.")

iface.launch()