import gradio as gr
from transformers import pipeline
from datasets import load_dataset
import requests
import traceback
import os

dataset = load_dataset("SaladSlayer00/twin_matcher")

image_classifier = pipeline("image-classification", model="SaladSlayer00/twin_matcher")

def fetch_info(celebrity_label):
    try:
        parts = celebrity_label.split("_")
        formatted_label = " ".join([part.capitalize() for part in parts])
        api_url = f'https://api.api-ninjas.com/v1/celebrity?name={formatted_label}'
        token = os.getenv('TOKEN')
        response = requests.get(api_url, headers={'X-Api-Key': token})
        if response.status_code == 200:
            return response.text
        else:
            return "Description not available."
    except Exception as e:
        print(f"Error fetching information: {e}")
        traceback.print_exc()
        return "Description not available."

# Function to fetch images for the predicted label from the Hugging Face dataset
def fetch_images_for_label(label):
    label_data = dataset['train'].filter(lambda example: example['label'] == label)
    images = [example['image'] for example in label_data]
    return images

# Function to process the image, predict, and fetch images and info
def predict_and_fetch_images(input_image):
    try:
        # Use the image classifier pipeline
        predictions = image_classifier(input_image)
        top_prediction = max(predictions, key=lambda x: x['score'])
        label, score = top_prediction['label'], top_prediction['score']

        # Fetch images for the predicted label
        images = fetch_images_for_label(label)

        # Fetch information for the predicted label
        info = fetch_info(label)

        return label, score, images, info, "No Error"
    except Exception as e:
        print(f"Error during prediction: {e}")
        traceback.print_exc()
        return "Error during prediction", 0, [], "N/A", str(e)



# Gradio interface
iface = gr.Interface(
    fn=predict_and_fetch_images,
    inputs=gr.Image(type="pil", label="Upload or Take a Snapshot"),
    outputs=[
        "text",  # Predicted label
        "number",  # Prediction score
        gr.Gallery(label="Lookalike Images"),  # Slideshow component for images
        "text",  # Info/Description
        gr.Textbox(type="text", label="Feedback", placeholder="Provide feedback here")  # Feedback textbox
    ],
    live=True,
    title="Celebrity Lookalike Predictor with Description and Feedback",
    description="Take a snapshot or upload an image to see which celebrity you look like!"
)

# Launch the Gradio interface
iface.launch()