#!/usr/bin/env python3
"""CPU post-processing for the six raw outputs of best_bpu.onnx."""

from __future__ import annotations

import numpy as np


def _nms(boxes: np.ndarray, scores: np.ndarray, iou_threshold: float) -> np.ndarray:
    order = scores.argsort()[::-1]
    kept: list[int] = []
    while order.size:
        current = int(order[0])
        kept.append(current)
        if order.size == 1:
            break
        rest = order[1:]
        xx1 = np.maximum(boxes[current, 0], boxes[rest, 0])
        yy1 = np.maximum(boxes[current, 1], boxes[rest, 1])
        xx2 = np.minimum(boxes[current, 2], boxes[rest, 2])
        yy2 = np.minimum(boxes[current, 3], boxes[rest, 3])
        inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
        area = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
        union = area[current] + area[rest] - inter
        order = rest[(inter / np.maximum(union, 1e-9)) <= iou_threshold]
    return np.asarray(kept, dtype=np.int64)


def postprocess(
    outputs: list[np.ndarray] | tuple[np.ndarray, ...],
    classes: int = 4,
    score_threshold: float = 0.25,
    iou_threshold: float = 0.7,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Decode `(cls_s8, box_s8, cls_s16, box_s16, cls_s32, box_s32)` outputs.

    Returns `(boxes_xyxy, scores, class_ids)` in the 640x640 model coordinate
    system. Scale boxes back to the camera image after this function when using
    letterbox preprocessing.
    """
    if len(outputs) != 6:
        raise ValueError(f"expected six outputs, got {len(outputs)}")
    conf_raw = -np.log(1.0 / score_threshold - 1.0)
    detections: list[np.ndarray] = []
    for index, stride in enumerate((8, 16, 32)):
        cls = np.asarray(outputs[index * 2]).reshape(-1, classes).astype(np.float32)
        box = np.asarray(outputs[index * 2 + 1]).reshape(-1, 4).astype(np.float32)
        max_logits = cls.max(axis=1)
        valid = np.flatnonzero(max_logits >= conf_raw)
        if valid.size == 0:
            continue
        height = width = 640 // stride
        grid_y, grid_x = np.indices((height, width))
        grid = np.stack((grid_x, grid_y), axis=-1).reshape(-1, 2).astype(np.float32) + 0.5
        anchor = grid[valid]
        distances = box[valid]
        xyxy = np.column_stack(
            ((anchor[:, 0] - distances[:, 0]) * stride,
             (anchor[:, 1] - distances[:, 1]) * stride,
             (anchor[:, 0] + distances[:, 2]) * stride,
             (anchor[:, 1] + distances[:, 3]) * stride)
        )
        scores = 1.0 / (1.0 + np.exp(-max_logits[valid]))
        ids = cls[valid].argmax(axis=1).astype(np.float32)
        detections.append(np.column_stack((xyxy, scores, ids)))

    if not detections:
        return np.empty((0, 4), np.float32), np.empty((0,), np.float32), np.empty((0,), np.int32)
    dets = np.concatenate(detections, axis=0)
    kept: list[np.ndarray] = []
    for class_id in np.unique(dets[:, 5]).astype(np.int32):
        class_dets = dets[dets[:, 5] == class_id]
        indices = _nms(class_dets[:, :4], class_dets[:, 4], iou_threshold)
        kept.append(class_dets[indices])
    if not kept:
        return np.empty((0, 4), np.float32), np.empty((0,), np.float32), np.empty((0,), np.int32)
    result = np.concatenate(kept, axis=0)
    return result[:, :4], result[:, 4], result[:, 5].astype(np.int32)