"""
scrfd_detector.py -- SCRFD face detection with TensorRT FP16 + ONNX fallback.

SCRFD (Sample and Computation Redistribution for Face Detection) produces
9 output tensors across 3 strides (8, 16, 32), each with score, bbox, and
keypoint branches.  This module handles anchor generation, bbox/keypoint
decoding, and NMS to produce final face detections.
"""

import os
import logging
from typing import Optional

import numpy as np
import cv2

logger = logging.getLogger(__name__)

_STRIDES = [8, 16, 32]
_NUM_ANCHORS = 2  # anchors per cell per stride


# -- Inference backends --------------------------------------------------------

class _TRTBackend:
    """TensorRT inference engine for SCRFD."""

    def __init__(self, engine_path: str):
        import tensorrt as trt
        import pycuda.driver as cuda
        import pycuda.autoinit  # noqa: F401

        self._cuda = cuda
        trt_logger = trt.Logger(trt.Logger.WARNING)
        with open(engine_path, 'rb') as f, trt.Runtime(trt_logger) as runtime:
            self._engine = runtime.deserialize_cuda_engine(f.read())
        self._context = self._engine.create_execution_context()

        self._inputs = []
        self._outputs = []
        self._output_names = []
        self._bindings = []
        for i in range(self._engine.num_io_tensors):
            name = self._engine.get_tensor_name(i)
            dtype = trt.nptype(self._engine.get_tensor_dtype(name))
            shape = tuple(self._engine.get_tensor_shape(name))
            nbytes = int(np.prod(shape)) * np.dtype(dtype).itemsize
            host_mem = cuda.pagelocked_empty(shape, dtype)
            device_mem = cuda.mem_alloc(nbytes)
            self._bindings.append(int(device_mem))
            if self._engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT:
                self._inputs.append({'host': host_mem, 'device': device_mem})
            else:
                self._outputs.append({'host': host_mem, 'device': device_mem,
                                      'shape': shape})
                self._output_names.append(name)
        self._stream = cuda.Stream()

    def infer(self, input_data: np.ndarray) -> list[np.ndarray]:
        """Run inference and return output tensors."""
        np.copyto(self._inputs[0]['host'], input_data.ravel())
        self._cuda.memcpy_htod_async(
            self._inputs[0]['device'], self._inputs[0]['host'], self._stream)
        self._context.execute_async_v2(self._bindings, self._stream.handle)
        results = []
        for out in self._outputs:
            self._cuda.memcpy_dtoh_async(out['host'], out['device'], self._stream)
        self._stream.synchronize()
        for out in self._outputs:
            results.append(out['host'].reshape(out['shape']).copy())
        return results


class _ONNXBackend:
    """ONNX Runtime inference (CUDA EP with CPU fallback)."""

    def __init__(self, onnx_path: str):
        import onnxruntime as ort
        providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
        self._session = ort.InferenceSession(onnx_path, providers=providers)
        self._input_name = self._session.get_inputs()[0].name
        self._output_names = [o.name for o in self._session.get_outputs()]

    def infer(self, input_data: np.ndarray) -> list[np.ndarray]:
        """Run inference and return output tensors."""
        return self._session.run(None, {self._input_name: input_data})


# -- NMS ----------------------------------------------------------------------

def _nms(boxes: np.ndarray, scores: np.ndarray, iou_thresh: float) -> list[int]:
    """Non-maximum suppression. boxes: [N, 4] as x1,y1,x2,y2."""
    if len(boxes) == 0:
        return []

    x1, y1, x2, y2 = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3]
    areas = (x2 - x1) * (y2 - y1)
    order = scores.argsort()[::-1]
    keep = []

    while order.size > 0:
        i = order[0]
        keep.append(int(i))
        if order.size == 1:
            break
        xx1 = np.maximum(x1[i], x1[order[1:]])
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        yy2 = np.minimum(y2[i], y2[order[1:]])
        inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
        iou = inter / (areas[i] + areas[order[1:]] - inter + 1e-6)
        remaining = np.where(iou <= iou_thresh)[0]
        order = order[remaining + 1]

    return keep


# -- Anchor generation ---------------------------------------------------------

def _generate_anchors(input_h: int, input_w: int) -> dict[int, np.ndarray]:
    """Generate anchor centers for each stride.

    Returns dict mapping stride -> array of shape [H*W*num_anchors, 2],
    where each row is (cx, cy) in input pixel coordinates.
    """
    anchors = {}
    for stride in _STRIDES:
        feat_h = input_h // stride
        feat_w = input_w // stride
        grid_y, grid_x = np.mgrid[:feat_h, :feat_w]
        centers = np.stack([grid_x.ravel(), grid_y.ravel()], axis=1).astype(np.float32)
        centers = (centers + 0.5) * stride
        # Repeat for num_anchors per cell
        centers = np.tile(centers, (_NUM_ANCHORS, 1))  # [H*W*2, 2]
        # Interleave properly: [anchor0_cell0, anchor1_cell0, anchor0_cell1, ...]
        centers = np.repeat(
            np.stack([grid_x.ravel(), grid_y.ravel()], axis=1).astype(np.float32),
            _NUM_ANCHORS, axis=0
        )
        centers = (centers + 0.5) * stride
        anchors[stride] = centers
    return anchors


# -- Main detector class -------------------------------------------------------

class SCRFDDetector:
    """SCRFD face detector with TensorRT FP16 and ONNX fallback.

    Args:
        engine_path: Path to TensorRT engine file.
        onnx_path: Path to ONNX model file (used if engine not available).
        conf_thresh: Minimum confidence for detections.
        nms_iou: IoU threshold for NMS.
        input_size: Model input resolution (square).
    """

    def __init__(
        self,
        engine_path: str = '',
        onnx_path: str = '',
        conf_thresh: float = 0.5,
        nms_iou: float = 0.4,
        input_size: int = 640,
    ):
        self._conf_thresh = conf_thresh
        self._nms_iou = nms_iou
        self._input_size = input_size
        self._backend: Optional[_TRTBackend | _ONNXBackend] = None
        self._anchors = _generate_anchors(input_size, input_size)

        # Try TRT first, then ONNX
        if engine_path and os.path.isfile(engine_path):
            try:
                self._backend = _TRTBackend(engine_path)
                logger.info('SCRFD TensorRT backend loaded: %s', engine_path)
                return
            except Exception as e:
                logger.warning('SCRFD TRT load failed (%s), trying ONNX', e)

        if onnx_path and os.path.isfile(onnx_path):
            try:
                self._backend = _ONNXBackend(onnx_path)
                logger.info('SCRFD ONNX backend loaded: %s', onnx_path)
                return
            except Exception as e:
                logger.error('SCRFD ONNX load failed: %s', e)

        logger.error('No SCRFD model loaded. Detection will be unavailable.')

    @property
    def is_loaded(self) -> bool:
        """Return True if a backend is loaded and ready."""
        return self._backend is not None

    def detect(self, bgr: np.ndarray) -> list[dict]:
        """Detect faces in a BGR image.

        Args:
            bgr: Input image in BGR format, shape (H, W, 3).

        Returns:
            List of dicts with keys:
                bbox: [x1, y1, x2, y2] in original image coordinates
                kps: [x0,y0, x1,y1, ..., x4,y4] — 10 floats, 5 landmarks
                score: detection confidence
        """
        if self._backend is None:
            return []

        orig_h, orig_w = bgr.shape[:2]
        tensor, scale, pad_w, pad_h = self._preprocess(bgr)
        outputs = self._backend.infer(tensor)
        detections = self._decode_outputs(outputs)
        detections = self._rescale(detections, scale, pad_w, pad_h, orig_w, orig_h)
        return detections

    def _preprocess(self, bgr: np.ndarray) -> tuple[np.ndarray, float, int, int]:
        """Resize to input_size x input_size with letterbox padding, normalize."""
        h, w = bgr.shape[:2]
        size = self._input_size
        scale = min(size / h, size / w)
        new_w, new_h = int(w * scale), int(h * scale)
        pad_w = (size - new_w) // 2
        pad_h = (size - new_h) // 2

        resized = cv2.resize(bgr, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
        canvas = np.full((size, size, 3), 0, dtype=np.uint8)
        canvas[pad_h:pad_h + new_h, pad_w:pad_w + new_w] = resized

        # Normalize: subtract 127.5, divide 128.0
        blob = canvas.astype(np.float32)
        blob = (blob - 127.5) / 128.0
        # HWC -> NCHW
        blob = blob.transpose(2, 0, 1)[np.newaxis]
        blob = np.ascontiguousarray(blob)
        return blob, scale, pad_w, pad_h

    def _decode_outputs(self, outputs: list[np.ndarray]) -> list[dict]:
        """Decode SCRFD 9-output format into face detections.

        SCRFD outputs 9 tensors, 3 per stride (score, bbox, kps):
            score_8, bbox_8, kps_8, score_16, bbox_16, kps_16, score_32, bbox_32, kps_32
        """
        all_scores = []
        all_bboxes = []
        all_kps = []

        for idx, stride in enumerate(_STRIDES):
            score_out = outputs[idx * 3].squeeze()       # [H*W*num_anchors]
            bbox_out = outputs[idx * 3 + 1].squeeze()    # [H*W*num_anchors, 4]
            kps_out = outputs[idx * 3 + 2].squeeze()     # [H*W*num_anchors, 10]

            if score_out.ndim == 0:
                continue

            # Ensure proper shapes
            if score_out.ndim == 1:
                n = score_out.shape[0]
            else:
                n = score_out.shape[0]
                score_out = score_out.ravel()

            if bbox_out.ndim == 1:
                bbox_out = bbox_out.reshape(-1, 4)
            if kps_out.ndim == 1:
                kps_out = kps_out.reshape(-1, 10)

            # Filter by confidence
            mask = score_out > self._conf_thresh
            if not mask.any():
                continue

            scores = score_out[mask]
            bboxes = bbox_out[mask]
            kps = kps_out[mask]

            anchors = self._anchors[stride]
            # Trim or pad anchors to match output count
            if anchors.shape[0] > n:
                anchors = anchors[:n]
            elif anchors.shape[0] < n:
                continue

            anchors = anchors[mask]

            # Decode bboxes: center = anchor + pred[:2]*stride, size = exp(pred[2:])*stride
            cx = anchors[:, 0] + bboxes[:, 0] * stride
            cy = anchors[:, 1] + bboxes[:, 1] * stride
            w = np.exp(bboxes[:, 2]) * stride
            h = np.exp(bboxes[:, 3]) * stride
            x1 = cx - w / 2.0
            y1 = cy - h / 2.0
            x2 = cx + w / 2.0
            y2 = cy + h / 2.0
            decoded_bboxes = np.stack([x1, y1, x2, y2], axis=1)

            # Decode keypoints: kp = anchor + pred * stride
            decoded_kps = np.zeros_like(kps)
            for k in range(5):
                decoded_kps[:, k * 2] = anchors[:, 0] + kps[:, k * 2] * stride
                decoded_kps[:, k * 2 + 1] = anchors[:, 1] + kps[:, k * 2 + 1] * stride

            all_scores.append(scores)
            all_bboxes.append(decoded_bboxes)
            all_kps.append(decoded_kps)

        if not all_scores:
            return []

        scores = np.concatenate(all_scores)
        bboxes = np.concatenate(all_bboxes)
        kps = np.concatenate(all_kps)

        # NMS
        keep = _nms(bboxes, scores, self._nms_iou)

        results = []
        for i in keep:
            results.append({
                'bbox': bboxes[i].tolist(),
                'kps': kps[i].tolist(),
                'score': float(scores[i]),
            })
        return results

    def _rescale(
        self,
        detections: list[dict],
        scale: float,
        pad_w: int,
        pad_h: int,
        orig_w: int,
        orig_h: int,
    ) -> list[dict]:
        """Rescale detections from model input space to original image space."""
        for det in detections:
            bbox = det['bbox']
            bbox[0] = max(0.0, (bbox[0] - pad_w) / scale)
            bbox[1] = max(0.0, (bbox[1] - pad_h) / scale)
            bbox[2] = min(float(orig_w), (bbox[2] - pad_w) / scale)
            bbox[3] = min(float(orig_h), (bbox[3] - pad_h) / scale)
            det['bbox'] = bbox

            kps = det['kps']
            for k in range(5):
                kps[k * 2] = (kps[k * 2] - pad_w) / scale
                kps[k * 2 + 1] = (kps[k * 2 + 1] - pad_h) / scale
            det['kps'] = kps
        return detections
