import PIL.Image
import gradio as gr
import numpy as np
import random
import PIL
import subprocess
import sys
import os

# os.system("ls -a")
# print("1")
# os.system("pip install ./dlib-19.24.99-cp36-cp36m-linux_x86_64.whl")
# print("2")
# os.system("pip install dlib-19.24.99-cp36-cp36m-linux_x86_64.whl")
# os.system("pip install face-recognition")
# import face_detection

# import spaces #[uncomment to use ZeroGPU]
from diffusers import DiffusionPipeline
import torch

import face_detection

# # Function to display the uploaded image
# def process_image(image : PIL.Image.Image):
#     outputs = face_detection.getCroppedImages(image)
#     # do AI stuff here
#     return gr.Image(outputs[0])

model_repo_id = "CSSE416-final-project/faceRecogModel"
weight_file_id = "matTemp.bin"


# 1. Load the model from Hugging Face Hub
def load_model(repo_id):
    # Download the model weights from the repo
    weights_path = hf_hub_download(repo_id=model_repo_id, filename=weight_file_id)

    # Initialize the ResNet-18 architecture
    model = models.resnet18(pretrained=True) # TODO: does it matter if this is set to true or false?
    num_ftrs = model.fc.in_features
    model.fc = nn.Linear(num_ftrs, 100)  # Adjust for your task (e.g., 128 classes)
    # TODO: check if this number^^ corresponds to the number of classes

    # Load the model weights
    state_dict = torch.load(weights_path, map_location=torch.device("cpu"))
    model.load_state_dict(state_dict)
    model.eval()  # Set the model to evaluation mode
    return model


# 2. Load model
model = load_model(model_repo_id)


# 3. Define how to transform image
transforms = transforms.Compose(
[
    transforms.ToTensor()
])


# 4. Preprocess and display the image
def process_image_str(groupImageFilePath: str):
    groupImage = PIL.Image.open(groupImageFilePath)
    locations, images = face_detection.getCroppedImages(groupImage)

    outputLabels = []
    for image in images:
        # Process the image
        intputTensor = transforms(image).unsqueeze(0) # unsqueeze? add batch dimension??

        # do AI stuff here
        with torch.no_grad():
            outputs_t = model(intputTensor)
            _, pred_t = torch.max(outputs_t, dim=1)
            outputLabels.append(pred_t.item())


    #return gr.Image(image)
    return outputLabels.pop(0)


# 5. Create the Gradio interface
interface = gr.Interface(
    fn=process_image_str,  # Function to process the image
    inputs=gr.Image(type='filepath'),  # Upload input
    # outputs=gr.Image(),  # Display output
    outputs='text',
    allow_flagging='never',
    title="Celebrity Face Detector",
    description="Upload a picture of a celebrity or group of celebrities to identify them"
)

# 6. Launch the app
if __name__ == "__main__":
    interface.launch()