import os

import cv2
import numpy as np
import torch
import torchvision.models as models
from torchvision import transforms

BASE = os.path.dirname(os.path.abspath(__file__))
MODELOS = os.path.join(BASE, "modelos")
CONOCIDOS = os.path.join(BASE, "conocidos")
YUNET_PATH = os.path.join(MODELOS, "face_detection_yunet.onnx")
ARCFACE_PATH = os.path.join(MODELOS, "buffalo_s", "w600k_mbf.onnx")
YOLO_PATH = os.path.join(BASE, "yolov8n.pt")

UMBRAL_CARA = 0.40
UMBRAL_CUERPO = 0.55


class Reconocedor:
    def __init__(self):
        self.yunet = cv2.FaceDetectorYN_create(YUNET_PATH, "", (320, 320))
        self.yunet.setScoreThreshold(0.6)
        self.arcface = cv2.dnn.readNetFromONNX(ARCFACE_PATH)
        self.body = self._cargar_body()
        self.tf_body = transforms.Compose([
            transforms.ToPILImage(),
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
        ])
        self.conocidos = self._cargar_conocidos()

    def _cargar_body(self):
        m = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1)
        m.fc = torch.nn.Identity()
        m.eval()
        return m

    def _cargar_conocidos(self):
        conocidos = {}
        if not os.path.isdir(CONOCIDOS):
            return conocidos
        for nombre in os.listdir(CONOCIDOS):
            carpeta = os.path.join(CONOCIDOS, nombre)
            if not os.path.isdir(carpeta):
                continue
            emb_cara = self._cargar_npy(carpeta, "face")
            emb_cuerpo = self._cargar_npy(carpeta, "body")
            if emb_cara is not None or emb_cuerpo is not None:
                conocidos[nombre] = (emb_cara, emb_cuerpo)
        return conocidos

    def _cargar_npy(self, carpeta, tipo):
        ruta = os.path.join(carpeta, f"{tipo}.npy")
        if os.path.exists(ruta):
            return np.load(ruta)
        return None

    # ------------------------------------------------------------------
    # Cara
    # ------------------------------------------------------------------
    def detectar_caras(self, frame):
        h, w = frame.shape[:2]
        self.yunet.setInputSize((w, h))
        _, caras = self.yunet.detect(frame)
        if caras is None:
            return []
        return [c for c in caras if c[-1] >= 0.6]

    def embedding_cara(self, frame, caja):
        x, y, w, h = (int(v) for v in caja[:4])
        x, y = max(0, x), max(0, y)
        recorte = frame[y:y + h, x:x + w]
        if recorte.size == 0:
            return None
        blob = cv2.dnn.blobFromImage(
            recorte, 1.0 / 128, (112, 112), (127.5, 127.5, 127.5), swapRB=True
        )
        self.arcface.setInput(blob)
        emb = self.arcface.forward().flatten()
        norm = np.linalg.norm(emb)
        return emb / norm if norm > 0 else None

    # ------------------------------------------------------------------
    # Cuerpo
    # ------------------------------------------------------------------
    def embedding_cuerpo(self, recorte_bgr):
        if recorte_bgr.size == 0:
            return None
        rgb = cv2.cvtColor(recorte_bgr, cv2.COLOR_BGR2RGB)
        tensor = self.tf_body(rgb).unsqueeze(0)
        with torch.no_grad():
            emb = self.body(tensor).numpy().flatten()
        norm = np.linalg.norm(emb)
        return emb / norm if norm > 0 else None

    # ------------------------------------------------------------------
    # Identificacion
    # ------------------------------------------------------------------
    def identificar_cara(self, emb):
        if emb is None or not self.conocidos:
            return None, 0.0
        mejor, score = None, -1.0
        for nombre, (emb_cara, _) in self.conocidos.items():
            if emb_cara is not None:
                sim = float(np.dot(emb, emb_cara))
                if sim > score:
                    score = sim
                    mejor = nombre
        return (mejor, score) if score >= UMBRAL_CARA else (None, score)

    def identificar_cuerpo(self, emb):
        if emb is None or not self.conocidos:
            return None, 0.0
        mejor, score = None, -1.0
        for nombre, (_, emb_cuerpo) in self.conocidos.items():
            if emb_cuerpo is not None:
                sim = float(np.dot(emb, emb_cuerpo))
                if sim > score:
                    score = sim
                    mejor = nombre
        return (mejor, score) if score >= UMBRAL_CUERPO else (None, score)

    def analizar_frame(self, frame):
        """Devuelve [{nombre:'X'} | {'desconocido':True}] por cada persona del frame."""
        personas = self._detectar_personas(frame)
        if not personas:
            return []
        resultados = []
        for p in personas:
            x, y, w, h = p
            persona = self._identificar_recorte(frame, x, y, w, h)
            resultados.append(persona)
        return resultados

    def _detectar_personas(self, frame):
        try:
            from ultralytics import YOLO
            model = self._modelo_yolo()
            r = model.predict(frame, verbose=False, conf=0.4)
            cajas = []
            for box in r[0].boxes:
                if int(box.cls) == 0:
                    x1, y1, x2, y2 = box.xyxy[0].tolist()
                    cajas.append([int(x1), int(y1), int(x2 - x1), int(y2 - y1)])
            return cajas
        except Exception:
            return []

    def _modelo_yolo(self):
        if not hasattr(self, "_yolo"):
            from ultralytics import YOLO
            self._yolo = YOLO(YOLO_PATH)
        return self._yolo

    def _identificar_recorte(self, frame, x, y, w, h):
        x, y = max(0, x), max(0, y)
        recorte = frame[y:y + h, x:x + w]
        if recorte.size == 0:
            return {"desconocido": True}

        mejor_nombre, mejor_score = None, 0.0
        for cara in self.detectar_caras(frame):
            cx, cy, cw, ch = (int(v) for v in cara[:4])
            if cx < x + w and cx + cw > x and cy < y + h and cy + ch > y:
                emb = self.embedding_cara(frame, cara)
                nombre, score = self.identificar_cara(emb)
                if nombre and score > mejor_score:
                    mejor_nombre, mejor_score = nombre, score

        if mejor_nombre:
            return {"nombre": mejor_nombre, "score": mejor_score}

        emb_cuerpo = self.embedding_cuerpo(recorte)
        nombre, score = self.identificar_cuerpo(emb_cuerpo)
        if nombre:
            return {"nombre": nombre, "score": score, "via": "cuerpo"}

        return {"desconocido": True}


def cargar_reconocedor():
    return Reconocedor()