# Copyright 2025-2026 Yakhyokhuja Valikhujaev
# Author: Yakhyokhuja Valikhujaev
# GitHub: https://github.com/yakhyo

from typing import Any, Literal

import numpy as np

from uniface.common import letterbox_resize, non_max_suppression
from uniface.constants import YOLOv5FaceWeights
from uniface.log import Logger
from uniface.model_store import verify_model_weights
from uniface.onnx_utils import create_onnx_session
from uniface.types import Face

from .base import BaseDetector

# Optional torchvision import for faster NMS
try:
    import torch
    import torchvision

    TORCHVISION_AVAILABLE = True
except ImportError:
    TORCHVISION_AVAILABLE = False

__all__ = ['YOLOv5Face']


class YOLOv5Face(BaseDetector):
    """
    Face detector based on the YOLOv5-Face architecture.

    Title: "YOLO5Face: Why Reinventing a Face Detector"
    Paper: https://arxiv.org/abs/2105.12931
    Code: https://github.com/yakhyo/yolov5-face-onnx-inference (ONNX inference implementation)

    Args:
        model_name (YOLOv5FaceWeights): Predefined model enum (e.g., `YOLOV5S`).
            Specifies the YOLOv5-Face variant to load. Defaults to YOLOV5S.
        confidence_threshold (float): Confidence threshold for filtering detections. Defaults to 0.6.
        nms_threshold (float): Non-Maximum Suppression threshold. Defaults to 0.5.
        input_size (int): Input image size. Defaults to 640.
            Note: ONNX model is fixed at 640. Changing this will cause inference errors.
        nms_mode (str): NMS calculation method. Options: 'torchvision' (faster, requires torch)
            or 'numpy' (no dependencies). Defaults to 'numpy'.
        providers (list[str] | None): ONNX Runtime execution providers. If None, auto-detects
            the best available provider. Example: ['CPUExecutionProvider'] to force CPU.
        **kwargs: Advanced options:
            max_det (int): Maximum number of detections to return. Defaults to 750.

    Attributes:
        model_name (YOLOv5FaceWeights): Selected model variant.
        confidence_threshold (float): Threshold used to filter low-confidence detections.
        nms_threshold (float): Threshold used during NMS to suppress overlapping boxes.
        input_size (int): Image size to which inputs are resized before inference.
        nms_mode (str): NMS calculation method being used.
        max_det (int): Maximum number of detections to return.
        _model_path (str): Absolute path to the downloaded/verified model weights.

    Raises:
        ValueError: If the model weights are invalid or not found.
        RuntimeError: If the ONNX model fails to load or initialize.
    """

    def __init__(
        self,
        *,
        model_name: YOLOv5FaceWeights = YOLOv5FaceWeights.YOLOV5S,
        confidence_threshold: float = 0.6,
        nms_threshold: float = 0.5,
        input_size: int = 640,
        nms_mode: Literal['torchvision', 'numpy'] = 'numpy',
        providers: list[str] | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(
            model_name=model_name,
            confidence_threshold=confidence_threshold,
            nms_threshold=nms_threshold,
            input_size=input_size,
            nms_mode=nms_mode,
            providers=providers,
            **kwargs,
        )
        self._supports_landmarks = True  # YOLOv5-Face supports landmarks

        # Validate input size
        if input_size != 640:
            raise ValueError(
                f'YOLOv5Face only supports input_size=640 (got {input_size}). The ONNX model has a fixed input shape.'
            )

        self.model_name = model_name
        self.confidence_threshold = confidence_threshold
        self.nms_threshold = nms_threshold
        self.input_size = input_size
        self.providers = providers

        # Set NMS mode with automatic fallback
        if nms_mode == 'torchvision' and not TORCHVISION_AVAILABLE:
            Logger.warning('torchvision not available, falling back to numpy NMS')
            self.nms_mode = 'numpy'
        else:
            self.nms_mode = nms_mode

        # Advanced options from kwargs
        self.max_det = kwargs.get('max_det', 750)

        Logger.info(
            f'Initializing YOLOv5Face with model={self.model_name}, confidence_threshold={self.confidence_threshold}, '
            f'nms_threshold={self.nms_threshold}, input_size={self.input_size}, nms_mode={self.nms_mode}'
        )

        # Get path to model weights
        self._model_path = verify_model_weights(self.model_name)
        Logger.info(f'Verified model weights located at: {self._model_path}')

        # Initialize model
        self._initialize_model(self._model_path)

    def _initialize_model(self, model_path: str) -> None:
        """
        Initializes an ONNX model session from the given path.

        Args:
            model_path (str): The file path to the ONNX model.

        Raises:
            RuntimeError: If the model fails to load, logs an error and raises an exception.
        """
        try:
            self.session = create_onnx_session(model_path, providers=self.providers)
            self.input_names = self.session.get_inputs()[0].name
            self.output_names = [x.name for x in self.session.get_outputs()]
            Logger.info(f'Successfully initialized the model from {model_path}')
        except Exception as e:
            Logger.error(f"Failed to load model from '{model_path}': {e}", exc_info=True)
            raise RuntimeError(f"Failed to initialize model session for '{model_path}'") from e

    def preprocess(self, image: np.ndarray) -> tuple[np.ndarray, float, tuple[int, int]]:
        """Preprocess image using letterbox resize.

        Args:
            image: Input image in BGR format.

        Returns:
            Tuple of (preprocessed_tensor, scale_ratio, padding).
        """
        return letterbox_resize(image, self.input_size)

    def inference(self, input_tensor: np.ndarray) -> list[np.ndarray]:
        """Perform model inference on the preprocessed image tensor.

        Args:
            input_tensor (np.ndarray): Preprocessed input tensor.

        Returns:
            List[np.ndarray]: Raw model outputs.
        """
        return self.session.run(self.output_names, {self.input_names: input_tensor})

    def postprocess(
        self,
        predictions: np.ndarray,
        scale: float,
        padding: tuple[int, int],
    ) -> tuple[np.ndarray, np.ndarray]:
        """
        Postprocess model predictions.

        Args:
            predictions (np.ndarray): Raw model output
            scale (float): Scale ratio used in preprocessing
            padding (Tuple[int, int]): Padding used in preprocessing

        Returns:
            Tuple[np.ndarray, np.ndarray]: Filtered detections and landmarks
                - detections: [x1, y1, x2, y2, conf]
                - landmarks: [5, 2] for each detection
        """
        # predictions shape: (1, 25200, 16)
        # 16 = [x, y, w, h, obj_conf, cls_conf, 10 landmarks (5 points * 2 coords)]

        predictions = predictions[0]  # Remove batch dimension

        # Filter by confidence
        mask = predictions[:, 4] >= self.confidence_threshold
        predictions = predictions[mask]

        if len(predictions) == 0:
            return np.array([]), np.array([])

        # Convert from xywh to xyxy
        boxes = self._xywh2xyxy(predictions[:, :4])

        # Get confidence scores
        scores = predictions[:, 4]

        # Get landmarks (5 points, 10 coordinates)
        landmarks = predictions[:, 5:15].copy()

        # Apply NMS based on selected mode
        if self.nms_mode == 'torchvision':
            keep = torchvision.ops.nms(
                torch.tensor(boxes, dtype=torch.float32),
                torch.tensor(scores, dtype=torch.float32),
                self.nms_threshold,
            ).numpy()
        else:
            detections = np.hstack((boxes, scores[:, None])).astype(np.float32, copy=False)
            keep = non_max_suppression(detections, self.nms_threshold)

        if len(keep) == 0:
            return np.array([]), np.array([])

        # Filter detections and limit to max_det
        keep = keep[: self.max_det]
        boxes = boxes[keep]
        scores = scores[keep]
        landmarks = landmarks[keep]

        # Scale back to original image coordinates
        pad_w, pad_h = padding
        boxes[:, [0, 2]] = (boxes[:, [0, 2]] - pad_w) / scale
        boxes[:, [1, 3]] = (boxes[:, [1, 3]] - pad_h) / scale

        # Scale landmarks
        for i in range(5):
            landmarks[:, i * 2] = (landmarks[:, i * 2] - pad_w) / scale
            landmarks[:, i * 2 + 1] = (landmarks[:, i * 2 + 1] - pad_h) / scale

        # Reshape landmarks to (N, 5, 2)
        landmarks = landmarks.reshape(-1, 5, 2)

        # Combine results
        detections = np.concatenate([boxes, scores[:, None]], axis=1)

        return detections, landmarks

    def _xywh2xyxy(self, x: np.ndarray) -> np.ndarray:
        """
        Convert bounding box format from xywh to xyxy.

        Args:
            x (np.ndarray): Boxes in [x, y, w, h] format

        Returns:
            np.ndarray: Boxes in [x1, y1, x2, y2] format
        """
        y = np.copy(x)
        y[..., 0] = x[..., 0] - x[..., 2] / 2  # x1
        y[..., 1] = x[..., 1] - x[..., 3] / 2  # y1
        y[..., 2] = x[..., 0] + x[..., 2] / 2  # x2
        y[..., 3] = x[..., 1] + x[..., 3] / 2  # y2
        return y

    def detect(
        self,
        image: np.ndarray,
        *,
        max_num: int = 0,
        metric: Literal['default', 'max'] = 'max',
        center_weight: float = 2.0,
    ) -> list[Face]:
        """
        Perform face detection on an input image and return bounding boxes and facial landmarks.

        Args:
            image (np.ndarray): Input image as a NumPy array of shape (H, W, C).
            max_num (int): Maximum number of detections to return. Use 0 to return all detections. Defaults to 0.
            metric (Literal["default", "max"]): Metric for ranking detections when `max_num` is limited.
                - "default": Prioritize detections closer to the image center.
                - "max": Prioritize detections with larger bounding box areas.
            center_weight (float): Weight for penalizing detections farther from the image center
                when using the "default" metric. Defaults to 2.0.

        Returns:
            List[Face]: List of Face objects, each containing:
                - bbox (np.ndarray): Bounding box coordinates with shape (4,) as [x1, y1, x2, y2]
                - confidence (float): Detection confidence score (0.0 to 1.0)
                - landmarks (np.ndarray): 5-point facial landmarks with shape (5, 2)

        Example:
            >>> faces = detector.detect(image)
            >>> for face in faces:
            ...     bbox = face.bbox  # np.ndarray with shape (4,)
            ...     confidence = face.confidence  # float
            ...     landmarks = face.landmarks  # np.ndarray with shape (5, 2)
            ...     # Can pass landmarks directly to recognition
            ...     embedding = recognizer.get_normalized_embedding(image, face.landmarks)
        """

        original_height, original_width = image.shape[:2]

        # Preprocess
        image_tensor, scale, padding = self.preprocess(image)

        # ONNXRuntime inference
        outputs = self.inference(image_tensor)

        # Postprocess
        detections, landmarks = self.postprocess(outputs[0], scale, padding)

        # Handle case when no faces are detected
        if len(detections) == 0:
            return []

        # Filter to top max_num faces if requested
        detections, landmarks = self._select_top_detections(
            detections, landmarks, max_num, (original_height, original_width), metric, center_weight
        )

        return self._detections_to_faces(detections, landmarks)
