"""Shared helpers for NXVA validators (match_predictions, etc.)."""

import numpy as np
import torch

from .metrics import box_iou


def match_predictions(
    pred_classes: torch.Tensor,
    true_classes: torch.Tensor,
    iou: torch.Tensor,
    iouv: torch.Tensor,
    use_scipy: bool = False,
) -> torch.Tensor:
    """
    Match predictions to ground truth using IoU. Same logic as Ultralytics BaseValidator.match_predictions.
    Returns correct tensor of shape (N_pred, len(iouv)).
    """
    correct = np.zeros((pred_classes.shape[0], iouv.shape[0]), dtype=bool)
    correct_class = true_classes[:, None] == pred_classes
    iou_np = (iou * correct_class).cpu().numpy()
    thresh_list = iouv.cpu().tolist()
    for i, threshold in enumerate(thresh_list):
        matches = np.nonzero(iou_np >= threshold)
        matches = np.array(matches).T
        if matches.shape[0]:
            if matches.shape[0] > 1:
                matches = matches[iou_np[matches[:, 0], matches[:, 1]].argsort()[::-1]]
                matches = matches[np.unique(matches[:, 1], return_index=True)[1]]
                matches = matches[np.unique(matches[:, 0], return_index=True)[1]]
            correct[matches[:, 1].astype(int), i] = True
    return torch.tensor(correct, dtype=torch.bool, device=pred_classes.device)


def det_gt_from_yolo_labels(labels: np.ndarray, orig_shape: tuple) -> tuple:
    """
    Convert YOLO detect labels (N, 5) xywhn to xyxy bboxes and cls for validator.
    labels: (N, 5) with (x_center_norm, y_center_norm, w_norm, h_norm, class_id).
    orig_shape: (height, width).
    Returns (bboxes_xyxy torch (M,4), cls torch (M,)) after filtering invalid rows.
    """
    if labels is None or (hasattr(labels, "shape") and labels.shape[0] == 0):
        return torch.zeros((0, 4), dtype=torch.float32), torch.zeros(0, dtype=torch.long)
    if hasattr(labels, "shape") and len(labels.shape) == 3:
        labels = labels[0]
    labels = np.asarray(labels, dtype=np.float64)
    valid = (labels[:, 2] > 0) & (labels[:, 3] > 0)
    labels = labels[valid]
    if labels.shape[0] == 0:
        return torch.zeros((0, 4), dtype=torch.float32), torch.zeros(0, dtype=torch.long)
    h, w = orig_shape[0], orig_shape[1]
    xc = labels[:, 1] * w
    yc = labels[:, 2] * h
    ww = labels[:, 3] * w
    hh = labels[:, 4] * h
    x1 = xc - ww / 2
    y1 = yc - hh / 2
    x2 = xc + ww / 2
    y2 = yc + hh / 2
    bboxes = np.stack([x1, y1, x2, y2], axis=1).astype(np.float32)
    cls = labels[:, 0].astype(np.int64)
    return torch.from_numpy(bboxes), torch.from_numpy(cls)
