"""Hungarian bipartite matching between predicted queries and GT lanes.

Replaces the paper's vague "distance-based positive/negative" anchor rule. This is
load-bearing for training (not an optional improvement): DETR-style architectures
with fixed slot assignment typically fail to train due to the permutation problem.
"""
from __future__ import annotations

import torch
import torch.nn.functional as F
from scipy.optimize import linear_sum_assignment

from models.head import eval_curve


@torch.no_grad()
def hungarian_match(
    cls_logits: torch.Tensor,       # (B, N, 2)
    curve_coeffs: torch.Tensor,     # (B, N, 4)
    target_xs: torch.Tensor,        # (B, N, S)
    target_valid_mask: torch.Tensor,  # (B, N, S) bool
    target_lane_valid: torch.Tensor,  # (B, N) bool
    sample_ys: torch.Tensor,        # (S,)
    cls_weight: float = 1.0,
    reg_weight: float = 5.0,
) -> list[tuple[torch.Tensor, torch.Tensor]]:
    """Returns, per batch item, (pred_indices, gt_indices) of the matched pairs."""
    B, N, _ = cls_logits.shape
    probs = F.softmax(cls_logits, dim=-1)[..., 1]  # (B, N) prob of "lane"

    pred_xs = eval_curve(curve_coeffs, sample_ys.view(1, 1, -1).expand(B, N, -1))  # (B, N, S)

    results = []
    for b in range(B):
        valid_gt = target_lane_valid[b].nonzero(as_tuple=True)[0]  # (M,)
        if valid_gt.numel() == 0:
            results.append((torch.empty(0, dtype=torch.long), torch.empty(0, dtype=torch.long)))
            continue

        gt_xs = target_xs[b, valid_gt]            # (M, S)
        gt_mask = target_valid_mask[b, valid_gt]  # (M, S)

        # (N, M) cost matrix
        diff = (pred_xs[b].unsqueeze(1) - gt_xs.unsqueeze(0)).abs()  # (N, M, S)
        mask = gt_mask.unsqueeze(0).float()                          # (1, M, S)
        denom = mask.sum(dim=-1).clamp(min=1.0)                      # (N broadcast, M)
        reg_cost = (diff * mask).sum(dim=-1) / denom                 # (N, M)

        cls_cost = -probs[b].unsqueeze(1).expand(-1, valid_gt.numel())  # (N, M)

        cost = cls_weight * cls_cost + reg_weight * reg_cost
        cost_np = cost.cpu().numpy()

        pred_idx, gt_idx_local = linear_sum_assignment(cost_np)
        gt_idx = valid_gt[gt_idx_local]

        results.append((torch.as_tensor(pred_idx, dtype=torch.long),
                         torch.as_tensor(gt_idx, dtype=torch.long)))
    return results