import streamlit as st
import requests
from transformers import DetrImageProcessor, DetrForObjectDetection
import torch
from PIL import Image, ImageDraw
import math  # Importer la bibliothèque math pour arrondir au supérieur

# Fonction pour générer du texte avec Gemini
def generate_text_with_gemini(api_key, prompt):
    url = "https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash-latest:generateContent"
    headers = {
        "Content-Type": "application/json",
    }
    params = {
        "key": api_key
    }
    data = {
        "contents": [
            {
                "parts": [
                    {"text": prompt}
                ]
            }
        ]
    }
    response = requests.post(url, headers=headers, params=params, json=data)
    
    # Debug: Print the raw response
    print(f"Gemini response status code: {response.status_code}")
    print(f"Gemini response content: {response.text}")

    try:
        response_json = response.json()
    except ValueError:
        response_json = {}
        print("Failed to decode JSON response")
    
    # Renvoyer le texte généré ou un message d'erreur
    if response.status_code == 200 and "candidates" in response_json:
        return response_json["candidates"][0]["content"]["parts"][0]["text"]
    else:
        return "Erreur lors de la génération du texte."

# Fonction de traitement de l'image et d'encadrement des personnes détectées
def process_image(image, confidence_threshold=0.7):
    processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50", revision="no_timm")
    model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50", revision="no_timm")
    
    # Préparer l'image pour le modèle
    inputs = processor(images=image, return_tensors="pt")
    outputs = model(**inputs)
    
    # Post-traitement des résultats pour obtenir les boîtes englobantes et les labels
    target_sizes = torch.tensor([image.size[::-1]])
    results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=confidence_threshold)[0]

    # Convertir l'image pour dessiner
    draw_image = image.copy()
    draw = ImageDraw.Draw(draw_image)

    # Compter le nombre de personnes détectées
    persons = 0
    for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
        if score >= confidence_threshold and model.config.id2label[label.item()] == 'person':
            persons += 1
            box = [round(i, 2) for i in box.tolist()]
            # Dessiner le rectangle rouge autour de la personne détectée
            draw.rectangle(box, outline="red", width=2)
    
    return draw_image, persons

# Interface Streamlit
st.title("Analyse d'Images pour Détection de Personnes")

# Téléchargement des images
uploaded_files = st.file_uploader("Choisissez des images (vous pouvez en ajouter plusieurs)", type=["jpg", "jpeg", "png"], accept_multiple_files=True)

# Saisie du seuil de confiance
confidence_threshold = st.slider("Seuil de confiance pour la détection (0 à 1)", 0.0, 1.0, 0.7, 0.01)

# Traitement des images et calcul de la moyenne
if uploaded_files:
    total_persons = 0
    total_images = len(uploaded_files)
    
    for i, uploaded_file in enumerate(uploaded_files):
        # Lire l'image téléchargée
        image = Image.open(uploaded_file)
        st.image(image, caption=f'Image téléchargée : {uploaded_file.name}', use_column_width=True)
        
        # Traitement de l'image
        modified_image, detected_persons = process_image(image, confidence_threshold=confidence_threshold)
        
        # Ajouter au total des personnes détectées
        total_persons += detected_persons
        
        # Afficher l'image modifiée
        st.image(modified_image, caption=f'Image avec personnes détectées : {uploaded_file.name}', use_column_width=True)
        st.write(f"**Personnes détectées sur l'image {i + 1} ({uploaded_file.name}) : {detected_persons}**")

    # Calcul de la moyenne et arrondi au supérieur
    if total_images > 0:
        average_persons = math.ceil(total_persons / total_images)  # Arrondir au supérieur
        st.write(f"**Moyenne des personnes détectées sur toutes les images : {average_persons}**")

    # Générer un e-mail avec Gemini basé sur la moyenne arrondie
    prompt = f"Génère-moi un mail destiné à l'attaché de promotion d'une école, pour lui signaler que la moyenne des personnes détectées sur les images est de {average_persons} élèves pour la classe 'ESGI 2'."
    api_key = "AIzaSyCJVUr-76I4TL42X98d9U0THbRD3MKlB38"  # Clé d'API

    email_text = generate_text_with_gemini(api_key, prompt)
    st.write("**E-mail généré :**")
    st.write(email_text)