src/streamlit_app_stable.py

import streamlit as st
import os
from PIL import Image
import torch
import gc
import time
import psutil
from transformers import AutoProcessor, Gemma3nForConditionalGeneration

# Configuration de la page
st.set_page_config(
    page_title="AgriLens AI - Analyse de Plantes",
    page_icon="🌱",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Initialisation des variables de session
if 'model_loaded' not in st.session_state:
    st.session_state.model_loaded = False
if 'model' not in st.session_state:
    st.session_state.model = None
if 'processor' not in st.session_state:
    st.session_state.processor = None
if 'model_status' not in st.session_state:
    st.session_state.model_status = "Non chargé"
if 'model_load_time' not in st.session_state:
    st.session_state.model_load_time = None
if 'language' not in st.session_state:
    st.session_state.language = "fr"
if 'load_attempt_count' not in st.session_state:
    st.session_state.load_attempt_count = 0
if 'device' not in st.session_state:
    st.session_state.device = "cpu"

# Fonctions d'aide système
def check_model_health():
    """Vérifie si le modèle et le processeur sont chargés et semblent opérationnels."""
    try:
        return (st.session_state.model is not None and
                st.session_state.processor is not None and
                hasattr(st.session_state.model, 'device'))
    except Exception:
        return False

def diagnose_loading_issues():
    """Diagnostique les problèmes potentiels avant le chargement du modèle."""
    issues = []

    try:
        ram = psutil.virtual_memory()
        ram_gb = ram.total / (1024**3)
        if ram_gb < 8:
            issues.append(f"⚠️ RAM faible: {ram_gb:.1f}GB (recommandé: 8GB+ pour ce modèle)")
    except Exception as e:
        issues.append(f"⚠️ Impossible de vérifier la RAM : {e}")

    try:
        disk_usage = psutil.disk_usage('/')
        disk_gb = disk_usage.free / (1024**3)
        if disk_gb < 10:
            issues.append(f"⚠️ Espace disque faible: {disk_gb:.1f}GB libre sur '/'")
    except Exception as e:
        issues.append(f"⚠️ Impossible de vérifier l'espace disque : {e}")

    if torch.cuda.is_available():
        try:
            gpu_memory = torch.cuda.get_device_properties(0).total_memory / (1024**3)
            if gpu_memory < 6:
                issues.append(f"⚠️ GPU mémoire faible: {gpu_memory:.1f}GB (recommandé: 6GB+)")
        except Exception as e:
            issues.append(f"⚠️ Erreur lors de la vérification de la mémoire GPU : {e}")
    else:
        issues.append("ℹ️ CUDA non disponible - Le modèle fonctionnera sur CPU (lentement)")

    return issues

def resize_image_if_needed(image, max_size=(1024, 1024)):
    """Redimensionne l'image si ses dimensions dépassent max_size."""
    original_size = image.size
    if image.size[0] > max_size[0] or image.size[1] > max_size[1]:
        image.thumbnail(max_size, Image.Resampling.LANCZOS)
        return image, True
    return image, False

def afficher_ram_disponible(context=""):
    """Affiche l'utilisation de la RAM de manière lisible."""
    try:
        ram = psutil.virtual_memory()
        ram_used_gb = ram.used / (1024**3)
        ram_total_gb = ram.total / (1024**3)
        ram_percent = ram.percent
        st.write(f"💾 RAM {context}: {ram_used_gb:.1f}GB / {ram_total_gb:.1f}GB ({ram_percent:.1f}%)")
    except Exception as e:
        st.write(f"💾 Impossible d'afficher l'utilisation de la RAM {context}: {e}")

# Gestion des traductions
def t(key):
    """Fonction pour gérer les traductions."""
    translations = {
        "fr": {
            "title": "🌱 AgriLens AI - Assistant d'Analyse de Plantes",
            "subtitle": "Analysez vos plantes avec l'IA pour détecter les maladies",
            "tabs": ["📸 Analyse d'Image", "📝 Analyse de Texte", "⚙️ Configuration", "ℹ️ À Propos"],
            "image_analysis_title": "📸 Analyse d'Image de Plante",
            "image_analysis_desc": "Téléchargez ou capturez une image de votre plante pour obtenir un diagnostic.",
            "choose_image": "Choisissez une image de plante...",
            "text_analysis_title": "📝 Analyse de Description Textuelle",
            "text_analysis_desc": "Décrivez les symptômes de votre plante pour obtenir un diagnostic personnalisé.",
            "enter_description": "Décrivez les symptômes de votre plante ici...",
            "config_title": "⚙️ Configuration & Informations",
            "about_title": "ℹ️ À Propos de l'Application",
            "load_model": "Charger le Modèle IA",
            "model_status": "Statut du Modèle IA"
        },
        "en": {
            "title": "🌱 AgriLens AI - Plant Analysis Assistant",
            "subtitle": "Analyze your plants with AI to detect diseases",
            "tabs": ["📸 Image Analysis", "📝 Text Analysis", "⚙️ Configuration", "ℹ️ About"],
            "image_analysis_title": "📸 Plant Image Analysis",
            "image_analysis_desc": "Upload or capture an image of your plant for a diagnosis.",
            "choose_image": "Choose a plant image...",
            "text_analysis_title": "📝 Textual Description Analysis",
            "text_analysis_desc": "Describe your plant's symptoms for a personalized diagnosis.",
            "enter_description": "Describe your plant's symptoms here...",
            "config_title": "⚙️ Configuration & Information",
            "about_title": "ℹ️ About the Application",
            "load_model": "Load AI Model",
            "model_status": "AI Model Status"
        }
    }
    return translations[st.session_state.language].get(key, key)

# Fonctions de chargement et d'analyse du modèle
MODEL_ID_HF = "google/gemma-3n-E4B-it" # ID du modèle sur Hugging Face Hub

def get_device_map():
    """Détermine si le modèle doit être chargé sur GPU ou CPU."""
    if torch.cuda.is_available():
        st.session_state.device = "cuda"
        return "auto"
    else:
        st.session_state.device = "cpu"
        return "cpu"

def load_model():
    """
    Charge le modèle Gemma 3n et son processeur associé depuis Hugging Face Hub.
    Gère les erreurs et les tentatives de chargement.
    """
    try:
        if st.session_state.load_attempt_count >= 3:
            st.error("❌ Trop de tentatives de chargement ont échoué. Veuillez vérifier votre configuration et redémarrer l'application.")
            return None, None

        st.session_state.load_attempt_count += 1
        st.info("🔍 Diagnostic de l'environnement avant chargement...")
        issues = diagnose_loading_issues()
        if issues:
            with st.expander("📊 Diagnostic système", expanded=False):
                for issue in issues:
                    st.write(issue)

        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        device_map = get_device_map()

        try:
            st.info(f"Chargement du modèle depuis Hugging Face Hub : {MODEL_ID_HF}...")
            processor = AutoProcessor.from_pretrained(MODEL_ID_HF, trust_remote_code=True)
            model = Gemma3nForConditionalGeneration.from_pretrained(
                MODEL_ID_HF,
                torch_dtype=torch.bfloat16 if device_map == "auto" else torch.float32,
                trust_remote_code=True,
                low_cpu_mem_usage=True,
                device_map=device_map
            )
            st.success(f"✅ Modèle chargé avec succès depuis Hugging Face Hub ({MODEL_ID_HF}).")
            st.session_state.model_status = "Chargé (Hub)"

            st.session_state.model = model
            st.session_state.processor = processor
            st.session_state.model_loaded = True
            st.session_state.model_load_time = time.time()
            st.session_state.load_attempt_count = 0

            return model, processor

        except Exception as e:
            st.error(f"❌ Échec du chargement du modèle depuis Hugging Face Hub : {e}")
            return None, None

    except ImportError:
        st.error("❌ Erreur : Les bibliothèques `transformers` ou `torch` ne sont pas installées.")
        return None, None
    except Exception as e:
        st.error(f"❌ Erreur générale lors du chargement du modèle : {e}")
        return None, None

def analyze_image_multilingual(image, prompt_text=""):
    """
    Analyse une image de plante en utilisant le modèle Gemma et un prompt personnalisé.
    Retourne le résultat de l'analyse.
    """
    if not st.session_state.model_loaded or not check_model_health():
        st.error("❌ Modèle IA non chargé ou non fonctionnel. Veuillez le charger via la barre latérale.")
        return None

    try:
        if image.mode != 'RGB':
            image = image.convert('RGB')

        if not prompt_text:
            prompt_text = """<image>
Analyse cette image de plante et fournis un diagnostic complet :
1. **État général de la plante :** Décris son apparence globale et sa vitalité.
2. **Identification des problèmes :** Liste les maladies, parasites ou carences visibles.
3. **Diagnostic probable :** Indique la maladie ou le problème le plus probable.
4. **Causes possibles :** Explique ce qui a pu causer ce problème.
5. **Recommandations de traitement :** Propose des solutions concrètes et adaptées.
6. **Conseils préventifs :** Donne des astuces pour éviter que le problème ne revienne.
Réponds de manière structurée et claire en français."""

        if "<image>" not in prompt_text:
            prompt_text = "<image>\n" + prompt_text

        inputs = st.session_state.processor(
            images=[image],
            text=prompt_text,
            return_tensors="pt"
        )

        inputs = {key: val.to(st.session_state.model.device) for key, val in inputs.items()}

        with st.spinner("🔍 Analyse d'image en cours..."):
            input_len = inputs["input_ids"].shape[-1]
            outputs = st.session_state.model.generate(
                **inputs,
                max_new_tokens=512,
                do_sample=True,
                temperature=0.7,
                top_p=0.9
            )
            generation = outputs[0][input_len:]
            response = st.session_state.processor.decode(generation, skip_special_tokens=True)

        return response.strip()

    except Exception as e:
        st.error(f"❌ Erreur lors de l'analyse de l'image : {e}")
        return None

def analyze_text_multilingual(text_description):
    """
    Analyse une description textuelle des symptômes d'une plante en utilisant le modèle Gemma.
    Retourne le diagnostic et les recommandations.
    """
    if not st.session_state.model_loaded or not check_model_health():
        st.error("❌ Modèle IA non chargé ou non fonctionnel. Veuillez le charger via la barre latérale.")
        return None

    try:
        prompt = f"""Analyse la description des symptômes de cette plante et fournis un diagnostic détaillé :
**Description des symptômes :**
{text_description}
**Instructions :**
1. **Diagnostic probable :** Quel est le problème principal ?
2. **Causes possibles :** Pourquoi ce problème survient-il ?
3. **Traitement recommandé :** Comment le résoudre ?
4. **Conseils préventifs :** Comment l'éviter à l'avenir ?
Réponds en français de manière claire et structurée."""

        inputs = st.session_state.processor(
            text=prompt,
            return_tensors="pt"
        )

        inputs = {key: val.to(st.session_state.model.device) for key, val in inputs.items()}

        with st.spinner("🔍 Analyse textuelle en cours..."):
            outputs = st.session_state.model.generate(
                **inputs,
                max_new_tokens=512,
                do_sample=True,
                temperature=0.7,
                top_p=0.9
            )

        response = st.session_state.processor.decode(outputs[0], skip_special_tokens=True)

        if prompt.strip() in response:
            response_only = response.split(prompt.strip())[-1].strip()
        else:
            response_only = response.strip()

        return response_only

    except Exception as e:
        st.error(f"❌ Erreur lors de l'analyse textuelle : {e}")
        return None

# Interface Utilisateur Streamlit
st.title(t("title"))
st.markdown(t("subtitle"))

# Barre Latérale (Sidebar) pour la Configuration
with st.sidebar:
    st.header(t("config_title"))

    lang_selector_options = ["Français", "English"]
    current_lang_index = 0 if st.session_state.language == "fr" else 1
    language_selected = st.selectbox(
        "🌐 Langue / Language",
        lang_selector_options,
        index=current_lang_index,
        help="Sélectionnez la langue de l'interface et des réponses."
    )
    st.session_state.language = "fr" if language_selected == "Français" else "en"

    st.divider()

    st.header(t("model_status"))

    if st.session_state.model_loaded and check_model_health():
        st.success("✅ Modèle chargé et fonctionnel")
        st.write(f"**Statut :** `{st.session_state.model_status}`")
        if st.session_state.model_load_time:
            load_time_str = time.strftime('%H:%M:%S', time.localtime(st.session_state.model_load_time))
            st.write(f"**Heure de chargement :** {load_time_str}")

        if st.button("🔄 Recharger le modèle", type="secondary"):
            st.session_state.model_loaded = False
            st.session_state.model = None
            st.session_state.processor = None
            st.session_state.model_status = "Non chargé"
            st.session_state.load_attempt_count = 0
            st.info("Modèle déchargé. Cliquez sur 'Charger le modèle IA' pour le recharger.")
    else:
        st.warning("⚠️ Modèle IA non chargé")

        if st.button(t("load_model"), type="primary"):
            with st.spinner("🔄 Chargement du modèle IA en cours..."):
                model_loaded_success = load_model()
                if model_loaded_success[0] is not None and model_loaded_success[1] is not None:
                    st.success("✅ Modèle IA chargé avec succès !")
                else:
                    st.error("❌ Échec du chargement du modèle IA.")

    st.divider()
    st.subheader("📊 Ressources Système")
    afficher_ram_disponible()
    if torch.cuda.is_available():
        try:
            gpu_memory = torch.cuda.memory_allocated(0) / (1024**3)
            gpu_total = torch.cuda.get_device_properties(0).total_memory / (1024**3)
            st.write(f"🚀 GPU actuel : {gpu_memory:.1f}GB / {gpu_total:.1f}GB")
        except Exception as e:
            st.write(f"🚀 Erreur lors de la récupération des informations GPU : {e}")
    else:
        st.write("🚀 GPU : Non disponible (utilisation CPU)")

# Onglets Principaux
tab1, tab2, tab3, tab4 = st.tabs(t("tabs"))

with tab1:
    st.header(t("image_analysis_title"))
    st.markdown(t("image_analysis_desc"))

    capture_option = st.radio(
        "Choisissez votre méthode de capture :",
        ["📁 Upload d'image" if st.session_state.language == "fr" else "📁 Upload Image",
         "📷 Capture par webcam" if st.session_state.language == "fr" else "📷 Webcam Capture"],
        horizontal=True,
        key="image_capture_method"
    )

    uploaded_file = None
    captured_image = None

    if capture_option == ("📁 Upload d'image" if st.session_state.language == "fr" else "📁 Upload Image"):
        uploaded_file = st.file_uploader(
            t("choose_image"),
            type=['png', 'jpg', 'jpeg'],
            help="Formats acceptés : PNG, JPG, JPEG (taille max recommandée : 10MB)."
        )
        if uploaded_file is not None and uploaded_file.size > 10 * 1024 * 1024:
            st.warning("Le fichier est très volumineux. Il est recommandé d'utiliser des images de taille raisonnable pour une analyse plus rapide.")
    else:
        st.markdown("**📷 Capture d'image par webcam**")
        st.info("Positionnez votre plante malade devant la webcam et cliquez sur 'Prendre une photo'.")
        captured_image = st.camera_input("Prendre une photo de la plante")

    image_to_analyze = None
    if uploaded_file is not None:
        try:
            image_to_analyze = Image.open(uploaded_file)
        except Exception as e:
            st.error(f"❌ Erreur lors du traitement du fichier uploadé : {e}")
    elif captured_image is not None:
        try:
            image_to_analyze = Image.open(captured_image)
        except Exception as e:
            st.error(f"❌ Erreur lors du traitement de l'image capturée : {e}")

    if image_to_analyze is not None:
        original_size = image_to_analyze.size
        resized_image, was_resized = resize_image_if_needed(image_to_analyze)

        col1, col2 = st.columns([1, 1])
        with col1:
            st.image(resized_image, caption="Image à analyser", use_container_width=True)
            if was_resized:
                st.info(f"ℹ️ Image redimensionnée de {original_size} à {resized_image.size} pour l'analyse.")

        with col2:
            if st.session_state.model_loaded and check_model_health():
                st.subheader("Options d'analyse")
                analysis_type = st.selectbox(
                    "Type d'analyse souhaitée :",
                    [
                        "Diagnostic complet (maladie, soins, prévention)" if st.session_state.language == "fr" else "Full Diagnosis (disease, care, prevention)",
                        "Identification et diagnostic de maladies" if st.session_state.language == "fr" else "Disease Identification and Diagnosis",
                        "Conseils de soins et d'entretien" if st.session_state.language == "fr" else "Care and Maintenance Advice"
                    ],
                    key="image_analysis_type_selector"
                )

                custom_prompt_input = st.text_area(
                    "Prompt personnalisé (optionnel) :",
                    value="",
                    height=100,
                    placeholder="Entrez une requête spécifique ici si besoin..."
                )

                if st.button("🔍 Analyser l'image", type="primary", key="analyze_image_button"):
                    final_prompt = custom_prompt_input.strip()
                    if not final_prompt:
                        if analysis_type.startswith("Diagnostic complet"):
                            final_prompt = """Analyse cette image de plante et fournis un diagnostic complet :
1. **État général de la plante :** Décris son apparence globale et sa vitalité.
2. **Identification des problèmes :** Liste les maladies, parasites ou carences visibles.
3. **Diagnostic probable :** Indique la maladie ou le problème le plus probable.
4. **Causes possibles :** Explique ce qui a pu causer ce problème.
5. **Recommandations de traitement :** Propose des solutions concrètes et adaptées.
6. **Conseils préventifs :** Donne des astuces pour éviter que le problème ne revienne.
Réponds de manière structurée et claire en français."""
                        elif analysis_type.startswith("Identification et diagnostic de maladies"):
                            final_prompt = """Diagnostique cette plante en te concentrant sur les maladies et parasites :
1. Identifie les symptômes visuels spécifiques aux maladies ou parasites.
2. Détermine la maladie ou le parasite le plus probable.
3. Explique les conditions favorisant leur développement.
4. Propose des traitements ciblés et des méthodes de lutte.
Réponds en français de manière structurée."""
                        else:
                            final_prompt = """Analyse cette plante et donne des conseils de soins détaillés :
1. État général de la plante : Évalue sa santé actuelle.
2. Besoins spécifiques : Précise ses besoins en eau, lumière, nutriments et substrat.
3. Conseils d'entretien : Donne des instructions pour l'arrosage, la fertilisation et la taille.
4. Améliorations recommandées : Suggère des actions pour optimiser sa croissance et sa santé.
Réponds en français de manière structurée."""

                    analysis_result = analyze_image_multilingual(resized_image, prompt_text=final_prompt)

                    if analysis_result:
                        st.success("✅ Analyse terminée !")
                        st.markdown("### 📋 Résultats de l'analyse")
                        st.markdown(analysis_result)
                    else:
                        st.error("❌ Échec de l'analyse de l'image.")
            else:
                st.warning("⚠️ Modèle IA non chargé. Veuillez d'abord charger le modèle depuis la barre latérale.")

with tab2:
    st.header(t("text_analysis_title"))
    st.markdown(t("text_analysis_desc"))

    text_description_input = st.text_area(
        t("enter_description"),
        height=200,
        placeholder="Décrivez ici les symptômes observés sur votre plante : feuilles jaunes, taches, flétrissement, présence d'insectes, etc."
    )

    if st.button("🔍 Analyser la description", type="primary", key="analyze_text_button"):
        if text_description_input.strip():
            if st.session_state.model_loaded and check_model_health():
                analysis_result = analyze_text_multilingual(text_description_input)

                if analysis_result:
                    st.success("✅ Analyse terminée !")
                    st.markdown("### 📋 Résultats de l'analyse")
                    st.markdown(analysis_result)
                else:
                    st.error("❌ Échec de l'analyse textuelle.")
            else:
                st.warning("⚠️ Modèle IA non chargé. Veuillez d'abord charger le modèle depuis la barre latérale.")
        else:
            st.warning("⚠️ Veuillez entrer une description des symptômes de votre plante.")

with tab3:
    st.header(t("config_title"))

    col1, col2 = st.columns(2)

    with col1:
        st.subheader("🔧 Informations Système")

        try:
            ram = psutil.virtual_memory()
            st.write(f"**RAM Totale :** {ram.total / (1024**3):.1f} GB")
            st.write(f"**RAM Utilisée :** {ram.used / (1024**3):.1f} GB ({ram.percent:.1f}%)")

            disk = psutil.disk_usage('/')
            st.write(f"**Espace Disque Libre (/) :** {disk.free / (1024**3):.1f} GB")

            if torch.cuda.is_available():
                st.write(f"**GPU Détecté :** {torch.cuda.get_device_name(0)}")
                gpu_props = torch.cuda.get_device_properties(0)
                st.write(f"**Mémoire GPU Totale :** {gpu_props.total_memory / (1024**3):.1f} GB")
            else:
                st.write("**GPU :** Non disponible (fonctionnement sur CPU)")
        except Exception as e:
            st.error(f"Erreur lors de la récupération des informations système : {e}")

    with col2:
        st.subheader("📊 Statistiques du Modèle IA")

        if st.session_state.model_loaded and check_model_health():
            st.write("**Statut :** ✅ Chargé et fonctionnel")
            if st.session_state.model is not None:
                st.write(f"**Type de modèle :** `{type(st.session_state.model).__name__}`")
                if hasattr(st.session_state.model, 'device'):
                    st.write(f"**Device utilisé :** `{st.session_state.model.device}`")
            if st.session_state.model_load_time:
                load_time_str = time.strftime('%H:%M:%S', time.localtime(st.session_state.model_load_time))
                st.write(f"**Heure de chargement :** {load_time_str}")
        else:
            st.write("**Statut :** ❌ Non chargé")
            st.write("**Type de modèle :** N/A")
            st.write("**Device utilisé :** N/A")

with tab4:
    st.header(t("about_title"))

    st.markdown("""
    ## 🌱 AgriLens AI : Votre Assistant d'Analyse de Plantes
    **AgriLens AI** est une application alimentée par l'intelligence artificielle qui vous aide à identifier les problèmes de vos plantes, à diagnostiquer les maladies et à obtenir des conseils de traitement personnalisés. Que vous soyez un jardinier débutant ou expérimenté, AgriLens AI est là pour vous accompagner.
    ### 🚀 Fonctionnalités Principales :
    - **Analyse d'Images de Plantes :** Téléchargez une photo de votre plante malade, et notre IA analysera les symptômes visuels pour fournir un diagnostic précis.
    - **Diagnostic par Description Textuelle :** Décrivez les symptômes que vous observez, et obtenez un avis d'expert basé sur une large base de connaissances.
    - **Recommandations Précises :** Recevez des conseils de traitement, des suggestions de soins et des mesures préventives adaptées à chaque situation.
    - **Support Multilingue :** L'interface et les réponses sont disponibles en français et en anglais pour une accessibilité maximale.
    ### 🤖 Technologie Utilisée :
    - **Modèle IA :** Google Gemma 3n E4B IT, un modèle multimodal performant pour l'analyse de plantes.
    - **Bibliothèques :** `transformers`, `torch`, `streamlit`, `Pillow`, `psutil`, `requests`, `huggingface-hub`.
    ### 📝 Comment Utiliser AgriLens AI :
    1. **Chargez le Modèle IA :** Dans la barre latérale, cliquez sur "Charger le Modèle IA". Attendez que le statut passe à "Chargé et fonctionnel".
    2. **Analysez votre Plante :**
        - **Via Image :** Allez à l'onglet "📸 Analyse d'Image". Uploadez une photo de votre plante ou utilisez votre webcam. Choisissez le type d'analyse et cliquez sur "Analyser l'image".
        - **Via Texte :** Allez à l'onglet "📝 Analyse de Texte". Décrivez les symptômes de votre plante dans la zone de texte et cliquez sur "Analyser la description".
    3. **Interprétez les Résultats :** Lisez attentivement le diagnostic, les causes probables, et les recommandations fournies par l'IA.
    ### 🔧 Support et Optimisation :
    - **Gestion des Ressources :** Le modèle nécessite une quantité significative de RAM et idéalement un GPU pour des performances optimales. L'application s'adapte à votre environnement (CPU/GPU).
    - **Performance :** Les temps de réponse dépendent de votre matériel. Un GPU est fortement recommandé pour une expérience fluide.
    ---
    *Développé avec passion pour aider les jardiniers et amoureux des plantes à prendre soin de leurs cultures grâce à la puissance de l'IA.*
    """)

# Pied de page
st.divider()
st.markdown("""
<div style='text-align: center; color: #666;'>
  🌱 AgriLens AI - Assistant d'Analyse de Plantes |
  <a href='#' target='_blank'>Documentation</a> |
  <a href='#' target='_blank'>Support</a> |
  <a href='https://huggingface.co/google/gemma-3n-E4B-it' target='_blank'>Modèle Gemma 3n</a>
</div>
""", unsafe_allow_html=True)