import streamlit as st
import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from PIL import Image
import time

# Chargement du modèle pré-entraîné
model = tf.keras.models.load_model("babong_kidney_classification_model_Tensorflow.h5")

# Redéfinition des classes (nouvelles légendes)
class_names = ["Cyst", "Normal", "Stone", "Tumor"]

# Configuration de la page et style global
st.set_page_config(page_title="Analyse d'Imagerie Rénale", page_icon="🧬", layout="wide")
st.markdown("""
    <style>
        body {
            background-color: #e0f7fa;
            font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
        }
        .main {
            background-color: #e0f7fa;
        }
        h1, h2, h3 {
            color: #006064;
            text-align: center;
        }
        .stButton>button {
            background-color: #006064;
            color: white;
            border-radius: 5px;
            padding: 10px 20px;
        }
        .upload-area {
            border: 2px dashed #006064;
            border-radius: 10px;
            padding: 20px;
            text-align: center;
            background-color: #ffffff;
            margin-bottom: 20px;
        }
        .result-card {
            background-color: #ffffff;
            border-radius: 10px;
            padding: 20px;
            margin-bottom: 20px;
            box-shadow: 2px 2px 10px rgba(0,0,0,0.1);
        }
    </style>
""", unsafe_allow_html=True)

# Titre et description principale
st.title("Analyse d'Imagerie Rénale 🧬")
st.subheader("Diagnostic assisté par Intelligence Artificielle")
st.write(
    "Chargez vos radiographies rénales pour obtenir une prédiction détaillée sur la présence de Kyste, Calcul, Tumeur ou un état Normal.")

# Zone de téléversement dans la barre latérale
st.sidebar.header("Téléversement des Images")
uploaded_files = st.sidebar.file_uploader(
    "Glissez-déposez vos images ou cliquez ici (formats acceptés: JPG, PNG, JPEG)",
    type=["jpg", "png", "jpeg"],
    accept_multiple_files=True,
    key="fileUploader"
)

if uploaded_files:
    images = []
    predictions_list = []

    # Traitement de chaque image téléversée
    for uploaded_file in uploaded_files:
        with st.container():
            st.markdown(f"<div class='result-card'><strong>Analyse de l'image : {uploaded_file.name}</strong></div>",
                        unsafe_allow_html=True)
            image = Image.open(uploaded_file)
            st.image(image, caption=f"Image fournie : {uploaded_file.name}", use_column_width=True)

            # Prétraitement de l'image
            image_resized = image.resize((224, 224))
            image_array = np.array(image_resized)
            if image_array.shape[-1] == 4:  # Gestion des images avec canal alpha
                image_array = image_array[..., :3]
            image_array = np.expand_dims(image_array, axis=0) / 255.0

            # Prédiction avec animation de chargement
            with st.spinner("Analyse en cours..."):
                time.sleep(2)
                predictions = model.predict(image_array)

            index = np.argmax(predictions)
            predicted_label = class_names[index]
            confidence = np.max(predictions)

            predictions_list.append((predicted_label, confidence, predictions[0]))
            images.append(image)

            st.success(f"**Résultat :** {predicted_label}")
            st.info(f"**Niveau de Confiance :** {confidence:.2f}")

            # Affichage graphique des scores
            fig, ax = plt.subplots(figsize=(6, 4))
            sns.barplot(x=class_names, y=predictions[0], palette="viridis", ax=ax)
            ax.set_ylim(0, 1)
            ax.set_title("Répartition des Scores de Diagnostic", fontsize=14, color="#004d40")
            ax.set_xlabel("Catégories", fontsize=12, color="#004d40")
            ax.set_ylabel("Score", fontsize=12, color="#004d40")
            st.pyplot(fig)

    # Si deux images sont téléversées, comparaison côte à côte
    if len(images) == 2:
        st.markdown("<hr>", unsafe_allow_html=True)
        st.subheader("Comparaison Visuelle et Statistique")
        col1, col2 = st.columns(2)
        with col1:
            st.image(images[0], caption="Image A", use_column_width=True)
            st.write(f"**Diagnostic :** {predictions_list[0][0]}")
            st.write(f"**Score :** {predictions_list[0][1]:.2f}")
        with col2:
            st.image(images[1], caption="Image B", use_column_width=True)
            st.write(f"**Diagnostic :** {predictions_list[1][0]}")
            st.write(f"**Score :** {predictions_list[1][1]:.2f}")

        # Graphique comparatif des deux images
        fig, ax = plt.subplots(figsize=(8, 4))
        x = np.arange(len(class_names))
        width = 0.35
        ax.bar(x - width / 2, predictions_list[0][2], width, label="Image A", color="#00796b")
        ax.bar(x + width / 2, predictions_list[1][2], width, label="Image B", color="#c62828")
        ax.set_xticks(x)
        ax.set_xticklabels(class_names)
        ax.legend(title="Images", fontsize=10)
        ax.set_title("Comparaison des Scores de Classification", fontsize=14, color="#004d40")
        ax.set_ylabel("Score de Classification", fontsize=12, color="#004d40")
        st.pyplot(fig)

# Note informative en bas de page
st.markdown("""
    <div style='text-align: center; padding: 20px;'>
        <em>Attention : Ce système fournit une aide au diagnostic basée sur l'IA. Pour un diagnostic définitif, consultez un professionnel de santé.</em>
    </div>
""", unsafe_allow_html=True)