import os
import gradio as gr
import torch
from torch import nn
import torchvision.models as models
import torchvision.transforms as transforms
from PIL import Image
from datasets import load_dataset
import numpy as np
from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline

# Veri setini yükleme
dataset = load_dataset("wanglab/chest-agent-bench")

# Model tanımı - önceden eğitilmiş bir DenseNet modeli kullanacağız
class ChestXRayModel(nn.Module):
    def __init__(self, num_classes=14):
        super(ChestXRayModel, self).__init__()
        self.densenet = models.densenet121(pretrained=True)
        num_features = self.densenet.classifier.in_features
        self.densenet.classifier = nn.Sequential(
            nn.Linear(num_features, num_classes),
            nn.Sigmoid()
        )
    
    def forward(self, x):
        return self.densenet(x)

# Model yükleme fonksiyonu (gerçek bir uygulamada önceden eğitilmiş model kullanılacaktır)
def load_model():
    model = ChestXRayModel()
    # Gerçek uygulamada: model.load_state_dict(torch.load('path/to/model.pth'))
    model.eval()
    return model

# Görüntü dönüştürme
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# Sınıf etiketleri (CheXpert'ten alındı)
class_names = [
    'Atelektazi', 'Kardiyomegali', 'Effüzyon', 'İnfiltrasyon', 'Kitle',
    'Nodül', 'Pnömoni', 'Pnömotoraks', 'Konsolidasyon', 'Ödem',
    'Amfizem', 'Fibrozis', 'Plevral Kalınlaşma', 'Hernia'
]

# Rapor oluşturma için doğal dil işleme modeli
def load_nlp_model():
    tokenizer = AutoTokenizer.from_pretrained("facebook/bart-large-cnn")
    summarizer = pipeline("summarization", model="facebook/bart-large-cnn", tokenizer=tokenizer)
    return summarizer

# Röntgen analiz fonksiyonu
def analyze_xray(image):
    # Görüntü ön işleme
    if image is None:
        return None, "Lütfen bir görüntü yükleyin."
    
    img = Image.fromarray(image).convert('RGB')
    img_tensor = transform(img).unsqueeze(0)
    
    # Model tahminleri 
    model = load_model()
    with torch.no_grad():
        outputs = model(img_tensor)
    
    # Tahminleri işleme
    probabilities = outputs[0].numpy()
    
    # Pozitif bulguların tespiti (0.5 eşik değeri ile)
    findings = []
    for i, prob in enumerate(probabilities):
        if prob > 0.5:
            findings.append(f"{class_names[i]} (%.2f%%)" % (prob * 100))
    
    if not findings:
        findings.append("Belirgin patoloji tespit edilmedi.")
    
    # Görselle ilgili özet rapor oluşturma (gerçek uygulama için geliştirilebilir)
    summary = f"**Röntgen Analiz Raporu**\n\n"
    summary += "**Tespit Edilen Bulgular:**\n"
    for finding in findings:
        summary += f"- {finding}\n"
    
    # Gerçek bir uygulamada NLP tabanlı açıklama eklenir
    summary += "\n**Yorum:**\n"
    if "Belirgin patoloji tespit edilmedi." in findings:
        summary += "Radyografide belirgin patolojik bulgu izlenmedi. Klinik bulgular ışığında değerlendirilmesi önerilir."
    else:
        summary += "Yukarıda belirtilen bulgular saptandı. Klinik bulgularla korele değerlendirilmesi önerilir."
    
    # Doktora öneriler
    summary += "\n\n**Öneriler:**\n"
    if any(x in ' '.join(findings) for x in ['Pnömoni', 'Atelektazi', 'İnfiltrasyon']):
        summary += "- Enfeksiyon belirtileri dikkate alınmalı\n"
        summary += "- Antibiyotik tedavisi değerlendirilmeli\n"
    if 'Nodül' in ' '.join(findings) or 'Kitle' in ' '.join(findings):
        summary += "- İleri görüntüleme (BT) önerilir\n"
        summary += "- Takip görüntüleme planlanmalı\n"
    if 'Effüzyon' in ' '.join(findings) or 'Pnömotoraks' in ' '.join(findings):
        summary += "- Torasentez düşünülebilir\n"
    
    return img, summary

# Veri setinden örnek yükleme fonksiyonu
def load_sample():
    # Rastgele bir örnek seçelim
    random_idx = np.random.randint(0, len(dataset["train"]))
    sample = dataset["train"][random_idx]
    image = sample["image"].convert('RGB')
    image_array = np.array(image)
    return image_array

# Gradio arayüzü
with gr.Blocks(title="Göğüs Röntgeni Analiz Asistanı", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🏥 Göğüs Röntgeni Analiz Asistanı")
    gr.Markdown("Bu araç, göğüs röntgenlerini analiz ederek olası bulguları tespit etmeye yardımcı olur.")
    
    with gr.Row():
        with gr.Column(scale=1):
            image_input = gr.Image(type="numpy", label="Röntgen Görüntüsü")
            with gr.Row():
                analyze_btn = gr.Button("Analiz Et", variant="primary")
                sample_btn = gr.Button("Örnek Yükle")
        
        with gr.Column(scale=1):
            image_output = gr.Image(label="İşaretlenmiş Görüntü")
            text_output = gr.Markdown(label="Rapor")
    
    analyze_btn.click(analyze_xray, inputs=image_input, outputs=[image_output, text_output])
    sample_btn.click(load_sample, inputs=None, outputs=image_input)
    
    gr.Markdown("## Nasıl Kullanılır")
    gr.Markdown("""
    1. Göğüs röntgeni yükleyin veya 'Örnek Yükle' butonuna tıklayın
    2. 'Analiz Et' butonuna tıklayın
    3. Sistem röntgeni analiz edecek ve olası bulguları listeleyecektir
    
    **Not:** Bu uygulama sadece yardımcı bir araçtır ve kesin teşhis için uzman radyolog görüşü gereklidir.
    """)

# Uygulamayı başlat
if __name__ == "__main__":
    demo.launch(share=True)