import gradio as gr
from transformers import pipeline
from PIL import Image
import torch

# Load the model
try:
    classifier = pipeline("image-classification", model="Parveshiiii/breast-cancer-detector")
except Exception as e:
    print(f"Error loading model: {e}")
    classifier = None

ID_TO_LABEL = {
    "0": "benign",
    "1": "malignant",
    "2": "normal"
}

def predict(img):
    if classifier is None:
        return {"Error": "Model not loaded"}
    
    results = classifier(img)
    output = {}
    for res in results:
        label_id = res['label']
        score = res['score']
        
        if label_id.startswith("LABEL_"):
            clean_id = label_id.replace("LABEL_", "")
            display_label = ID_TO_LABEL.get(clean_id, label_id)
        elif label_id in ID_TO_LABEL:
            display_label = ID_TO_LABEL[label_id]
        else:
            display_label = label_id
            
        output[display_label.capitalize()] = float(score)
    
    return output

# Professional, Tech-Sleek CSS
custom_css = """
body, .gradio-container {
    background-color: #0f172a !important; /* Dark Slate Full Background */
    margin: 0 !important;
    padding: 0 !important;
}
.gradio-container {
    font-family: 'Inter', -apple-system, system-ui, sans-serif !important;
    max-width: none !important;
    width: 100% !important;
    min-height: 100vh !important;
    padding: 40px !important;
    color: #f8fafc !important;
}
.main-header {
    text-align: center;
    padding: 40px 0;
    border-bottom: 1px solid #1e293b;
    margin-bottom: 40px;
}
.main-header h1 {
    font-weight: 700;
    font-size: 2.5rem;
    color: #ffffff !important;
    margin: 0;
}
.main-header p {
    color: #94a3b8 !important;
    font-size: 1.1rem;
    margin-top: 10px;
}
.input-section, .output-section {
    background: #1e293b !important;
    border-radius: 12px !important;
    border: 1px solid #334155 !important;
    padding: 24px !important;
}
button.primary-btn {
    background: #3b82f6 !important;
    color: white !important;
    border: none !important;
    padding: 12px 24px !important;
    font-weight: 600 !important;
    border-radius: 8px !important;
    cursor: pointer !important;
    transition: all 0.3s ease !important;
}
button.primary-btn:hover {
    background: #2563eb !important;
    box-shadow: 0 0 20px rgba(59, 130, 246, 0.4) !important;
}
.label-box {
    color: #ffffff !important;
}
.footer {
    text-align: center;
    margin-top: 60px;
    padding: 20px;
    border-top: 1px solid #1e293b;
}
.footer a {
    color: #3b82f6 !important;
    text-decoration: none !important;
    font-weight: 500;
}
.footer a:hover {
    text-decoration: underline !important;
}
"""

with gr.Blocks(theme=gr.themes.Default(primary_hue="blue", secondary_hue="slate")) as demo:
    with gr.Column(elem_classes="main-header"):
        gr.Markdown("# Image Diagnosis Interface")
        gr.Markdown("Vision Transformer analysis system for high-resolution ultrasound imagery.")
    
    with gr.Row(variant="compact"):
        with gr.Column(scale=1, elem_classes="input-section"):
            image_input = gr.Image(type="pil", label="Source Imagery")
            predict_btn = gr.Button("Execute Analysis", variant="primary", elem_classes="primary-btn")
        
        with gr.Column(scale=1, elem_classes="output-section"):
            label_output = gr.Label(num_top_classes=3, label="Probability Distribution")
            
    gr.Markdown("### Technical Documentation")
    gr.Markdown("""
    The system utilizes a Vision Transformer (ViT) architecture fine-tuned for classification tasks across three primary domains:
    - **Benign**: Identification of stable structural formations.
    - **Malignant**: Detection of invasive cellular development.
    - **Normal**: Verification of standard cellular architecture.
    
    Notice: This interface is for analytical demonstration. Professional review required for outcome validation.
    """)
    
    with gr.Column(elem_classes="footer"):
        gr.Markdown("System Architecture: Hugging Face & Gradio Open-Source Frameworks")
        gr.Markdown('Interested in how this model was made? Read the technical blog: <a href="https://parveshiiii.github.io/blogs/breast-cancer-detector/" target="_blank">Enabling Early Detection: Breast Cancer Detector</a>')

    predict_btn.click(fn=predict, inputs=image_input, outputs=label_output)

if __name__ == "__main__":
    demo.launch(css=custom_css)