import streamlit as st
import numpy as np
from PIL import Image, ImageEnhance
from collections import Counter

# Ensembles pour stocker les pixels validés et non validés
oe = set()  # Pixels validés
nn = set()  # Pixels non validés

# Fonction pour traiter et redimensionner l'image importée
def process_uploaded_image(uploaded_file):
    image = Image.open(uploaded_file).convert("RGB").resize((300, 300))  # Conversion en RGB et redimensionnement à 300x300
    pixel_data = np.array(image)
    return pixel_data

# Fonction pour déterminer la couleur dominante
def dominant_color(surrounding_pixels):
    colors = [tuple(pixel) for pixel in surrounding_pixels]
    color_counts = Counter(colors)
    return np.array(color_counts.most_common(1)[0][0], dtype=np.uint8)

# Fonction pour générer une image en fonction des pixels voisins
def generate_image_from_neighbors(base_pixels):
    height, width, _ = base_pixels.shape
    generated_image = np.zeros((50, 50, 3), dtype=np.uint8)  # Image générée de 50x50 pixels

    # Itération sur chaque pixel de l'image générée
    for i in range(50):  # Hauteur de l'image générée
        for j in range(50):  # Largeur de l'image générée
            # Calculer les indices correspondants dans l'image de base
            orig_i = int(i * (height / 50))  # Mappage de 50 à 300
            orig_j = int(j * (width / 50))    # Mappage de 50 à 300

            surrounding_pixels = []
            # Collecte des couleurs des pixels environnants
            for di in range(-1, 2):  # -1, 0, 1
                for dj in range(-1, 2):  # -1, 0, 1
                    # Vérifier les limites
                    if 0 <= orig_i + di < height and 0 <= orig_j + dj < width:
                        surrounding_pixels.append(base_pixels[orig_i + di, orig_j + dj])

            # Déterminer la couleur dominante parmi les pixels environnants
            new_pixel_color = dominant_color(surrounding_pixels)

            # Assignation de la couleur dominante au pixel généré
            generated_image[i, j] = new_pixel_color

    return Image.fromarray(generated_image)

# Fonction pour appliquer un filtre à l'image
def apply_filter(image, filter_type):
    if filter_type == "Contraste":
        return ImageEnhance.Contrast(image).enhance(1.5)
    return image

# Interface utilisateur pour télécharger et afficher l'image importée
st.title("Générateur d'images d'animaux basé sur l'image importée")
st.write("Choisissez un animal pour générer une image ou téléchargez une image comme référence.")

# Téléchargement de l'image importée
uploaded_file = st.file_uploader("Téléchargez une image de référence (JPG ou PNG)", type=["jpg", "jpeg", "png"])

# Initialisation d'une image générée
generated_image = None

# Historique des images générées
image_history = []

# Génération d'image selon la référence téléchargée
if uploaded_file:
    base_pixels = process_uploaded_image(uploaded_file)
    st.image(Image.fromarray(base_pixels), caption="Image importée de référence", use_column_width=True)

    # Options de filtres
    filter_option = st.selectbox("Choisissez un filtre à appliquer :", ["Aucun", "Contraste"])
    if filter_option != "Aucun":
        base_image = Image.fromarray(base_pixels)
        base_image = apply_filter(base_image, filter_option)
        st.image(base_image, caption=f"Image avec filtre : {filter_option}", use_column_width=True)
        base_pixels = np.array(base_image)  # Mettez à jour les pixels de base

    # Bouton pour générer une image
    if st.button("Générer une image"):
        generated_image = generate_image_from_neighbors(base_pixels)
        st.image(generated_image, caption="Image générée", use_column_width=True)
        image_history.append(generated_image)  # Ajoutez l'image générée à l'historique

    # Options de validation de l'image générée
    if generated_image is not None:
        col1, col2 = st.columns(2)  # Créez deux colonnes pour les boutons
        with col1:
            if st.button("Oui"):
                # Enregistrer les pixels de l'image générée dans l'ensemble "oe"
                new_pixels = set(map(tuple, np.array(generated_image.getdata())))  # Récupérer et convertir les pixels
                oe.update(new_pixels)  # Ajouter les nouveaux pixels à l'ensemble "oe"
                st.success("Image validée ! Les pixels sont enregistrés.")

                # Réinitialiser l'image générée après validation
                generated_image = None

        with col2:
            if st.button("Non"):
                # Enregistrer les pixels de l'image générée dans l'ensemble "nn"
                new_pixels = set(map(tuple, np.array(generated_image.getdata())))  # Récupérer et convertir les pixels
                nn.update(new_pixels)  # Ajouter les nouveaux pixels à l'ensemble "nn"
                st.error("Image rejetée ! Les pixels sont enregistrés.")

                # Réinitialiser l'image générée après rejet
                generated_image = None

    # Affichage des ensembles de pixels validés et non validés
    st.write("Pixels validés (oe):", len(oe), "pixels enregistrés.")
    st.write("Pixels non validés (nn):", len(nn), "pixels enregistrés.")

    # Affichage de l'historique des images générées
    st.write("Historique des images générées:")
    for img in image_history:
        st.image(img, use_column_width=True)

# Option pour réinitialiser les ensembles
if st.button("Réinitialiser les ensembles"):
    oe.clear()
    nn.clear()
    st.success("Les ensembles de pixels ont été réinitialisés.")