# Copyright (c) OpenMMLab. All rights reserved.
from typing import List, Union

import torch
from mmengine import ConfigDict
from mmengine.structures import InstanceData
from scipy.optimize import linear_sum_assignment
from torch.cuda.amp import autocast

from mmseg.registry import TASK_UTILS
from .base_assigner import BaseAssigner


@TASK_UTILS.register_module()
class HungarianAssigner(BaseAssigner):
    """Computes one-to-one matching between prediction masks and ground truth.

    This class uses bipartite matching-based assignment to computes an
    assignment between the prediction masks and the ground truth. The
    assignment result is based on the weighted sum of match costs. The
    Hungarian algorithm is used to calculate the best matching with the
    minimum cost. The prediction masks that are not matched are classified
    as background.

    Args:
        match_costs (ConfigDict|List[ConfigDict]): Match cost configs.
    """

    def __init__(
        self, match_costs: Union[List[Union[dict, ConfigDict]], dict,
                                 ConfigDict]
    ) -> None:

        if isinstance(match_costs, dict):
            match_costs = [match_costs]
        elif isinstance(match_costs, list):
            assert len(match_costs) > 0, \
                'match_costs must not be a empty list.'

        self.match_costs = [
            TASK_UTILS.build(match_cost) for match_cost in match_costs
        ]

    def assign(self, pred_instances: InstanceData, gt_instances: InstanceData,
               **kwargs):
        """Computes one-to-one matching based on the weighted costs.

        This method assign each query prediction to a ground truth or
        background. The assignment first calculates the cost for each
        category assigned to each query mask, and then uses the
        Hungarian algorithm to calculate the minimum cost as the best
        match.

        Args:
            pred_instances (InstanceData): Instances of model
                predictions. It includes "masks", with shape
                (n, h, w) or (n, l), and "cls", with shape (n, num_classes+1)
            gt_instances (InstanceData): Ground truth of instance
                annotations. It includes "labels", with shape (k, ),
                and "masks", with shape (k, h, w) or (k, l).

        Returns:
            matched_quiery_inds (Tensor): The indexes of matched quieres.
            matched_label_inds (Tensor): The indexes of matched labels.
        """
        # compute weighted cost
        cost_list = []
        with autocast(enabled=False):
            for match_cost in self.match_costs:
                cost = match_cost(
                    pred_instances=pred_instances, gt_instances=gt_instances)
                cost_list.append(cost)
            cost = torch.stack(cost_list).sum(dim=0)

        device = cost.device
        # do Hungarian matching on CPU using linear_sum_assignment
        cost = cost.detach().cpu()
        if linear_sum_assignment is None:
            raise ImportError('Please run "pip install scipy" '
                              'to install scipy first.')

        matched_quiery_inds, matched_label_inds = linear_sum_assignment(cost)
        matched_quiery_inds = torch.from_numpy(matched_quiery_inds).to(device)
        matched_label_inds = torch.from_numpy(matched_label_inds).to(device)

        return matched_quiery_inds, matched_label_inds