import tensorflow as tf
import numpy as np
from PIL import Image
import gradio as gr
import zipfile
import os
import tempfile

# Charger le modèle
model = tf.keras.models.load_model("CNN_model.h5")

# Prédiction sur une image unique
def predict(image):
    image = image.resize((64, 64))
    image_array = np.array(image) / 255.0
    image_array = image_array.reshape(1, 64, 64, 3)
    prediction = model.predict(image_array)[0][0]
    return {
        "Normal": float(1 - prediction),
        "Tuberculose": float(prediction)
    }

# Résultat formaté lisible
def format_prediction_output(image):
    result = predict(image)
    normal = result["Normal"]
    tb = result["Tuberculose"]
    
    if tb > normal:
        verdict = f"🔴 Tuberculose probable avec une confiance de {tb:.2%}"
    else:
        verdict = f"🟢 Poumons normaux avec une confiance de {normal:.2%}"
    
    return verdict

# Prédiction sur un dossier ZIP
def predict_folder(zip_file):
    with tempfile.TemporaryDirectory() as tmp_input_dir, tempfile.TemporaryDirectory() as tmp_output_dir:
        try:
            with zipfile.ZipFile(zip_file.name, 'r') as zip_ref:
                zip_ref.extractall(tmp_input_dir)
        except Exception as e:
            print(" Erreur extraction ZIP :", e)
            return None

        for fname in os.listdir(tmp_input_dir):
            if fname.lower().endswith(('.png', '.jpg', '.jpeg', '.jfif')):
                try:
                    img_path = os.path.join(tmp_input_dir, fname)
                    img = Image.open(img_path).convert("RGB")
                    probs = predict(img)
                    txt_filename = os.path.splitext(fname)[0] + ".txt"
                    with open(os.path.join(tmp_output_dir, txt_filename), "w") as f:
                        for cls, prob in probs.items():
                            f.write(f"{cls}: {prob:.4f}\n")
                except Exception as e:
                    print(f"Erreur image {fname} :", e)

        output_zip_path = os.path.join(tempfile.gettempdir(), "predictions_text.zip")
        with zipfile.ZipFile(output_zip_path, 'w') as zip_out:
            for file_out in os.listdir(tmp_output_dir):
                zip_out.write(os.path.join(tmp_output_dir, file_out), arcname=file_out)
        return output_zip_path

# Exemples d’images
examples = [["Patient1.jfif"], ["Patient2.jfif"], ["Patient3.jfif"]]

# Interface Gradio
with gr.Blocks(theme=gr.themes.Glass()) as demo:
    gr.Markdown("<h1 style='text-align: center;'>🩺 Détection automatisée de la tuberculose</h1>")
    gr.Markdown("<p style='text-align: center;'>Analyse des images radiographiques pulmonaires pour détecter les signes de tuberculose.</p>")

    with gr.Tabs():
        with gr.TabItem("📷 Importer une image"):
            with gr.Row():
                with gr.Column(scale=2):
                    image_input = gr.Image(type="pil", label="Radiographie pulmonaire")
                    btn_predict = gr.Button("🔍 Analyser l’image")
                with gr.Column(scale=1):
                    output_label = gr.Label(label="🧪 Résultat d’analyse")
            btn_predict.click(fn=format_prediction_output, inputs=image_input, outputs=output_label)

            with gr.Accordion("🖼️ Exemples de radiographies", open=False):
                gr.Examples(
                    examples=examples,
                    inputs=[image_input],
                    label="Cliquez sur une image pour la tester"
                )

        with gr.TabItem("📁 Dossier ZIP d'images"):
            zip_input = gr.File(label="Télécharger un ZIP contenant plusieurs images", file_types=[".zip"])
            zip_output = gr.File(label="Télécharger les résultats (ZIP de fichiers texte)")
            btn_process = gr.Button("🧪 Analyser le dossier")
            btn_process.click(fn=predict_folder, inputs=zip_input, outputs=zip_output)

    gr.Markdown("---")

    with gr.Accordion("💡 Astuces & Prévention contre la tuberculose", open=False):
        gr.Markdown("""
- 🛡️ **Vaccination BCG** dès la naissance dans les zones à risque.
- 🧼 **Hygiène respiratoire** : tousser dans le coude, porter un masque.
- 🏠 **Ventilation** : bien aérer les lieux clos.
- 🩺 **Dépistage précoce** : consulter si toux prolongée, fièvre, perte de poids.
- 💊 **Suivi médical strict** : ne jamais interrompre le traitement.
        """)

    gr.Markdown("---")
    gr.Markdown("### 🧠 Suggestions")
    gr.Markdown("""
✅ Combinez ce modèle de l'IA avec un avis médical réel.  
✅ Consultez un professionnel de santé en cas de doute.  
✅ Cette application est un **outil d’aide au diagnostic**, pas un substitut médical.
    """)

demo.launch()