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

dataset = load_dataset("SaladSlayer00/twin_matcher")

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

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.environ.get('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:
        
        predictions = image_classifier(input_image)
        top_prediction = max(predictions, key=lambda x: x['score'])
        label, score = top_prediction['label'], top_prediction['score']

        
        images = fetch_images_for_label(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)


iface = gr.Interface(
    fn=predict_and_fetch_images,
    inputs=gr.Image(type="pil", label="Upload or Take a Snapshot"),
    outputs=[
        "text",  
        "number",  
        gr.Gallery(label="Lookalike Images"),  
        "html",  
        gr.Textbox(type="text", label="Feedback", placeholder="Provide feedback here")  # Feedback textbox
    ],
    live=True,
    title="Celebrity Lookalike Predictor",
    description="Take a snapshot or upload an image to see which celebrity you look like!"
)


iface.launch()