import gradio as gr
import torch
from transformers import ViTForImageClassification, ViTImageProcessor
from PIL import Image
import numpy as np

# Senin model'inin sınıfları (config'den)
class_names = {
    0: "Enfeksiyonel",
    1: "Ekzama",
    2: "Akne", 
    3: "Pigment",
    4: "Benign",
    5: "Malign",
    6: "Acne",
    7: "Actinic Keratosis",
    8: "Basal Cell Carcinoma",
    9: "Benign Keratosis",
    10: "Dermatofibroma",
    11: "Melanocytic Nevus",
    12: "Melanoma",
    13: "Vascular Lesion",
    14: "Warts/Molluscum"
}

# Risk assessment mapping
risk_categories = {
    "high_risk": [5, 12, 8],  # Malign, Melanoma, Basal Cell Carcinoma
    "medium_risk": [7, 0, 13],  # Actinic Keratosis, Enfeksiyonel, Vascular
    "low_risk": [1, 2, 3, 4, 6, 9, 10, 11, 14]  # Diğerleri
}

def load_model():
    """Senin trained ViT modelini yükle"""
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    # Model ve processor yükle (local files'dan)
    model = ViTForImageClassification.from_pretrained(
        "./", 
        local_files_only=True,
        num_labels=15,
        ignore_mismatched_sizes=True)
    processor = ViTImageProcessor.from_pretrained("VitModel")
    
    model.to(device)
    model.eval()
    
    return model, processor, device

def predict_skin_condition(image):
    """Ana prediction fonksiyonu"""
    
    if image is None:
        return {}, "Lütfen bir görüntü yükleyin"
    
    try:
        # Model yükle
        model, processor, device = load_model()
        
        # Image preprocessing
        inputs = processor(images=image, return_tensors="pt")
        inputs = {k: v.to(device) for k, v in inputs.items()}
        
        # Prediction
        with torch.no_grad():
            outputs = model(**inputs)
            logits = outputs.logits
            probabilities = torch.softmax(logits, dim=-1)[0]
        
        # Results dictionary oluştur
        results = {}
        for idx, prob in enumerate(probabilities):
            if idx in class_names:
                results[class_names[idx]] = float(prob)
        
        # En yüksek tahmin
        top_pred_idx = torch.argmax(probabilities).item()
        top_class = class_names[top_pred_idx]
        confidence = float(probabilities[top_pred_idx])
        
        # Risk assessment
        risk_html = get_risk_assessment(top_pred_idx, confidence, top_class)
        
        return results, risk_html
        
    except Exception as e:
        return {}, f"Hata oluştu: {str(e)}"

def get_risk_assessment(pred_idx, confidence, class_name):
    """Risk değerlendirmesi yap"""
    
    if pred_idx in risk_categories["high_risk"] and confidence > 0.7:
        risk_level = "🚨 YÜKSEK RİSK"
        message = f"<strong>{class_name}</strong> tespit edildi. Dermatolog konsültasyonu ÖNERİLİR."
        color = "#FF4444"
    elif pred_idx in risk_categories["medium_risk"] and confidence > 0.5:
        risk_level = "⚠️ ORTA RİSK"
        message = f"<strong>{class_name}</strong> tespit edildi. Takip önerilir."
        color = "#FF8800"
    else:
        risk_level = "✅ DÜŞÜK RİSK"
        message = f"<strong>{class_name}</strong> tespit edildi. Rutin kontrol yeterli."
        color = "#00AA44"
    
    return f"""
    <div style='padding: 15px; background-color: {color}; color: white; border-radius: 10px; text-align: center; margin: 10px 0;'>
        <h3 style='margin: 0 0 10px 0;'>{risk_level}</h3>
        <p style='margin: 0; font-size: 16px;'>{message}</p>
        <p style='margin: 5px 0 0 0; font-size: 14px;'>Güven Skoru: {confidence:.1%}</p>
    </div>
    """

# CSS styling
css = """
.gradio-container {
    max-width: 1000px;
    margin: 0 auto;
    font-family: 'Segoe UI', sans-serif;
}
.title {
    text-align: center;
    color: #2E86AB;
    margin-bottom: 20px;
}
"""

# Gradio interface
with gr.Blocks(css=css, title="ViT Cilt Hastalığı Analizi") as demo:
    
    gr.Markdown("""
    # 🔬 ViT Cilt Hastalığı AI Analizi
    ### 15 farklı cilt hastalığını tespit eden yapay zeka sistemi
    **Accuracy: %97 | Model: Vision Transformer | Classes: 15**
    
    ---
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            input_image = gr.Image(
                label="📸 Cilt Görüntüsü Yükleyin",
                type="pil",
                height=400
            )
            
            predict_btn = gr.Button(
                "🔍 AI Analizi Başlat",
                variant="primary", 
                size="lg"
            )
            
        with gr.Column(scale=1):
            prediction_output = gr.Label(
                label="📊 Tahmin Sonuçları",
                num_top_classes=5
            )
            
            risk_output = gr.HTML(
                label="⚠️ Risk Değerlendirmesi"
            )
    
    # Sınıf açıklamaları
    gr.Markdown("""
    ### 📋 Tespit Edilen Hastalık Kategorileri:
    
    **Yüksek Risk:** Malign, Melanoma, Basal Cell Carcinoma  
    **Orta Risk:** Actinic Keratosis, Enfeksiyonel, Vascular Lesion  
    **Düşük Risk:** Ekzama, Akne, Pigment, Benign, Dermatofibroma, vb.
    
    ⚠️ **Önemli:** Bu sistem sadece ön değerlendirme içindir. Kesin tanı için doktora başvurun.
    """)
    
    # Event handling
    predict_btn.click(
        fn=predict_skin_condition,
        inputs=input_image,
        outputs=[prediction_output, risk_output]
    )

if __name__ == "__main__":
    demo.launch(
        share=True,
        server_name="0.0.0.0", 
        server_port=7860
    )