import streamlit as st
from PIL import Image
import torch
from torchvision import transforms
from facenet_pytorch import InceptionResnetV1
import numpy as np
import json
import pickle
from sklearn.metrics.pairwise import cosine_similarity
from mtcnn import MTCNN

st.set_page_config(page_title="FaceTwin AI", page_icon="🤖", layout="centered")

st.title("FaceTwin AI 🤖")
st.subheader("Find your closest celebrity match through AI-powered facial similarity 📸 ⭐")
st.write("---")

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


@st.cache_resource(show_spinner="Loading VGGFace2 model...")
def load_vgg_model():
    model = InceptionResnetV1(pretrained='vggface2').eval().to(device)
    return model

@st.cache_resource(show_spinner="Initializing face detector...")
def load_face_detector():
    return MTCNN()

@st.cache_data(show_spinner="Loading embeddings and metadata...")
def load_feature_data():
    features_list = np.load("features.npy")
    labels_list = np.load("labels.npy")
    with open("image_paths.pkl", "rb") as f:
        imgs_path = pickle.load(f)
    with open("label_dict.json", "r") as f:
        label_dict = json.load(f)
    return features_list, labels_list, imgs_path, label_dict

vgg_feature_extractor = load_vgg_model()
face_detector = load_face_detector()
features_list, labels_list, imgs_path, label_dict = load_feature_data()



transformer = transforms.Compose([
    transforms.Resize((160, 160)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5, 0.5, 0.5],
                         std=[0.5, 0.5, 0.5])
])




def predict_actor(img):
    img = np.array(img)
    results = face_detector.detect_faces(img)

    if len(results) == 0:
        return None, None, "😕 No face detected. Please try another image with a clear front-facing face."

    x, y, width, height = results[0]['box']
    face = img[y:y + height, x:x + width]
    face = Image.fromarray(face).convert("RGB")

    img_tensor = transformer(face).unsqueeze(0).to(device)

    with torch.no_grad():
        features = vgg_feature_extractor(img_tensor)

    flattened_features = features.flatten().cpu().numpy().astype('float16')

    similarities = cosine_similarity([flattened_features], features_list)
    best_index = np.argmax(similarities)

    label = str(labels_list[best_index])
    name = label_dict[label]
    img_path = imgs_path[best_index]

    return name, img_path, None


# ===============================
# 🎛️ STREAMLIT UI
# ===============================
upload_img = st.file_uploader("Upload your image", type=["jpg", "jpeg", "png"])

if st.button("Find my Celebrity Match ⭐"):
    if upload_img is not None:
        img = Image.open(upload_img)

        with st.spinner("Analyzing your face..."):
            name, img2_path, error_msg = predict_actor(img)

        if error_msg:
            st.warning(error_msg)
        else:
            img2 = Image.open(f"celeb_pictures/{name}.jpg")
            img3 = Image.open(img2_path)

            TARGET_SIZE = (300, 350)
            img = img.resize(TARGET_SIZE)
            img2 = img2.resize(TARGET_SIZE)
            img3 = img3.resize(TARGET_SIZE)

            col1, col2, col3 = st.columns(3)
            with col1:
                st.image(img, caption="You", use_container_width=True)
            with col2:
                st.image(img2, caption=f"Your Celebrity Twin: {name}", use_container_width=True)
            with col3:
                st.image(img3, caption="Exact Look", use_container_width=True)

            st.success(f"🎉 You look like **{name}**!")
    else:
        st.warning("⚠️ Please upload an image first.")