from __future__ import annotations

from collections.abc import Sequence

import numpy as np
import torch
from scipy.optimize import linear_sum_assignment
from torch import Tensor

from .boxes import box_cxcywh_to_xyxy, generalized_box_iou_batched


def hungarian_match(
    outputs: dict[str, Tensor],
    targets: Sequence[dict[str, Tensor]],
    class_cost: float = 2.0,
    bbox_cost: float = 5.0,
    giou_cost: float = 2.0,
) -> list[tuple[Tensor, Tensor]]:
    """Match predictions to targets with a detached CPU Hungarian solve."""
    return hungarian_match_layers([outputs], targets, class_cost, bbox_cost, giou_cost)[0]


def hungarian_match_layers(
    outputs: Sequence[dict[str, Tensor]],
    targets: Sequence[dict[str, Tensor]],
    class_cost: float = 2.0,
    bbox_cost: float = 5.0,
    giou_cost: float = 2.0,
) -> list[list[tuple[Tensor, Tensor]]]:
    """Match all decoder layers with one GPU-to-CPU synchronization."""
    if not outputs:
        return []

    device = outputs[0]["pred_boxes"].device
    target_sizes = [len(target["labels"]) for target in targets]
    max_targets = max(target_sizes, default=0)
    if max_targets == 0:
        empty = torch.empty(0, dtype=torch.int64, device=device)
        return [[(empty, empty) for _ in targets] for _ in outputs]

    batch_size = len(targets)
    padded_boxes = torch.zeros(batch_size, max_targets, 4, dtype=torch.float32, device=device)
    padded_labels = torch.zeros(batch_size, max_targets, dtype=torch.int64, device=device)
    for batch_index, target in enumerate(targets):
        count = target_sizes[batch_index]
        if count > 0:
            padded_boxes[batch_index, :count] = target["boxes"]
            padded_labels[batch_index, :count] = target["labels"]

    padded_costs = []
    for layer_output in outputs:
        probabilities = layer_output["pred_logits"].sigmoid().detach()
        boxes = layer_output["pred_boxes"].detach()
        num_queries = boxes.shape[1]
        gather_index = padded_labels[:, None, :].expand(-1, num_queries, -1)
        class_term = -probabilities.gather(2, gather_index)
        bbox_term = torch.cdist(boxes, padded_boxes, p=1)
        giou_term = -generalized_box_iou_batched(
            box_cxcywh_to_xyxy(boxes), box_cxcywh_to_xyxy(padded_boxes)
        )
        cost = (class_cost * class_term + bbox_cost * bbox_term + giou_cost * giou_term).float()
        for batch_index in range(batch_size):
            padded_costs.append(cost[batch_index])

    costs = torch.stack(padded_costs).cpu().numpy()

    cpu_matches: list[tuple[np.ndarray, np.ndarray]] = []
    for cost, target_size in zip(costs, target_sizes * len(outputs), strict=True):
        if target_size == 0:
            empty = np.empty(0, dtype=np.int64)
            cpu_matches.append((empty, empty))
        else:
            cpu_matches.append(linear_sum_assignment(cost[:, :target_size]))

    lengths = [len(rows) for rows, _ in cpu_matches]
    all_rows = torch.as_tensor(
        np.concatenate([rows for rows, _ in cpu_matches]),
        dtype=torch.int64,
        device=device,
    )
    all_cols = torch.as_tensor(
        np.concatenate([cols for _, cols in cpu_matches]),
        dtype=torch.int64,
        device=device,
    )
    flat_matches = []
    offset = 0
    for length in lengths:
        flat_matches.append(
            (all_rows[offset : offset + length], all_cols[offset : offset + length])
        )
        offset += length

    batch_size = len(targets)
    return [
        flat_matches[start : start + batch_size]
        for start in range(0, len(flat_matches), batch_size)
    ]