import os
import gradio as gr
import re
from fastai.vision.all import *
from groq import Groq
from PIL import Image

# Load the trained model for grapevine disease classification
learn = load_learner('export.pkl')  # Assumes you have a trained model for grapevine diseases
labels = learn.dls.vocab

# Initialize Groq client
client = Groq(
    api_key=os.environ.get("GROQ_API_KEY"),
)

# Your existing functions remain the same
def clean_disease_name(name):
    """Clean disease name by removing numbers and special characters, and fix formatting"""
    # Remove numbers and dots at the beginning
    cleaned = re.sub(r'^\d+\.', '', name)
    # Replace underscores with spaces
    cleaned = cleaned.replace('_', ' ')
    # Remove any remaining special characters
    cleaned = re.sub(r'[^\w\s]', '', cleaned)
    # Fix spacing
    cleaned = ' '.join(cleaned.split())
    return cleaned

def format_disease_info(raw_info):
    """Improve the formatting of disease information"""
    # Add proper line breaks between sections and ensure consistent heading levels
    formatted = raw_info
   
    # Replace markdown headings with HTML headings for better control
    formatted = re.sub(r'#+\s+(.*)', r'<h3>\1</h3>', formatted)
   
    # Add paragraph tags for better spacing
    formatted = re.sub(r'\n\*\s+(.*)', r'<p>• \1</p>', formatted)
    formatted = re.sub(r'\n([^<\n].*)', r'<p>\1</p>', formatted)
   
    # Remove any duplicate paragraph tags
    formatted = formatted.replace('<p><p>', '<p>')
    formatted = formatted.replace('</p></p>', '</p>')
   
    return formatted

def get_disease_treatment(disease_name, language="swahili"):
    """Get direct treatment measures for the disease"""
    clean_name = clean_disease_name(disease_name)
   
    # Return healthy message if the prediction is healthy
    if clean_name.lower() == "healthy" or "healthy" in clean_name.lower():
        if language.lower() == "swahili":
            return """
            ### Hali ya Jani: Afya nzuri
           
            Jani hili liko katika hali nzuri ya afya. Hakuna dalili za magonjwa zinazoonekana. Endelea kutunza mizabibu yako vizuri kwa:
           
            • Kumwagilia maji ipasavyo
            • Kuhakikisha rutuba ya udongo
            • Kupunguza matawi kwa usahihi
            • Kufuatilia mara kwa mara kwa dalili za magonjwa
            • Kudumisha hewa nzuri kuzunguka mmea
            """
        else:
            return """
            ### Leaf Status: Healthy
           
            This leaf is in a healthy state. No signs of disease are present. Continue to maintain your vines by:
           
            • Proper irrigation
            • Ensuring soil fertility
            • Appropriate pruning
            • Regular monitoring for disease symptoms
            • Maintaining good air circulation around the plant
            """
   
    prompt = f"""
    Provide a concise, practical list of DIRECT treatment measures for {clean_name} disease in grapevines.
   
    Format your response as a bulleted list with 5-7 specific actions a vineyard manager should take IMMEDIATELY to treat the disease.
   
    For each treatment measure:
    1. Start with an action verb
    2. Be specific about products/techniques to use
    3. Include application rates/timing when relevant
   
    The response should be in {"Swahili" if language.lower() == "swahili" else "English"}.
   
    Begin with a short title: "HATUA ZA MATIBABU:" (if Swahili) or "TREATMENT MEASURES:" (if English)
    """
   
    try:
        chat_completion = client.chat.completions.create(
            messages=[
                {
                    "role": "user",
                    "content": prompt,
                }
            ],
            model="llama-3.3-70b-versatile",
        )
        return chat_completion.choices[0].message.content
    except Exception as e:
        error_msg = "Hitilafu kupata taarifa za matibabu" if language.lower() == "swahili" else "Error fetching treatment information"
        return f"{error_msg}: {str(e)}"

def get_disease_info(disease_name, language="swahili"):
    """Get detailed information about a grapevine disease using Groq API"""
    clean_name = clean_disease_name(disease_name)
   
    # Return healthy message if the prediction is healthy
    if clean_name.lower() == "healthy" or "healthy" in clean_name.lower():
        if language.lower() == "swahili":
            return """
            # Hali ya Jani: Afya nzuri
           
            Jani hili liko katika hali nzuri ya afya. Hakuna dalili za magonjwa zinazoonekana. Hii ni ishara nzuri kwamba:
           
            - Mizabibu yako inapata virutubisho vya kutosha
            - Umwagiliaji wa maji uko katika kiasi cha kutosha
            - Hali ya hewa katika shamba lako ni nzuri
            - Hakuna mashambulizi ya wadudu waharibifu
           
            Endelea kudumisha utunzaji mzuri wa mizabibu yako ili kuhakikisha uzalishaji bora wa zabibu.
            """
        else:
            return """
            # Leaf Status: Healthy
           
            This leaf is in a healthy state. No signs of disease are present. This is a good indication that:
           
            - Your vines are receiving adequate nutrition
            - Irrigation levels are appropriate
            - Climate conditions in your vineyard are favorable
            - There are no pest invasions
           
            Continue to maintain good vineyard practices to ensure optimal grape production.
            """
   
    lang = "SWAHILI" if language.lower() == "swahili" else "ENGLISH"
   
    prompt = f"""
    Provide detailed information IN {lang} about {clean_name} disease in grapevines.
   
    Structure your response like this:
    1. A brief introduction about what this disease is (1-2 sentences)
    2. Main symptoms and how to identify the disease (short paragraph)
    3. FOCUS EXTENSIVELY on the treatment methods and prevention strategies (this should be the longest section)
    4. Conclusion with practical advice for vineyard managers
   
    Format your response in markdown for better readability.
    Make sure the entire response is in {lang}.
    """
   
    try:
        chat_completion = client.chat.completions.create(
            messages=[
                {
                    "role": "user",
                    "content": prompt,
                }
            ],
            model="llama-3.3-70b-versatile",
        )
        return chat_completion.choices[0].message.content
    except Exception as e:
        error_msg = "Hitilafu kupata taarifa" if language.lower() == "swahili" else "Error fetching information"
        return f"{error_msg}: {str(e)}"

def predict_and_get_info(img, language="swahili"):
    """Predict grapevine disease and get detailed information"""
    # Process the image
    img = PILImage.create(img)
   
    # Get prediction
    pred, pred_idx, probs = learn.predict(img)
   
    # Get top 5 predictions (or all if less than 5)
    num_classes = min(5, len(labels))
    top_indices = probs.argsort(descending=True)[:num_classes]
    top_probs = probs[top_indices]
    top_labels = [labels[i] for i in top_indices]
   
    # Format as dictionary with cleaned names for display
    prediction_results = {clean_disease_name(top_labels[i]): float(top_probs[i]) for i in range(num_classes)}
   
    # Get top prediction (original format for info retrieval)
    top_disease = str(pred)
    # Also keep a clean version for display
    clean_top_disease = clean_disease_name(top_disease)
   
    # Check if healthy
    is_healthy = clean_top_disease.lower() == "healthy" or "healthy" in clean_top_disease.lower()
   
    # Get detailed information about the top predicted disease
    disease_info = get_disease_info(top_disease, language)
    formatted_info = format_disease_info(disease_info)
   
    # Get direct treatment measures
    treatment_measures = get_disease_treatment(top_disease, language)
    formatted_treatment = format_disease_info(treatment_measures)
   
    # Create title text based on language and health status
    if language.lower() == "swahili":
        title_text = "Jani Zuri na Afya Njema" if is_healthy else f"Ugonjwa: {clean_top_disease}"
        advice_title = "Ushauri wa Haraka kwa Wakulima"
        advice_text = "Kumbuka: Fuata maelekezo ya kitaalamu na tumia dawa zilizoidhinishwa tu. Wasiliana na mtaalamu wa kilimo kwa ushauri zaidi."
        treatment_title = "Hatua za Matibabu"
    else:
        title_text = "Healthy Leaf" if is_healthy else f"Disease: {clean_top_disease}"
        advice_title = "Quick Advice for Vineyard Managers"
        advice_text = "Remember: Follow professional instructions and use only approved treatments. Contact an agricultural expert for more advice."
        treatment_title = "Treatment Measures"
   
    # Create well-formatted display with improved fonts and styling
    custom_css = """
    <style>
    .disease-container {
        font-family: 'Nunito', 'Segoe UI', sans-serif;
        padding: 20px;
        max-width: 100%;
        background-color: #ffffff;
        border-radius: 12px;
        box-shadow: 0 4px 12px rgba(0, 0, 0, 0.08);
    }
    .disease-title {
        color: #2a633d;
        font-size: 28px;
        font-weight: 700;
        margin-bottom: 20px;
        padding-bottom: 10px;
        border-bottom: 2px solid #e0e0e0;
    }
    .info-container {
        line-height: 1.8;
        font-size: 16px;
        color: #333333;
    }
    .info-container h3 {
        color: #2a633d;
        font-size: 22px;
        font-weight: 600;
        margin-top: 25px;
        margin-bottom: 15px;
        padding-bottom: 8px;
        border-bottom: 1px solid #e0e0e0;
    }
    .info-container p {
        margin-bottom: 15px;
        text-align: justify;
    }
    .treatment-section {
        background-color: #f1f8e9;
        padding: 20px;
        border-radius: 8px;
        margin: 25px 0;
        border-left: 5px solid #558b2f;
    }
    .treatment-section h3 {
        color: #33691e;
        border-bottom: none;
        margin-top: 0;
    }
    .direct-treatment {
        background-color: #e8f5e9;
        padding: 20px;
        border-radius: 8px;
        margin: 25px 0;
        border-left: 5px solid #2e7d32;
    }
    .direct-treatment h3 {
        color: #2e7d32;
        border-bottom: none;
        margin-top: 0;
    }
    .healthy-state {
        background-color: #e8f5e9;
        color: #2e7d32;
        padding: 15px;
        border-radius: 8px;
        text-align: center;
        font-weight: bold;
        font-size: 18px;
        margin: 15px 0;
    }
    /* Ensure good readability for text */
    body {
        font-feature-settings: "kern", "liga", "clig", "calt";
        text-rendering: optimizeLegibility;
    }
    </style>
    """
   
    # Healthy banner HTML if needed
    healthy_banner = ""
    if is_healthy:
        if language.lower() == "swahili":
            healthy_banner = '<div class="healthy-state">✓ Jani hili liko katika hali nzuri ya afya</div>'
        else:
            healthy_banner = '<div class="healthy-state">✓ This leaf is in a healthy state</div>'
   
    combined_info = f"""
    {custom_css}
    <div class="disease-container">
        <h2 class="disease-title">{title_text}</h2>
        {healthy_banner}
       
        <div class="info-container">
            {formatted_info}
        </div>
       
        <div class="direct-treatment">
            <h3>{treatment_title}</h3>
            {formatted_treatment}
        </div>
       
        <div class="treatment-section">
            <h3>{advice_title}</h3>
            <p>{advice_text}</p>
        </div>
    </div>
    """
   
    return prediction_results, combined_info, clean_top_disease

def follow_up_question(question, disease_name, language="swahili"):
    """Allow vineyard managers to ask follow-up questions about the identified disease"""
    if not question.strip() or not disease_name:
        return "Tafadhali tambua ugonjwa wa mizabibu kwanza na uulize swali mahususi kuuhusu." if language.lower() == "swahili" else "Please identify a grapevine disease first and ask a specific question about it."
   
    lang = "SWAHILI" if language.lower() == "swahili" else "ENGLISH"
   
    prompt = f"""
    The vineyard manager is asking about {disease_name} disease: "{question}"
   
    Provide a detailed answer IN {lang}, focusing on accurate plant pathology information.
    If the question relates to treatment options, preventive measures, or climate change impacts, emphasize those aspects in your response.
   
    IMPORTANT:
    - Do not repeat basic introductory information about the disease that would have already been provided
    - Directly answer the specific question asked
    - The entire response must be in {lang}
    - Focus on practical advice whenever possible
   
    Format your response in markdown for better readability.
    """
   
    try:
        chat_completion = client.chat.completions.create(
            messages=[
                {
                    "role": "user",
                    "content": prompt,
                }
            ],
            model="llama-3.3-70b-versatile",
        )
        return chat_completion.choices[0].message.content
    except Exception as e:
        error_msg = "Hitilafu kupata taarifa" if language.lower() == "swahili" else "Error fetching information"
        return f"{error_msg}: {str(e)}"

# Enhanced CSS for styling the improved language toggle
toggle_css = """
<style>
/* Main container for language toggle */
.language-toggle-container {
    display: flex;
    align-items: center;
    justify-content: flex-end;
    margin-bottom: 20px;
    padding: 10px;
    background-color: rgba(255, 255, 255, 0.9);
    border-radius: 8px;
    box-shadow: 0 2px 8px rgba(0, 0, 0, 0.05);
}

/* Toggle label styling */
.toggle-label {
    font-size: 16px;
    font-weight: 600;
    margin-right: 10px;
    color: #2a633d;
}

/* Custom styling for the language toggle buttons */
.language-switcher .gr-form {
    display: flex !important;
    flex-direction: row !important;
    background: #f0f4f1;
    border-radius: 30px;
    padding: 4px;
    box-shadow: inset 0 1px 3px rgba(0, 0, 0, 0.1);
    border: 1px solid #dbeadf;
}

.language-switcher .gr-radio {
    display: inline-flex !important;
    gap: 0 !important;
    background: transparent !important;
    border: none !important;
    box-shadow: none !important;
}

.language-switcher .gr-radio-row {
    display: flex !important;
    flex-direction: row !important;
}

.language-switcher .gr-radio label {
    border-radius: 20px !important;
    padding: 8px 16px !important;
    margin: 2px !important;
    font-weight: 600 !important;
    font-size: 14px !important;
    transition: all 0.2s ease !important;
    color: #555 !important;
}

.language-switcher .gr-radio input:checked + label {
    background-color: #2a633d !important;
    color: white !important;
    box-shadow: 0 2px 4px rgba(42, 99, 61, 0.3) !important;
}

.language-switcher .gr-radio input:not(:checked) + label:hover {
    background-color: rgba(42, 99, 61, 0.1) !important;
}

/* Flag icons for visual enhancement */
.flag-icon {
    display: inline-block;
    width: 18px;
    height: 18px;
    margin-right: 6px;
    vertical-align: middle;
    border-radius: 2px;
}

/* Add animation for language transition */
.language-transition {
    transition: all 0.3s ease-in-out;
}

/* Responsive adjustments */
@media (max-width: 768px) {
    .language-toggle-container {
        justify-content: center;
    }
    
    .language-switcher .gr-radio label {
        padding: 6px 12px !important;
        font-size: 13px !important;
    }
}
</style>
"""

# Create the Gradio interface
with gr.Blocks(theme=gr.themes.Soft(), css=toggle_css) as app:
    # State for current language
    current_language = gr.State("swahili")
   
    # Create language-specific text dictionary
    text_dict = {
        "swahili": {
            "title": "# Kitambulisho cha Magonjwa ya Mizabibu",
            "subtitle": "Pakia picha ili kutambua magonjwa ya mizabibu na kupata taarifa za kina kuhusu mikakati ya udhibiti na chaguo za matibabu.",
            "upload_label": "Pakia Picha ya Mizabibu",
            "identify_btn": "Tambua Ugonjwa",
            "prediction_label": "Utabiri wa Juu 5",
            "disease_info_label": "Maelezo ya Ugonjwa",
            "followup_title": "## Maswali ya Ufuatiliaji",
            "question_label": "Uliza swali kuhusu ugonjwa huu",
            "question_placeholder": "Mfano: Je, ni hatua gani bora za kuzuia ugonjwa huu?",
            "submit_btn": "Wasilisha Swali",
            "clear_btn": "Futa Mazungumzo",
            "toggle_label": "Lugha:",
            "error_no_image": "Tafadhali pakia picha",
            "error_processing": "Hitilafu katika kuchakata picha"
        },
        "english": {
            "title": "# Grapevine Disease Identification Tool",
            "subtitle": "Upload an image to identify grapevine diseases and get detailed information on management strategies and treatment options.",
            "upload_label": "Upload Grapevine Image",
            "identify_btn": "Identify Disease",
            "prediction_label": "Top 5 Predictions",
            "disease_info_label": "Disease Information",
            "followup_title": "## Follow-up Questions",
            "question_label": "Ask a question about this disease",
            "question_placeholder": "Example: What are the best preventative measures for this disease?",
            "submit_btn": "Submit Question",
            "clear_btn": "Clear Conversation",
            "toggle_label": "Language:",
            "error_no_image": "Please upload an image",
            "error_processing": "Error processing image"
        }
    }
   
    # UI Elements with default language
    title_md = gr.Markdown(text_dict["swahili"]["title"], elem_classes="language-transition")
    subtitle_md = gr.Markdown(text_dict["swahili"]["subtitle"], elem_classes="language-transition")
   
    # Store the current disease for context
    current_disease = gr.State("")
   
    # Add improved language toggle at the top
    with gr.Row(elem_classes="language-toggle-container"):
        toggle_label = gr.Markdown(text_dict["swahili"]["toggle_label"], elem_classes="toggle-label language-transition")
        
        # Custom HTML for flag icons
        flag_html = gr.HTML("""
        <div style="display:none">
            <img id="sw-flag" src="https://flagcdn.com/w20/tz.png" alt="Swahili Flag">
            <img id="en-flag" src="https://flagcdn.com/w20/gb.png" alt="English Flag">
        </div>
        """)
        
        with gr.Column(elem_classes="language-switcher", scale=1):
            language_toggle = gr.Radio(
                ["Swahili", "English"], 
                value="Swahili", 
                label="", 
                interactive=True,
                elem_id="language-toggle"
            )
   
    # Main identification section
    with gr.Row():
        with gr.Column(scale=1):
            input_image = gr.Image(type="pil", label=text_dict["swahili"]["upload_label"], elem_classes="language-transition")
            submit_btn = gr.Button(text_dict["swahili"]["identify_btn"], variant="primary", elem_classes="language-transition")
       
        with gr.Column(scale=2):
            prediction_output = gr.Label(label=text_dict["swahili"]["prediction_label"], num_top_classes=5, elem_classes="language-transition")
            disease_info_output = gr.HTML(label=text_dict["swahili"]["disease_info_label"], elem_classes="language-transition")
   
    # Clear divider
    gr.Markdown("---")
   
    # Follow-up question section with improved UI
    followup_title_md = gr.Markdown(text_dict["swahili"]["followup_title"], elem_classes="language-transition")
   
    conversation_history = gr.Markdown("", elem_classes="language-transition")
   
    with gr.Row():
        follow_up_input = gr.Textbox(
            label=text_dict["swahili"]["question_label"],
            placeholder=text_dict["swahili"]["question_placeholder"],
            lines=2,
            elem_classes="language-transition"
        )
   
    with gr.Row():
        follow_up_btn = gr.Button(text_dict["swahili"]["submit_btn"], variant="primary", elem_classes="language-transition")
        clear_btn = gr.Button(text_dict["swahili"]["clear_btn"], elem_classes="language-transition")
   
    # Enhanced function to switch language based on toggle
    def switch_language(selected_language, current_lang):
        # Map the radio value to our language code
        new_lang = "swahili" if selected_language == "Swahili" else "english"
        
        # If language hasn't changed, don't update anything
        if new_lang == current_lang:
            return [current_lang] + [gr.update() for _ in range(11)]
        
        # Debug message to verify the language switch
        print(f"Switching language from {current_lang} to {new_lang}")
        
        # Add custom JavaScript for smooth transition animation
        js = """
        document.querySelectorAll('.language-transition').forEach(el => {
            el.style.opacity = 0;
            setTimeout(() => { el.style.opacity = 1; }, 300);
        });
        """
        
        # Return updated values
        return (
            new_lang,
            text_dict[new_lang]["title"],
            text_dict[new_lang]["subtitle"],
            text_dict[new_lang]["followup_title"],
            gr.update(label=text_dict[new_lang]["upload_label"]),
            text_dict[new_lang]["identify_btn"],
            gr.update(label=text_dict[new_lang]["prediction_label"]),
            gr.update(label=text_dict[new_lang]["disease_info_label"]),
            gr.update(label=text_dict[new_lang]["question_label"], placeholder=text_dict[new_lang]["question_placeholder"]),
            text_dict[new_lang]["submit_btn"],
            text_dict[new_lang]["clear_btn"],
            text_dict[new_lang]["toggle_label"]
        )
   
    # Set up event handlers
    def process_image(img, lang):
        if img is None:
            return None, text_dict[lang]["error_no_image"], "", ""
       
        try:
            # Print debug message to verify language
            print(f"Processing image with language: {lang}")
            
            pred_results, info, clean_disease_name = predict_and_get_info(img, lang)
            return pred_results, info, clean_disease_name, ""
        except Exception as e:
            return None, f"{text_dict[lang]['error_processing']}: {str(e)}", "", ""
   
    def update_conversation(question, disease_name, history, lang):
        if not question.strip():
            return history
       
        answer = follow_up_question(question, disease_name, lang)
       
        # Format the conversation with clear separation
        question_label = "Swali:" if lang == "swahili" else "Question:"
        answer_label = "Jibu:" if lang == "swahili" else "Answer:"
       
        new_exchange = f"""
### {question_label}
{question}
### {answer_label}
{answer}
---
"""
        updated_history = new_exchange + history
        return updated_history
   
    def clear_conversation_history():
        return ""
   
    # Connect the language toggle with enhanced functionality
    language_toggle.change(
        switch_language,
        inputs=[language_toggle, current_language],
        outputs=[
            current_language,
            title_md,
            subtitle_md,
            followup_title_md,
            input_image,
            submit_btn,
            prediction_output,
            disease_info_output,
            follow_up_input,
            follow_up_btn,
            clear_btn,
            toggle_label
        ]
    )
   
    submit_btn.click(
        process_image,
        inputs=[input_image, current_language],
        outputs=[prediction_output, disease_info_output, current_disease, conversation_history]
    )
   
    follow_up_btn.click(
        update_conversation,
        inputs=[follow_up_input, current_disease, conversation_history, current_language],
        outputs=[conversation_history]
    ).then(
        lambda: "",
        outputs=follow_up_input
    )
   
    clear_btn.click(
        clear_conversation_history,
        outputs=[conversation_history]
    )

# Launch the app
app.launch(share=True)odel
model_path = "export.pkl"  # Change to "grapevine_model.pt" if converted
learner = load_learner(model_path)

# Dictionary of disease descriptions and treatments
disease_info = {
    "Powdery Mildew": {
        "description": "A fungal disease causing white, powdery spots on leaves and fruit.",
        "treatment": "Apply sulfur-based fungicides. Improve air circulation by pruning."
    },
    "Downy Mildew": {
        "description": "Yellow spots on leaves that turn brown and cause defoliation.",
        "treatment": "Use copper-based fungicides. Avoid overhead watering."
    },
    "Black Rot": {
        "description": "Causes black spots on leaves and shriveled grapes.",
        "treatment": "Remove infected leaves and fruit. Use fungicides containing myclobutanil."
    },
    "Botrytis Bunch Rot": {
        "description": "Gray mold on grape clusters, often in humid conditions.",
        "treatment": "Ensure good air circulation. Apply biological fungicides like Bacillus subtilis."
    },
    "Healthy": {
        "description": "Your grapevine appears to be healthy!",
        "treatment": "Continue regular monitoring and proper care."
    }
}

# Define prediction function
def predict(img):
    img = Image.fromarray(img)  # Convert NumPy array to PIL image
    pred, _, probs = learner.predict(img)
    disease = str(pred)
    
    # Get disease details
    info = disease_info.get(disease, {"description": "Unknown disease", "treatment": "Consult an expert."})
    
    return {
        "Disease": disease,
        "Description": info["description"],
        "Treatment": info["treatment"]
    }

# Create Gradio interface
iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="numpy"),
    outputs=[
        gr.Label(label="Disease"),
        gr.Textbox(label="Description"),
        gr.Textbox(label="Treatment")
    ],
    title="Grapevine Disease Detection",
    description="Upload an image of a grapevine leaf to identify the disease and receive treatment recommendations."
)

# Launch the app
if __name__ == "__main__":
    iface.launch()