import gradio as gr
import tensorflow as tf
import numpy as np
from PIL import Image
import json
from huggingface_hub import hf_hub_download

# ============================================
# Configurações do modelo
# ============================================
REPO_ID = "raysarocha/plant-resnet50-38classes"  # Seu modelo recém-deployado!

# ============================================
# Carregamento do modelo
# ============================================
def load_model():
    """Carrega o modelo uma única vez e mantém em cache"""
    print("🔄 Baixando e carregando modelo...")
    
    # Baixa os arquivos
    cfg_path = hf_hub_download(REPO_ID, "config.json")
    model_path = hf_hub_download(REPO_ID, "model.keras")
    
    # Carrega configuração
    with open(cfg_path, "r") as f:
        cfg = json.load(f)
    
    # Carrega modelo
    model = tf.keras.models.load_model(model_path)
    print("✅ Modelo carregado com sucesso!")
    
    return model, cfg

# Carrega o modelo
model, cfg = load_model()

# ============================================
# Dicionário de nomes amigáveis e recomendações
# ============================================
PLANT_INFO = {
    "Apple___Apple_scab": {
        "name": "🍎 Maçã - Sarna",
        "healthy": False,
        "description": "Doença fúngica que causa manchas escuras nas folhas e frutos",
        "treatment": "• Remova folhas infectadas\n• Aplique fungicida à base de cobre\n• Melhore a circulação de ar"
    },
    "Apple___Black_rot": {
        "name": "🍎 Maçã - Podridão Negra",
        "healthy": False,
        "description": "Infecção fúngica que causa podridão em frutos e folhas",
        "treatment": "• Pode galhos infectados\n• Remova frutos mumificados\n• Use fungicida preventivo"
    },
    "Apple___Cedar_apple_rust": {
        "name": "🍎 Maçã - Ferrugem do Cedro",
        "healthy": False,
        "description": "Doença que causa manchas alaranjadas nas folhas",
        "treatment": "• Remova cedros próximos se possível\n• Aplique fungicida na primavera\n• Escolha variedades resistentes"
    },
    "Apple___healthy": {
        "name": "🍎 Maçã - Saudável",
        "healthy": True,
        "description": "Planta em bom estado de saúde",
        "treatment": "• Continue com os cuidados regulares\n• Mantenha boa irrigação\n• Faça podas preventivas"
    },
    "Tomato___Bacterial_spot": {
        "name": "🍅 Tomate - Mancha Bacteriana",
        "healthy": False,
        "description": "Infecção bacteriana que causa manchas nas folhas",
        "treatment": "• Evite molhar as folhas\n• Use sementes certificadas\n• Aplique cobre preventivamente"
    },
    "Tomato___Early_blight": {
        "name": "🍅 Tomate - Requeima Precoce",
        "healthy": False,
        "description": "Doença fúngica com manchas em anéis concêntricos",
        "treatment": "• Remova folhas afetadas\n• Aplique fungicida\n• Faça rotação de culturas"
    },
    "Tomato___healthy": {
        "name": "🍅 Tomate - Saudável",
        "healthy": True,
        "description": "Planta em bom estado de saúde",
        "treatment": "• Mantenha irrigação regular\n• Adube adequadamente\n• Monitore pragas"
    },
    "Potato___Early_blight": {
        "name": "🥔 Batata - Requeima Precoce",
        "healthy": False,
        "description": "Manchas escuras com anéis concêntricos nas folhas",
        "treatment": "• Use fungicida preventivo\n• Evite irrigação excessiva\n• Destrua restos culturais"
    },
    "Potato___Late_blight": {
        "name": "🥔 Batata - Requeima Tardia",
        "healthy": False,
        "description": "Doença devastadora que causa manchas aquosas",
        "treatment": "• Aplique fungicida imediatamente\n• Remova plantas infectadas\n• Melhore drenagem do solo"
    },
    "Potato___healthy": {
        "name": "🥔 Batata - Saudável",
        "healthy": True,
        "description": "Planta em bom estado de saúde",
        "treatment": "• Continue monitorando\n• Faça amontoa regular\n• Controle irrigação"
    }
}

# Preenche informações padrão para classes não detalhadas
for class_key in cfg["label2id"].keys():
    if class_key not in PLANT_INFO:
        if "healthy" in class_key.lower():
            PLANT_INFO[class_key] = {
                "name": class_key.replace("___", " - ").replace("_", " "),
                "healthy": True,
                "description": "Planta aparentemente saudável",
                "treatment": "• Mantenha os cuidados regulares\n• Continue monitorando"
            }
        else:
            PLANT_INFO[class_key] = {
                "name": class_key.replace("___", " - ").replace("_", " "),
                "healthy": False,
                "description": "Possível problema detectado na planta",
                "treatment": "• Consulte um agrônomo\n• Isole a planta se possível\n• Evite excesso de umidade"
            }

# ============================================
# Função de classificação
# ============================================
def classify_plant(image):
    """Classifica a imagem e retorna diagnóstico detalhado"""
    
    if image is None:
        return None, "", ""
    
    # Preprocessamento
    img = Image.fromarray(image).convert("RGB")
    img = img.resize((cfg["image_size"], cfg["image_size"]))
    arr = np.array(img).astype("float32") * cfg["rescale"]
    arr = np.expand_dims(arr, axis=0)
    
    # Predição
    predictions = model(arr)
    probs = tf.nn.softmax(predictions[0]).numpy()
    
    # Top 5 resultados
    top_5_idx = np.argsort(probs)[-5:][::-1]
    
    # Formata resultados
    results = {}
    for idx in top_5_idx:
        class_name = cfg["id2label"][str(idx)]
        info = PLANT_INFO.get(class_name, {})
        friendly_name = info.get("name", class_name)
        confidence = float(probs[idx])
        results[friendly_name] = confidence
    
    # Pega informação da classe mais provável
    top_class = cfg["id2label"][str(top_5_idx[0])]
    top_confidence = float(probs[top_5_idx[0]])
    top_info = PLANT_INFO.get(top_class, {})
    
    # Cria card de diagnóstico
    if top_info.get("healthy", False):
        status_icon = "✅"
        status_text = "Planta Saudável"
        card_color = "#10b981"  # Verde
        bg_color = "#f0fdf4"
    else:
        status_icon = "⚠️"
        status_text = "Atenção Necessária"
        card_color = "#f59e0b"  # Laranja
        bg_color = "#fef3c7"
    
    diagnosis_html = f"""
    <div style='padding: 20px; border-radius: 12px; background: {bg_color}; border: 2px solid {card_color}; margin-bottom: 20px;'>
        <h2 style='color: {card_color}; margin: 0 0 10px 0; font-size: 24px;'>
            {status_icon} {status_text}
        </h2>
        <h3 style='color: #374151; margin: 10px 0;'>{top_info.get('name', top_class)}</h3>
        <p style='color: #6b7280; margin: 10px 0;'>Confiança: {top_confidence:.1%}</p>
        <hr style='border: 1px solid {card_color}; opacity: 0.3; margin: 15px 0;'>
        <p style='color: #4b5563; margin: 10px 0;'><strong>Descrição:</strong><br>{top_info.get('description', 'Informação não disponível')}</p>
    </div>
    """
    
    # Recomendações
    treatment_html = f"""
    <div style='padding: 15px; border-radius: 8px; background: #f9fafb; border: 1px solid #e5e7eb;'>
        <h3 style='color: #1f2937; margin: 0 0 10px 0;'>📋 Recomendações:</h3>
        <pre style='color: #4b5563; margin: 0; white-space: pre-wrap; font-family: sans-serif; line-height: 1.5;'>{top_info.get('treatment', '• Consulte um especialista para orientações específicas')}</pre>
    </div>
    """
    
    return results, diagnosis_html, treatment_html

# ============================================
# Interface Gradio
# ============================================
with gr.Blocks(
    title="Plant Identifier & Health Checker 🌿",
    theme=gr.themes.Soft(
        primary_hue="green",
        secondary_hue="emerald",
    )
) as demo:
    
    gr.Markdown("""
    # 🌿 Plant Identifier & Health Checker
    ### Identificação e Análise de Saúde de Plantas
    
    Upload uma foto da sua planta para receber um diagnóstico instantâneo e recomendações de tratamento.
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            input_image = gr.Image(
                label="📷 Upload da Imagem",
                type="numpy",
                height=400,
                sources=["upload", "webcam", "clipboard"]
            )
            
            with gr.Row():
                clear_btn = gr.Button("🗑️ Limpar", variant="secondary")
                submit_btn = gr.Button("🔍 Analisar", variant="primary", scale=2)
            
            gr.Markdown("""
            **💡 Dicas para melhores resultados:**
            - Use boa iluminação natural
            - Foque nas folhas afetadas
            - Evite sombras fortes
            - Aproxime-se da área problemática
            """)
            
        with gr.Column(scale=1):
            diagnosis_output = gr.HTML(
                label="Diagnóstico",
                value="<div style='padding: 40px; text-align: center; color: #9ca3af;'><h3>🔍 Aguardando imagem...</h3><p>Faça upload de uma foto para começar</p></div>"
            )
            treatment_output = gr.HTML(label="Tratamento Recomendado")
            confidence_output = gr.Label(
                label="📊 Análise de Probabilidade",
                num_top_classes=5
            )
    
    # # Exemplos
    # gr.Markdown("### 🖼️ Exemplos")
    # gr.Examples(
    #     examples=[
    #         ["https://www.gardeningknowhow.com/wp-content/uploads/2017/06/apple-scab-lesions.jpg"],
    #         ["https://www.gardeningknowhow.com/wp-content/uploads/2020/11/tomato-early-blight.jpg"],
    #         ["https://www.gardeningknowhow.com/wp-content/uploads/2019/08/potato-blight-400x300.jpg"],
    #     ],
    #     inputs=input_image,
    #     outputs=[confidence_output, diagnosis_output, treatment_output],
    #     fn=classify_plant,
    #     cache_examples=True
    # )
    
    gr.Markdown("""
    ---
    ### ℹ️ Sobre o Modelo
    
    - **Arquitetura:** ResNet50
    - **Classes:** 38 tipos de plantas e doenças
    - **Plantas incluídas:** Maçã, Tomate, Batata, Uva, Milho, Pêssego, Pimentão, Morango e mais
    - **Precisão:** ~98% no conjunto de validação
    
    ### 🔗 Links Úteis
    - [Modelo no HuggingFace](https://huggingface.co/raysarocha/plant-resnet50-38classes)
    - [Relatório Técnico](#)
    - [Dataset Original](https://www.kaggle.com/datasets/vipoooool/new-plant-diseases-dataset)
    ---
    <p style='text-align: center; color: #9ca3af;'>
        Desenvolvido com ❤️ por Raysa Rocha | 2024
    </p>
    """)
    
    # Eventos
    submit_btn.click(
        fn=classify_plant,
        inputs=input_image,
        outputs=[confidence_output, diagnosis_output, treatment_output]
    )
    
    clear_btn.click(
        fn=lambda: (None, 
                   "<div style='padding: 40px; text-align: center; color: #9ca3af;'><h3>🔍 Aguardando imagem...</h3></div>",
                   "", 
                   None),
        outputs=[input_image, diagnosis_output, treatment_output, confidence_output]
    )
    
    # Análise automática ao fazer upload
    input_image.change(
        fn=classify_plant,
        inputs=input_image,
        outputs=[confidence_output, diagnosis_output, treatment_output]
    )

# Inicia a aplicação
if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        share=False
    )