import gradio as gr
import torch
from torchvision import transforms
from PIL import Image
import timm
import numpy as np
import os
import google.generativeai as genai

# Load model function
def load_tb_detection_model():
    """Load the trained TB detection model"""
    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

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

# Gemini API Setup - 🔑 GANTI DENGAN API KEY ANDA!
GEMINI_API_KEY = "AIzaSyCAOpTxHISV5J3bZm0yzRWbBKr4fryJ93Y"
try:
    genai.configure(api_key=GEMINI_API_KEY)
    GEMINI_AVAILABLE = True
    print("✅ Gemini API connected")
except:
    GEMINI_AVAILABLE = False
    print("❌ Gemini API not available")

# Template recommendations
TEMPLATE_RECOMMENDATIONS = {
    "Tuberculosis": {
        "urgency": "TINGGI",
        "color": "#ff4444",
        "icon": "⚠️",
        "actions": [
            "🚑 Segera konsultasi dokter spesialis paru",
            "🔬 Lakukan tes dahak & rontgen thorax lengkap",
            "🏠 Isolasi diri untuk mencegah penularan",
            "💊 Konsultasi pengobatan DOTS dengan dokter"
        ]
    },
    "Normal": {
        "urgency": "RENDAH",
        "color": "#44ff44",
        "icon": "✅",
        "actions": [
            "👍 Pertahankan gaya hidup sehat",
            "📅 Lakukan pemeriksaan rutin tahunan",
            "💪 Monitor kesehatan pernapasan berkala"
        ]
    }
}

# Define transforms
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_path, model, transforms, class_names):
    """Perform TB diagnosis on a single image"""
    image = Image.open(image_path).convert("RGB")
    input_tensor = transforms(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]
    return diagnosis, confidence

def get_gemini_medical_advice(diagnosis, confidence):
    """Get AI-powered medical recommendations from Gemini"""
    if not GEMINI_AVAILABLE:
        return None

    try:
        model = genai.GenerativeModel("gemini-pro")

        prompt = f"""
        Sebagai asisten medis AI, berikan rekomendasi singkat (3-4 poin) dalam Bahasa Indonesia untuk:

        Diagnosis: {diagnosis}
        Tingkat Kepercayaan AI: {confidence:.1%}

        Format: poin-poin praktis tanpa disclaimer berulang.
        Contoh format:
        • Poin pertama
        • Poin kedua
        • Poin ketiga

        Fokus pada tindakan praktis yang bisa dilakukan pasien.
        """

        response = model.generate_content(prompt)
        return response.text

    except Exception as e:
        print(f"Gemini API error: {e}")
        return None

def predict_tuberculosis(image):
    """Enhanced prediction with template + AI recommendations"""
    try:
        # Save temporary image
        temp_path = "temp_medical_image.jpg"
        image.save(temp_path)

        # Get basic diagnosis
        diagnosis, confidence = diagnose_tuberculosis(temp_path, model, transform, class_names)

        # Get template recommendations
        template_rec = TEMPLATE_RECOMMENDATIONS.get(diagnosis, {})

        # Get AI recommendations from Gemini
        ai_advice = get_gemini_medical_advice(diagnosis, confidence) if GEMINI_AVAILABLE else None

        # Clean up
        if os.path.exists(temp_path):
            os.remove(temp_path)

        # Build result HTML
        result_html = f"""
        <div style="padding: 20px; border: 2px solid {template_rec.get('color', '#666')}; border-radius: 10px; background: {template_rec.get('color', '#666')}20;">
            <h3 style="color: {template_rec.get('color', '#666')}; margin-top: 0;">
                {template_rec.get('icon', '🔍')} Hasil Deteksi Tuberculosis
            </h3>
            <p><b>Diagnosis:</b> <span style="color: {template_rec.get('color', '#666')}; font-weight: bold;">{diagnosis}</span></p>
            <p><b>Tingkat Kepercayaan:</b> {confidence:.1%}</p>
            <p><b>Status Urgensi:</b> {template_rec.get('urgency', 'TIDAK DIKETAHUI')}</p>
        """

        # Add template recommendations
        result_html += f"""
            <div style="margin: 15px 0; padding: 15px; background: {template_rec.get('color', '#666')}10; border-radius: 8px;">
                <h4>📋 Rekomendasi Standar Medis:</h4>
                <ul>
                    {''.join([f'<li>{action}</li>' for action in template_rec.get('actions', [])])}
                </ul>
            </div>
        """

        # Add AI recommendations if available
        if ai_advice:
            result_html += f"""
            <div style="margin: 15px 0; padding: 15px; background: #ffeb3b20; border-radius: 8px;">
                <h4>🤖 Rekomendasi AI Cerdas:</h4>
                <p>{ai_advice.replace(chr(10), '<br>')}</p>
            </div>
            """
        else:
            result_html += """
            <div style="margin: 15px 0; padding: 10px; background: #ff980020; border-radius: 8px;">
                <small>⚠️ Fitur rekomendasi AI sementara tidak tersedia. Menggunakan rekomendasi standar medis.</small>
            </div>
            """

        # Add disclaimer
        result_html += """
            <div style="margin-top: 15px; padding: 10px; background: #f8f9fa; border-radius: 5px; font-size: 12px;">
                ⚠️ <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>"

# Create 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="AI-powered tuberculosis detection with **{peak_accuracy:.1f}% accuracy**. Upload a chest X-ray image for instant analysis.\n\n⚠️ **Medical Disclaimer:** This tool is for educational purposes only. Always consult healthcare professionals for medical diagnosis.",
    examples=[
        ["sample_Tuberculosis_0.jpg"],
        ["sample_Normal_0.jpg"]
    ] if all(os.path.exists(f) for f in ["sample_Tuberculosis_0.jpg", "sample_Normal_0.jpg"]) else []
)

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