import cv2
import time
import serial
from ultralytics import YOLO

# =========================
# MODÈLE IA
# =========================

model = YOLO(r"runs\detect\runs_perso\V1\weights\best.pt")

# =========================
# REGROUPEMENT DES CLASSES
# =========================
# Mets ici les noms EXACTS des classes de ton modèle.

classes_papier = [
    "papier",
]

classes_emballage = [
    "canette",
    "Carton",
    "Emballage papier",
    "Emballage plastique",
    "plastique transparant",
    "Couverts en bois",
]

classes_pas_ici = [
    "piles",
    "Agitateur magnetique",
    "Composants electronique",
]

classes_autre = [
    "caoutchouc",
    "stylo",
    "Vis",
    "tissu",
    "PVC",
    "plexiglass",
    "Papier a poncer",
    "Mouchoir",
    "jouet en plastique",
    "Filtre",
    "fil de fer",
    "eponge",
    "colle chaude",
    "capsule cafe",
]


def classe_vers_categorie(classe, confiance):
    """
    Transforme la classe détectée par YOLO en catégorie finale de Smart Trash.
    """

    # Détections trop faibles : on ignore
    if confiance < 0.1:
        return "rien"

    # Classes précises vers catégories finales
    if classe in classes_papier and confiance > 0.45:
        return "papier"

    elif classe in classes_emballage and confiance > 0.45:
        return "emballage"

    elif classe in classes_pas_ici and confiance > 0.45:
        return "Pas ici"

    elif classe in classes_autre and confiance > 0.30:
        return "autre"

    # Si YOLO détecte quelque chose mais que ce n'est pas une classe connue dans tes listes :
    # on l'envoie dans "autre".
    elif confiance > 0.30:
        return "autre"

    else:
        return "rien"


# =========================
# CONNEXION ARDUINO
# =========================

arduino = serial.Serial("COM5", 9600, timeout=1)
time.sleep(2)

# =========================
# CAMÉRA
# =========================

cap = cv2.VideoCapture(1, cv2.CAP_DSHOW)
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)

# =========================
# PARAMÈTRES STABILITÉ / BLOCAGE
# =========================

temps_fin_blocage = 0

ancienne_box = None
categorie_stable = None
compteur_stable = 0

frames_necessaires = 10
seuil_mouvement = 20

print("IA + Arduino lancés. Appuie sur Échap pour quitter.")

# =========================
# BOUCLE PRINCIPALE
# =========================

while True:
    ret, frame = cap.read()

    if not ret:
        print("Erreur : image non lue")
        break

    maintenant = time.time()

    results = model(frame, conf=0.1, verbose=False)

    categorie = "rien"
    classe_finale = "rien"
    confiance_finale = 0
    box_actuelle = None

    if len(results[0].boxes) > 0:
        # On prend la détection avec la meilleure confiance
        best_box = max(results[0].boxes, key=lambda box: float(box.conf[0]))

        x1, y1, x2, y2 = map(int, best_box.xyxy[0])
        box_actuelle = (x1, y1, x2, y2)

        classe_id = int(best_box.cls[0])
        confiance = float(best_box.conf[0])
        classe = model.names[classe_id]

        classe_finale = classe
        confiance_finale = confiance

        # Nouvelle logique V2 :
        # classe précise → catégorie finale
        categorie = classe_vers_categorie(classe, confiance)

        # Vérification : objet à l'arrêt ?
        if ancienne_box is not None:
            ax1, ay1, ax2, ay2 = ancienne_box

            mouvement = abs(x1 - ax1) + abs(y1 - ay1) + abs(x2 - ax2) + abs(y2 - ay2)

            if mouvement < seuil_mouvement and categorie == categorie_stable:
                compteur_stable += 1
            else:
                compteur_stable = 0
                categorie_stable = categorie
        else:
            compteur_stable = 0
            categorie_stable = categorie

        ancienne_box = box_actuelle

    else:
        ancienne_box = None
        categorie_stable = None
        compteur_stable = 0
        categorie = "rien"

    # =========================
    # CATÉGORIE → COMMANDE ARDUINO
    # =========================

    if categorie == "papier":
        commande = "P"
        temps_blocage = 2

    elif categorie == "emballage":
        commande = "E"
        temps_blocage = 4.61

    elif categorie == "Pas ici":
        commande = "N"
        temps_blocage = 2.21

    elif categorie == "rien":
        commande = None
        temps_blocage = 0

    else:
        commande = "A"
        temps_blocage = 6.55

    # =========================
    # ENVOI À L'ARDUINO
    # =========================

    en_blocage = maintenant < temps_fin_blocage
    objet_est_stable = compteur_stable >= frames_necessaires

    if commande is not None and objet_est_stable and not en_blocage:
        arduino.write(commande.encode())

        print(
            "Envoyé à Arduino :",
            commande,
            "| Classe IA :",
            classe_finale,
            "| Confiance :",
            round(confiance_finale, 2),
            "| Catégorie finale :",
            categorie,
            "| Temps blocage :",
            temps_blocage,
            "s"
        )

        if commande == "N":
            print("Attention, cet objet ne doit pas être jeté ici")

        temps_fin_blocage = time.time() + temps_blocage + 0.05

        compteur_stable = 0
        ancienne_box = None
        categorie_stable = None

    # =========================
    # AFFICHAGE
    # =========================

    annotated_frame = results[0].plot()

    cv2.putText(
        annotated_frame,
        f"Classe IA: {classe_finale} | Conf: {confiance_finale:.2f} | Cat: {categorie}",
        (20, 40),
        cv2.FONT_HERSHEY_SIMPLEX,
        0.7,
        (255, 255, 255),
        2
    )

    cv2.putText(
        annotated_frame,
        f"Stable: {compteur_stable}/{frames_necessaires}",
        (20, 75),
        cv2.FONT_HERSHEY_SIMPLEX,
        0.7,
        (255, 255, 255),
        2
    )

    if en_blocage:
        cv2.putText(
            annotated_frame,
            "Mouvement en cours : commandes bloquees",
            (20, 110),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.7,
            (255, 255, 255),
            2
        )

    cv2.imshow("IA + Arduino", annotated_frame)

    if cv2.waitKey(30) == 27:
        break

cap.release()
arduino.close()
cv2.destroyAllWindows()