"""
AGRICULTURAL DISEASE DIAGNOSIS CHATBOT - GRADIO VERSION
========================================================

This script implements a Gradio-based chatbot interface for diagnosing crop diseases 
based on user-described symptoms. It loads a pre-trained SetFit model, label mappings, 
and treatment recommendations to provide accurate diagnoses and actionable advice.

"""

import gradio as gr
from setfit import SetFitModel
import pickle
import json
import numpy as np
import torch

# =============================================================================
# LOAD MODEL AND DATA
# =============================================================================

print("Loading model and recommendations...")

# Load SetFit model from HuggingFace
model = SetFitModel.from_pretrained("Tree-Diagram/setfit_crop_disease_model_v1")

# Load label mappings
with open('label_mappings.pkl', 'rb') as f:
    label_info = pickle.load(f)

# Load recommendations
with open('diagnosis_recommendations.json', 'r') as f:
    recommendations = json.load(f)

print(f"✓ Loaded model with {label_info['n_classes']} disease classes")

# =============================================================================
# DIAGNOSIS FUNCTION
# =============================================================================

def diagnose_disease(crop, symptoms):
    """
    Main diagnosis function
    
    Args:
        crop (str): Selected crop (MAIZE, CASSAVA, or TOMATO)
        symptoms (str): Symptom description from user
    
    Returns:
        str: Formatted diagnosis and recommendations in Markdown
    """
    
    # Validate input
    if not symptoms or len(symptoms.strip()) < 10:
        return """
        ⚠️ **Error**: Please provide a detailed symptom description (at least 10 characters).
        
        **Example**: "The plant shows yellowing leaves, bore holes, and visible insects on the leaves."
        """
    
    try:
        # Make prediction
        prediction = model.predict([symptoms])
        
        # Get probabilities for confidence score
        try:
            probs = model.predict_proba([symptoms])
            
            # Convert to numpy if tensor
            if torch.is_tensor(probs):
                probs = probs.cpu().numpy()
            else:
                probs = np.array(probs)
            
            # Get top 3 predictions
            if len(probs.shape) > 1 and probs.shape[1] > 1:
                top_3_indices = np.argsort(-probs[0])[:3]
                top_3_probs = probs[0][top_3_indices]
            else:
                top_3_indices = [int(prediction[0])]
                top_3_probs = [1.0]
        
        except:
            # Fallback if probabilities fail
            top_3_indices = [int(prediction[0])]
            top_3_probs = [1.0]
        
        # Get diagnosis names
        idx_to_diagnosis = label_info['idx_to_diagnosis']
        primary_diagnosis = idx_to_diagnosis[int(top_3_indices[0])]
        confidence = float(top_3_probs[0])
        
        # Format output
        output = f"""
# 🎯 Diagnosis Results

## Primary Diagnosis: **{primary_diagnosis}**

- **Crop**: {crop}
- **Confidence**: {confidence*100:.1f}%
- **Reliability**: {"✅ High" if confidence >= 0.7 else "⚠️ Moderate" if confidence >= 0.5 else "❌ Low - Consult Expert"}

---

"""
        
        # Add alternative diagnoses if confidence is not very high
        if len(top_3_indices) > 1 and confidence < 0.9:
            output += "### 🔄 Alternative Possibilities:\n\n"
            for i in range(1, min(3, len(top_3_indices))):
                alt_diagnosis = idx_to_diagnosis[int(top_3_indices[i])]
                alt_prob = float(top_3_probs[i])
                if alt_prob > 0.1:
                    output += f"{i}. {alt_diagnosis} ({alt_prob*100:.1f}%)\n"
            output += "\n---\n\n"
        
        # Get recommendations
        if primary_diagnosis in recommendations:
            rec = recommendations[primary_diagnosis]
            
            output += f"""
# 💊 Treatment Recommendations

## 🚑 Immediate Action

{rec['current_treatment']}

---

## 🛡️ Prevention for Future

{rec['prevention']}

---

## 📋 Additional Information

- **Treatment Type**: {rec.get('recommendation_type', 'General')}
- **Recommendation Validity**: {rec.get('validity', 'Unknown')}

"""
        else:
            output += """
# ⚠️ No Specific Recommendations Available

**Next Steps:**
1. Contact your local agricultural extension officer
2. Take photos of the affected plant
3. Bring leaf samples for laboratory confirmation
4. Isolate affected plants to prevent spread

"""
        
        # Disclaimer
        output += """
---

## ⚠️ Important Disclaimer

This AI tool provides **preliminary guidance only**. For serious or persistent issues:

- ✅ Consult agricultural extension officers
- ✅ Seek laboratory confirmation
- ✅ Follow local agricultural guidelines
- ✅ Use as one source of information, not the only source

"""
        
        return output
    
    except Exception as e:
        return f"""
# ❌ Error During Diagnosis

An error occurred: {str(e)}

Please try again or contact support if the problem persists.
"""

# =============================================================================
# GRADIO INTERFACE
# =============================================================================

# Example inputs for quick testing
examples = [
    [
        "MAIZE",
        "The plant shows bore holes in leaves, chewed leaves, frass visible, and insects present. Some insects are feeding on the whorl."
    ],
    [
        "CASSAVA",
        "The leaves show yellow and green mosaic patterns, distorted new growth, and the plant is stunted compared to others."
    ],
    [
        "TOMATO",
        "Dark brown spots on lower leaves with yellow halos around them. The spots have concentric rings and older leaves are dying."
    ]
]

# Custom CSS for better styling
custom_css = """
.gradio-container {
    font-family: 'Arial', sans-serif;
}
h1 {
    text-align: center;
    color: #2E7D32;
}
"""

# Create Gradio interface
demo = gr.Interface(
    fn=diagnose_disease,
    inputs=[
        gr.Dropdown(
            choices=["MAIZE", "CASSAVA", "TOMATO"],
            label="🌱 Select Crop",
            value="MAIZE",
            info="Choose the crop you're diagnosing"
        ),
        gr.Textbox(
            lines=6,
            label="📝 Describe Symptoms",
            placeholder="Describe what you see on your plant...\n\nExample: The plant shows yellowing leaves, bore holes, visible insects, and frass. Some leaves are chewed.",
            info="Provide as much detail as possible"
        )
    ],
    outputs=gr.Markdown(
        label="🔬 Diagnosis & Recommendations"
    ),
    title="🌾 Agri Assist",
    description="""
    **AI-Powered Agricultural Diagnostic Tool**
    
    This system uses SetFit machine learning to identify crop diseases and provide treatment recommendations.
    Select your crop, describe the symptoms you observe, and get instant diagnosis with actionable advice.
    
    💡 **Tip**: The more detailed your description, the better the diagnosis!
    """,
    article="""
    ---
    
    ### 📊 About This Tool
    
    - **Model**: SetFit (Few-Shot Learning)
    - **Training**: Ghana Ministry of Agriculture Plant Clinic Data
    - **Diseases**: 66 diagnoses for Maize, Cassava, and Tomato
    - **Purpose**: Assist farmers with preliminary disease identification
    
    ### 🆘 Support
    
    For questions or issues, contact your local agricultural extension office.
    
    ---
    
    *Developed for Agricultural AI Research | 2024*
    """,
    examples=examples,
    cache_examples=False,
    theme=gr.themes.Soft(),
    css=custom_css
)

# =============================================================================
# LAUNCH
# =============================================================================

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