import math
import os
import sys
import time
from picamera2 import Picamera2
from picamera2.devices.hailo import Hailo, hailo_architecture

# Clases COCO: 15 = Cat (Gato), 16 = Dog (Perro)
TARGET_CLASSES = {15: "Gato", 16: "Perro"}
SPEED_THRESHOLD = 0.30  # Umbral de velocidad normalizada por segundo

def get_hailo_model_path():
    arch = hailo_architecture()
    candidates = []
    if arch == "HAILO8":
        candidates = [
            "/usr/share/hailo-models/yolov8s_h8.hef",
            "/usr/share/hailo-models/yolov8s_h8l.hef"
        ]
    else:
        candidates = [
            "/usr/share/hailo-models/yolov8s_h8l.hef",
            "/usr/share/hailo-models/yolov8s_h8.hef"
        ]
    for path in candidates:
        if os.path.exists(path):
            return path
    raise FileNotFoundError("No se encontró ningún modelo YOLOv8 HEF en /usr/share/hailo-models/")

def calculate_centroid(box):
    """
    Hailo YOLOv8 NMS devuelve [ymin, xmin, ymax, xmax, confidence].
    Calcula el centroide normalizado (x, y).
    """
    ymin, xmin, ymax, xmax = box[:4]
    cx = (xmin + xmax) / 2.0
    cy = (ymin + ymax) / 2.0
    return cx, cy

def main():
    model_path = get_hailo_model_path()
    print(f"[INFO] Cargando modelo Hailo: {model_path} ...")
    hailo = Hailo(model_path)
    model_h, model_w, _ = hailo.get_input_shape()
    print(f"[INFO] Modelo cargado. Resolución de entrada: {model_w}x{model_h}")

    picam2 = Picamera2()
    # Usamos configuración dual:
    # - main: 1280x720 para visualización / grabación / alertas nítidas
    # - lores: 640x640 directo al tamaño del modelo Hailo sin sobrecarga de CPU (escalado por hardware PiSP)
    config = picam2.create_video_configuration(
        main={"size": (1280, 720), "format": "RGB888"},
        lores={"size": (model_w, model_h), "format": "RGB888"},
        controls={"FrameRate": 30}
    )
    picam2.configure(config)
    picam2.start()

    print("[INFO] Pipeline Picamera2 + Hailo activo a 30 FPS.")
    print("[INFO] Monitorizando mascotas (Gatos y Perros)... Presiona Ctrl+C para salir.\n")

    last_centroid = None
    last_time = time.time()

    try:
        while True:
            # Capturamos el frame 'lores' escalado por hardware por el ISP
            lores_frame = picam2.capture_array("lores")
            current_time = time.time()
            dt = current_time - last_time

            # Inferencia por hardware en la NPU Hailo
            detections_by_class = hailo.run(lores_frame)

            pet_found = False
            for class_id, label in TARGET_CLASSES.items():
                if class_id < len(detections_by_class):
                    class_detections = detections_by_class[class_id]
                    for det in class_detections:
                        # det: [ymin, xmin, ymax, xmax, confidence]
                        confidence = det[4]

                        if confidence > 0.45:
                            pet_found = True
                            centroid = calculate_centroid(det)

                            if last_centroid and dt > 0:
                                dx = centroid[0] - last_centroid[0]
                                dy = centroid[1] - last_centroid[1]
                                dist = math.hypot(dx, dy)
                                speed = dist / dt

                                if speed > SPEED_THRESHOLD:
                                    print(f"⚠️ [ALERTA] ¡{label} corriendo! Vel: {speed:.2f} norm/s (Conf: {confidence:.2f})")
                                else:
                                    print(f"🐾 [OK] {label} detectado (calmado). Vel: {speed:.2f} norm/s (Conf: {confidence:.2f})")

                            last_centroid = centroid
                            break  # Evaluamos la mascota principal detectada
                if pet_found:
                    break

            if not pet_found:
                last_centroid = None

            last_time = current_time

    except KeyboardInterrupt:
        print("\n[INFO] Deteniendo monitor de mascotas...")
    finally:
        picam2.stop()
        hailo.close()
        print("[INFO] Cámara y acelerador Hailo liberados correctamente.")

if __name__ == "__main__":
    main()