import torch

from models.matcher import hungarian_match


def test_matcher_recovers_perfect_assignment():
    """If predicted curves exactly equal GT curves (in different query order),
    the matcher should recover the correct pred<->gt pairing."""
    torch.manual_seed(0)
    B, N, S = 1, 4, 10
    sample_ys = torch.linspace(0, 1, S)

    gt_xs = torch.rand(B, N, S)
    target_valid_mask = torch.ones(B, N, S, dtype=torch.bool)
    target_lane_valid = torch.tensor([[True, True, False, False]])

    # Build predictions as a *permuted* copy of the (valid) GT curves via exact
    # per-point fit isn't needed here -- construct predicted curve coeffs that
    # exactly reproduce gt_xs at the two valid GT rows, placed at query indices [2, 0].
    curve_coeffs = torch.zeros(B, N, 4)
    cls_logits = torch.zeros(B, N, 2)

    for q_idx, gt_idx in [(2, 0), (0, 1)]:
        # fit a cubic through the GT points exactly (deg-3 fit on 10 pts, small residual)
        ys_np = sample_ys.numpy()
        xs_np = gt_xs[0, gt_idx].numpy()
        coeffs = torch.tensor([0.0, 0.0, 0.0, 0.0])
        import numpy as np
        fit = np.polyfit(ys_np, xs_np, deg=3)
        coeffs = torch.tensor(fit, dtype=torch.float32)
        curve_coeffs[0, q_idx] = coeffs
        cls_logits[0, q_idx] = torch.tensor([-5.0, 5.0])  # confident "lane"

    matches = hungarian_match(cls_logits, curve_coeffs, gt_xs, target_valid_mask,
                               target_lane_valid, sample_ys)
    pred_idx, gt_idx = matches[0]
    pairs = set(zip(pred_idx.tolist(), gt_idx.tolist()))
    assert (2, 0) in pairs
    assert (0, 1) in pairs


def test_matcher_handles_no_gt_lanes():
    B, N, S = 1, 4, 10
    sample_ys = torch.linspace(0, 1, S)
    cls_logits = torch.zeros(B, N, 2)
    curve_coeffs = torch.zeros(B, N, 4)
    target_xs = torch.zeros(B, N, S)
    target_valid_mask = torch.zeros(B, N, S, dtype=torch.bool)
    target_lane_valid = torch.zeros(B, N, dtype=torch.bool)

    matches = hungarian_match(cls_logits, curve_coeffs, target_xs, target_valid_mask,
                               target_lane_valid, sample_ys)
    pred_idx, gt_idx = matches[0]
    assert pred_idx.numel() == 0 and gt_idx.numel() == 0