import gradio as gr
import torch
from torchvision import transforms
from PIL import Image
import timm
import numpy as np
import os
import requests
import json

# DeepSeek AI Class
class DeepSeekAI:
    def __init__(self, api_key=None):
        self.api_key = api_key or "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"
        self.base_url = "https://api.deepseek.com/v1/chat/completions"
        self.headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}"
        }
    
    def get_medical_recommendations(self, diagnosis, confidence, is_tb):
        """Get AI medical recommendations from DeepSeek - NO DEMO FALLBACK"""
        try:
            # Only use real API, no demo fallback
            if not self.api_key or self.api_key.startswith("sk-or-v1-"):
                return None  # Return None instead of demo
            
            prompt = f"Hasil analisis X-ray: Diagnosis: {diagnosis}, Confidence: {confidence:.1f}%. Berikan 5 rekomendasi praktis dalam Bahasa Indonesia."
            
            data = {
                "model": "deepseek-chat",
                "messages": [
                    {"role": "system", "content": "Anda adalah asisten medis AI."},
                    {"role": "user", "content": prompt}
                ],
                "max_tokens": 300,
                "temperature": 0.7
            }
            
            response = requests.post(self.base_url, headers=self.headers, json=data, timeout=10)
            
            if response.status_code == 200:
                result = response.json()
                return result["choices"][0]["message"]["content"]
            else:
                return None  # Return None on API error
                
        except Exception as e:
            print(f"DeepSeek API error: {e}")
            return None  # Return None on any error

deepseek = DeepSeekAI()

def load_tb_detection_model():
    """Load the trained TB detection model"""
    try:
        package = torch.load("tb_detection_complete_model.pth", map_location="cpu")
        model = timm.create_model(
            package["model_architecture"],
            pretrained=False,
            num_classes=len(package["class_names"])
        )
        model.load_state_dict(package["model_state_dict"])
        model.eval()
        return model, package
    except:
        model_weights = torch.load("optimal_clinical_model.pth", map_location="cpu")
        model = timm.create_model("inception_resnet_v2", pretrained=False, num_classes=2)
        model.load_state_dict(model_weights)
        model.eval()
        package = {
            "model_architecture": "inception_resnet_v2",
            "class_names": ["Normal", "Tuberculosis"],
            "peak_accuracy": 98.84,
            "image_size": 299,
            "normalization_mean": [0.5, 0.5, 0.5],
            "normalization_std": [0.5, 0.5, 0.5]
        }
        return model, package

model, config = load_tb_detection_model()
class_names = config["class_names"]
peak_accuracy = config["peak_accuracy"]

TEMPLATE_RECOMMENDATIONS = {
    "Tuberculosis": [
        "🚑 Segera konsultasi ke dokter spesialis paru dalam 48 jam",
        "🔬 Lakukan tes dahak (sputum test) untuk konfirmasi diagnosis",
        "🏠 Lakukan isolasi diri di ruangan terpisah untuk mencegah penularan",
        "💊 Konsultasi program pengobatan DOTS (Directly Observed Treatment)",
        "📊 Monitor gejala: batuk >2 minggu, demam, berkeringat malam, berat badan turun",
        "👨‍👩‍👧‍👦 Sarankan tes TB untuk anggota keluarga dan kontak dekat"
    ],
    "Normal": [
        "✅ Pertahankan gaya hidup sehat dan bersih",
        "💪 Lakukan olahraga rutin 3-5 kali seminggu untuk kesehatan paru",
        "🥦 Konsumsi makanan bergizi seimbang dengan cukup protein dan vitamin",
        "🚭 Hindari rokok dan paparan polusi udara",
        "😷 Gunakan masker di tempat umum yang ramai",
        "📅 Lakukan pemeriksaan kesehatan rutin minimal setahun sekali"
    ]
}

transform = transforms.Compose([
    transforms.Resize((config["image_size"], config["image_size"])),
    transforms.ToTensor(),
    transforms.Normalize(config["normalization_mean"], config["normalization_std"])
])

def diagnose_tuberculosis(image):
    image = Image.fromarray(image).convert("RGB") if isinstance(image, np.ndarray) else image.convert("RGB")
    input_tensor = transform(image).unsqueeze(0)

    model.eval()
    with torch.no_grad():
        output = model(input_tensor)
        probabilities = torch.nn.functional.softmax(output[0], dim=0)
        predicted_class = torch.argmax(probabilities).item()
        confidence = probabilities[predicted_class].item()

    diagnosis = class_names[predicted_class]
    is_tb = diagnosis == "Tuberculosis"
    return diagnosis, confidence, is_tb

def predict_tuberculosis(image):
    try:
        diagnosis, confidence, is_tb = diagnose_tuberculosis(image)
        template_recs = TEMPLATE_RECOMMENDATIONS.get(diagnosis, [])[:6]
        
        # Get AI recommendations - returns None if not available
        ai_advice = deepseek.get_medical_recommendations(diagnosis, confidence * 100, is_tb)
        
        color = "#e74c3c" if is_tb else "#27ae60"
        icon = "⚠️" if is_tb else "✅"
        
        result_html = f"""
        <div style="padding: 20px; border: 2px solid {color}; border-radius: 10px; background: {color}10;">
            <h3 style="color: {color}; margin-top: 0;">
                {icon} Hasil Deteksi Tuberculosis
            </h3>
            <p><b>Diagnosis:</b> <span style="color: {color}; font-weight: bold;">{diagnosis}</span></p>
            <p><b>Tingkat Kepercayaan:</b> <span style="color: {color}; font-weight: bold;">{confidence:.1%}</span></p>
        """

        # Template recommendations
        result_html += f"""
            <div style="margin: 15px 0; padding: 15px; background: {color}05; border-radius: 8px;">
                <h4>📋 Rekomendasi Template Medis:</h4>
                <ul>
                    {''.join([f'<li>{action}</li>' for action in template_recs])}
                </ul>
            </div>
        """

        # AI recommendations - show "Not Available" if None
        result_html += f"""
            <div style="margin: 15px 0; padding: 15px; background: #f0f0f0; border-radius: 8px;">
                <h4>🤖 Rekomendasi AI:</h4>
        """
        
        if ai_advice:
            result_html += f"<p>{ai_advice.replace(chr(10), '<br>')}</p>"
        else:
            result_html += """
                <p style="color: #666; font-style: italic;">
                    ⚠️ AI recommendations not available at the moment. 
                    This is a known limitation of the current implementation.
                </p>
            """
        
        result_html += """
            </div>
        """

        # Disclaimer dengan teks hitam
        result_html += """
            <div style="margin-top: 15px; padding: 10px; background: #f8f9fa; border-radius: 5px; font-size: 12px; color: #000000;">
                ⚠️ <b>Disclaimer Medis:</b> Hasil ini merupakan bantuan diagnosis awal.
                Konsultasi dengan tenaga medis profesional tetap diperlukan untuk diagnosis dan pengobatan yang akurat.
            </div>
        </div>
        """

        return result_html

    except Exception as e:
        return f"<div style='color: red; padding: 20px; border: 2px solid red;'>Error: {str(e)}</div>"

# Gradio Interface
demo = gr.Interface(
    fn=predict_tuberculosis,
    inputs=gr.Image(type="pil", label="📁 Upload Chest X-Ray Image"),
    outputs=gr.HTML(label="🎯 Detection Result"),
    title="🩺 Tuberculosis Detection from Chest X-Ray",
    description=f"AI-powered tuberculosis detection with **{peak_accuracy:.1f}% accuracy**. Upload a chest X-ray image for instant analysis.",
    examples=None
)

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