import torch


def nms(
    p_bars: torch.Tensor,
    p_ticks: torch.Tensor,
    radius_thresh: float
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Performs non-maximum suppression on predicted bars and ticks.

    :param p_bars: predicted bars [y, x, conf], shape [N, 3]
    :param p_ticks: predicted ticks [y, x, conf], shape [M, 3]
    :param radius_thresh: threshold distance for considering points as neighbors
    :return: tuple containing:

        - nms filtered bars [y, x, conf], shape [N, 3]
        - nms filtered ticks [y, x, conf], shape [M, 3]
    """
    bar_out = p_bars.new_empty((0, 3))
    tick_out = p_ticks.new_empty((0, 3))

    # Sort bars by confidence
    if p_bars.size(0) > 0:
        sorted_bars = p_bars[torch.argsort(p_bars[:, 2], descending=True)]
        checked_bars = torch.zeros(sorted_bars.size(0), dtype=torch.bool, device=sorted_bars.device)
        keep_bars = torch.zeros_like(checked_bars)

        # Iteratively remove low confidence bars within radius
        for k in range(sorted_bars.size(0)):
            if checked_bars[k]:
                continue
            pt = sorted_bars[k]
            # Keep current bar, neighbors within radius are removed
            distances = torch.abs(sorted_bars[:, 0] - pt[0]) + torch.abs(sorted_bars[:, 1] - pt[1])
            checked_bars |= distances.lt(radius_thresh)
            keep_bars[k] = True

        bar_out = sorted_bars[keep_bars]

    # Sort ticks by confidence
    if p_ticks.size(0) > 0:
        sorted_ticks = p_ticks[torch.argsort(p_ticks[:, 2], descending=True)]
        checked_ticks = torch.zeros(sorted_ticks.size(0), dtype=torch.bool, device=sorted_ticks.device)
        keep_ticks = torch.zeros_like(checked_ticks)

        # Iteratively remove low confidence ticks within radius
        for k in range(sorted_ticks.size(0)):
            if checked_ticks[k]:
                continue
            pt = sorted_ticks[k]
            # Keep current tick, neighbors within radius are removed
            distances = torch.abs(sorted_ticks[:, 0] - pt[0]) + torch.abs(sorted_ticks[:, 1] - pt[1])
            checked_ticks |= distances.lt(radius_thresh)
            keep_ticks[k] = True

        tick_out = sorted_ticks[keep_ticks]

    return bar_out, tick_out


def hungarian_match(
    gt_points: torch.Tensor,
    p_points: torch.Tensor,
    dist_thresh: float
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Matches GT and predictions with Hungarian matching under a distance threshold.

    :param gt_points: ground truth points [N, 2]
    :param p_points: predicted points [M, 2]
    :param dist_thresh: max distance for a valid match
    :return: tuple containing:

        - matched GT indices [K]
        - matched prediction indices [K]
    """
    device = gt_points.device
    empty = torch.empty((0,), dtype=torch.long, device=device)

    # Nothing to match if either side is empty
    if gt_points.size(0) == 0 or p_points.size(0) == 0:
        return empty, empty

    # Pairwise distances between gt and predicted coords
    diffs = gt_points.unsqueeze(1) - p_points.unsqueeze(0)
    distances = torch.sqrt((diffs * diffs).sum(dim=2))

    # Distances above threshold are not valid candidate matches
    invalid = distances.gt(dist_thresh)
    if bool(invalid.all()):
        return empty, empty

    # Coordinates in [0, 1], so max distance is sqrt(2)
    large_cost = (2.0 ** 0.5) + 1e-6
    cost = distances.masked_fill(invalid, large_cost)

    # Global one-to-one assignment minimizing total matching cost
    from scipy.optimize import linear_sum_assignment
    row_idx, col_idx = linear_sum_assignment(cost.detach().cpu().numpy())

    row_idx_t = torch.as_tensor(row_idx, device=device, dtype=torch.long)
    col_idx_t = torch.as_tensor(col_idx, device=device, dtype=torch.long)

    # Keep only assignments that truly satisfy the threshold
    valid_matches = distances[row_idx_t, col_idx_t].le(dist_thresh)
    return row_idx_t[valid_matches], col_idx_t[valid_matches]


def evaluate_gt_p_match(
    gt_bars: torch.Tensor,
    gt_ticks: torch.Tensor,
    p_bars: torch.Tensor,
    p_ticks: torch.Tensor,
    dist_thresh: float
) -> tuple[float, float, float, float, float, float]:
    """
    Evaluate predicted bars and ticks against ground truth.

    :param gt_bars: ground truth bars [y, x, value], shape [N, 3]
    :param gt_ticks: ground truth ticks [y, x, value], shape [M, 3]
    :param p_bars: predicted bars [y, x, conf], shape [N, 3]
    :param p_ticks: predicted ticks [y, x, conf], shape [M, 3]
    :param dist_thresh: max. distance threshold when considering matches
    :return: Tuple containing:

        - bar precision, recall, f1
        - tick precision, recall, f1
    """
    bar_rows, _ = hungarian_match(gt_bars[:, :2], p_bars[:, :2], dist_thresh)
    tick_rows, _ = hungarian_match(gt_ticks[:, :2], p_ticks[:, :2], dist_thresh)
    bar_matches = int(bar_rows.numel())
    tick_matches = int(tick_rows.numel())

    # Calculate precision and recall for bars
    bar_precision = (bar_matches / p_bars.size(0)) if p_bars.size(0) > 0 else 0.
    bar_recall = (bar_matches / gt_bars.size(0)) if gt_bars.size(0) > 0 else 0.
    bar_f1 = f1(bar_precision, bar_recall)

    # Calculate precision and recall for ticks
    tick_precision = (tick_matches / p_ticks.size(0)) if p_ticks.size(0) > 0 else 0.
    tick_recall = (tick_matches / gt_ticks.size(0)) if gt_ticks.size(0) > 0 else 0.
    tick_f1 = f1(tick_precision, tick_recall)

    return bar_precision, bar_recall, bar_f1, tick_precision, tick_recall, tick_f1


def evaluate_value_accuracy(
    gt_bar_yxv: torch.Tensor,
    p_bar_yxv: torch.Tensor,
    dist_thresh: float,
    hard_eps: float = 0.02,
    relaxed_eps: float = 0.05
) -> tuple[float, float]:
    """
    Evaluate bar-value accuracy with padded-count denominator.

    :param gt_bar_yxv: ground truth bars [N, 3] as [y, x, value]
    :param p_bar_yxv: predicted bars [M, 3] as [y, x, value]
    :param dist_thresh: max. spatial distance for matching bars
    :param hard_eps: strict relative error threshold
    :param relaxed_eps: relaxed relative error threshold
    :return: tuple containing:

        - hard-threshold accuracy
        - relaxed-threshold accuracy
    """
    gt_bars = gt_bar_yxv[:, :2]
    gt_values = gt_bar_yxv[:, 2]
    p_bars = p_bar_yxv[:, :2]
    p_values = p_bar_yxv[:, 2]

    # Ignore invalid gt values
    valid_gt_mask = torch.isfinite(gt_values)
    if not bool(valid_gt_mask.any()):
        return 0.0, 0.0

    p_count = int(min(p_bars.size(0), p_values.numel()))
    total = max(int(gt_bars[valid_gt_mask].size(0)), p_count)
    if total == 0 or p_count == 0:
        return 0.0, 0.0

    # Match predicted and gt bars using Hungarian assignment
    gt_idx, pred_idx = hungarian_match(
        gt_bars[valid_gt_mask],
        p_bars[:p_count],
        dist_thresh
    )
    if gt_idx.numel() == 0:
        return 0.0, 0.0

    # Criterion = abs(h_g - h_p) / h_g <= eps.
    gt_matched = gt_values[valid_gt_mask][gt_idx]
    p_matched = p_values[:p_count][pred_idx]
    denom = gt_matched.abs().clamp(min=1e-8)
    rel_err = (gt_matched - p_matched).abs() / denom

    # Defaults: hard = 0.02, relaxed = 0.05
    hard_correct = int(rel_err.le(hard_eps).sum().item())
    relaxed_correct = int(rel_err.le(relaxed_eps).sum().item())
    hard_acc = float(hard_correct / total)
    relaxed_acc = float(relaxed_correct / total)
    return hard_acc, relaxed_acc


def evaluate_confusion(
    gt_bars: torch.Tensor,
    gt_ticks: torch.Tensor,
    p_bars: torch.Tensor,
    p_ticks: torch.Tensor,
    dist_thresh: float
) -> tuple[int, int, int, int, int, int, int, int]:
    """
    Returns a full confusion matrix for bar/tick keypoint detection.

                     p_bar       p_tick       p_none
        gt_bar       tp_bar      bar_as_tick  fn_bar
        gt_tick      tick_as_bar tp_tick      fn_tick
        gt_none      fp_bar      fp_tick      –

    :param gt_bars: ground truth bars [y, x, value], shape [N, 3]
    :param gt_ticks: ground truth ticks [y, x, value], shape [M, 3]
    :param p_bars: predicted bars [y, x, conf], shape [P, 3]
    :param p_ticks: predicted ticks [y, x, conf], shape [Q, 3]
    :param dist_thresh: max. distance threshold for matching
    :return: tuple containing:

        - tp_bar: GT bars matched with predicted bars
        - bar_as_tick: GT bars matched with predicted ticks (cross-class)
        - fn_bar: GT bars with no match
        - fp_bar: predicted bars with no GT match
        - tp_tick: GT ticks matched with predicted ticks
        - tick_as_bar: GT ticks matched with predicted bars (cross-class)
        - fn_tick: GT ticks with no match
        - fp_tick: predicted ticks with no GT match
    """
    device = gt_bars.device
    n_gt_bars = gt_bars.size(0)
    n_gt_ticks = gt_ticks.size(0)
    n_p_bars = p_bars.size(0)
    n_p_ticks = p_ticks.size(0)

    empty = torch.empty(0, dtype=torch.long, device=device)

    # --- Pass 1: within-class matching ---
    if n_gt_bars > 0 and n_p_bars > 0:
        gt_bar_idx, p_bar_idx = hungarian_match(gt_bars[:, :2], p_bars[:, :2], dist_thresh)
    else:
        gt_bar_idx, p_bar_idx = empty, empty

    if n_gt_ticks > 0 and n_p_ticks > 0:
        gt_tick_idx, p_tick_idx = hungarian_match(gt_ticks[:, :2], p_ticks[:, :2], dist_thresh)
    else:
        gt_tick_idx, p_tick_idx = empty, empty

    tp_bar = int(gt_bar_idx.numel())
    tp_tick = int(gt_tick_idx.numel())

    # Build masks of unmatched points after pass 1
    unmatched_gt_bar = torch.ones(n_gt_bars, dtype=torch.bool, device=device)
    unmatched_p_bar = torch.ones(n_p_bars, dtype=torch.bool, device=device)
    unmatched_gt_tick = torch.ones(n_gt_ticks, dtype=torch.bool, device=device)
    unmatched_p_tick = torch.ones(n_p_ticks, dtype=torch.bool, device=device)

    if gt_bar_idx.numel() > 0:
        unmatched_gt_bar[gt_bar_idx] = False
        unmatched_p_bar[p_bar_idx] = False
    if gt_tick_idx.numel() > 0:
        unmatched_gt_tick[gt_tick_idx] = False
        unmatched_p_tick[p_tick_idx] = False

    # --- Pass 2: cross-class matching on remaining unmatched points ---
    ub_gt_bar = gt_bars[unmatched_gt_bar, :2]
    ub_p_tick = p_ticks[unmatched_p_tick, :2]
    ub_gt_tick = gt_ticks[unmatched_gt_tick, :2]
    ub_p_bar = p_bars[unmatched_p_bar, :2]

    # Unmatched GT bars vs unmatched pred ticks
    if ub_gt_bar.size(0) > 0 and ub_p_tick.size(0) > 0:
        xm_gt_bar, xm_p_tick = hungarian_match(ub_gt_bar, ub_p_tick, dist_thresh)
    else:
        xm_gt_bar, xm_p_tick = empty, empty

    bar_as_tick = int(xm_gt_bar.numel())

    # Unmatched GT ticks vs unmatched pred bars
    if ub_gt_tick.size(0) > 0 and ub_p_bar.size(0) > 0:
        xm_gt_tick, xm_p_bar = hungarian_match(ub_gt_tick, ub_p_bar, dist_thresh)
    else:
        xm_gt_tick, xm_p_bar = empty, empty

    tick_as_bar = int(xm_gt_tick.numel())

    # Update masks based on cross-class matches
    unmatched_gt_bar_idx = unmatched_gt_bar.nonzero(as_tuple=False).squeeze(1)
    unmatched_p_tick_idx = unmatched_p_tick.nonzero(as_tuple=False).squeeze(1)
    unmatched_gt_tick_idx = unmatched_gt_tick.nonzero(as_tuple=False).squeeze(1)
    unmatched_p_bar_idx = unmatched_p_bar.nonzero(as_tuple=False).squeeze(1)

    if xm_gt_bar.numel() > 0:
        unmatched_gt_bar[unmatched_gt_bar_idx[xm_gt_bar]] = False
        unmatched_p_tick[unmatched_p_tick_idx[xm_p_tick]] = False
    if xm_gt_tick.numel() > 0:
        unmatched_gt_tick[unmatched_gt_tick_idx[xm_gt_tick]] = False
        unmatched_p_bar[unmatched_p_bar_idx[xm_p_bar]] = False

    fn_bar = int(unmatched_gt_bar.sum())
    fn_tick = int(unmatched_gt_tick.sum())
    fp_bar = int(unmatched_p_bar.sum())
    fp_tick = int(unmatched_p_tick.sum())

    return tp_bar, bar_as_tick, fn_bar, fp_bar, tp_tick, tick_as_bar, fn_tick, fp_tick


def f1(
    precision: float,
    recall: float
) -> float:
    """
    Computes the F1 score from precision and recall.

    :param precision: precision value
    :param recall: recall value
    :return: F1 score
    """
    denominator = precision + recall
    if denominator <= 0:
        return 0.0

    numerator = 2.0 * precision * recall
    return numerator / denominator