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

dataset = load_dataset("SaladSlayer00/twin_matcher_data")

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

def format_info(info_json):
    try:
        
        info_data = json.loads(info_json)

        
        formatted_info = "<table style='border-collapse: collapse; width: 80%; margin: 20px;'>"
        
        formatted_info += "<tr style='background-color: #f2f2f2;'>"
        for key in info_data[0].keys():
            formatted_info += f"<th style='border: 1px solid #dddddd; text-align: left; padding: 8px;'><b>{key.capitalize()}</b></th>"
        formatted_info += "</tr>"

        
        for entry in info_data:
            formatted_info += "<tr>"
            for value in entry.values():
                formatted_info += f"<td style='border: 1px solid #dddddd; text-align: left; padding: 8px;'>{value}</td>"
            formatted_info += "</tr>"
        formatted_info += "</table>"
        return formatted_info
    except Exception as e:
        print(f"Error formatting info: {e}")
        return "Info not available."


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 format_info(response.text)
        else:
            return "Description not available."
    except Exception as e:
        print(f"Error fetching information: {e}")
        traceback.print_exc()
        return "Description not available."


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


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)


example_images = [
    "images/megan_fox.png",
    "images/chris_evans.png",
    "images/millie_bobby_brown.png",
    "images/alvaro_morte.png",
    "images/amber_heard.png"
]
# 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
        "html",  # Info/Description as HTML
        gr.Textbox(type="text", label="Feedback", placeholder="Provide feedback here")  # Feedback textbox
    ],
    examples=example_images,
    live=True,
    title="Celebrity Lookalike Predictor",
    description="Take a snapshot or upload an image to see which celebrity you look like!"
)


iface.launch()