from typing import Callable, Dict, Sequence, Union

import torch
from ignite.metrics.metric import Metric, reinit__is_reduced, sync_all_reduce
from monai.apps.detection.metrics.coco import COCOMetric
from monai.apps.detection.metrics.matching import matching_batch
from monai.data import box_utils

from .utils import detach_to_numpy


class IgniteCocoMetric(Metric):
    def __init__(
        self,
        coco_metric_monai: Union[None, COCOMetric] = None,
        box_key="box",
        label_key="label",
        pred_score_key="label_scores",
        output_transform: Callable = lambda x: x,
        device: Union[str, torch.device, None] = None,
        reduce_scalar: bool = True,
    ):
        r"""
        Computes coco detection metric in Ignite.

        Args:
            coco_metric_monai: the coco metric in monai.
                If not given, will asume COCOMetric(classes=[0], iou_list=[0.1], max_detection=[100])
            box_key: box key in the ground truth target dict and prediction dict.
            label_key: classification label key in the ground truth target dict and prediction dict.
            pred_score_key: classification score key in the prediction dict.
            output_transform: A callable that is used to transform the Engine’s
                process_function’s output into the form expected by the metric.
            device: specifies which device updates are accumulated on.
                Setting the metric’s device to be the same as your update arguments ensures
                the update method is non-blocking. By default, CPU.
            reduce_scalar: if True, will return the average value of coc metric values;
                if False, will return an dictionary of coc metric.

        Examples:
            To use with ``Engine`` and ``process_function``,
             simply attach the metric instance to the engine.
            The output of the engine's ``process_function`` needs to be in format of
            ``(y_pred, y)`` or ``{'y_pred': y_pred, 'y': y, ...}``.
            For more information on how metric works with :class:`~ignite.engine.engine.Engine`,
             visit :ref:`attach-engine`.
            .. include:: defaults.rst
                :start-after: :orphan:
            .. testcode::
                coco = IgniteCocoMetric()
                coco.attach(default_evaluator, 'coco')
                preds = [
                    {
                        'box': torch.Tensor([[1,1,1,2,2,2]]),
                        'label':torch.Tensor([0]),
                        'label_scores':torch.Tensor([0.8])
                    }
                ]
                target = [{'box': torch.Tensor([[1,1,1,2,2,2]]), 'label':torch.Tensor([0])}]
                state = default_evaluator.run([[preds, target]])
                print(state.metrics['coco'])
            .. testoutput::
                1.0...
        .. versionadded:: 0.4.3
        """
        self.box_key = box_key
        self.label_key = label_key
        self.pred_score_key = pred_score_key
        if coco_metric_monai is None:
            self.coco_metric = COCOMetric(classes=[0], iou_list=[0.1], max_detection=[100])
        else:
            self.coco_metric = coco_metric_monai
        self.reduce_scalar = reduce_scalar

        if device is None:
            device = torch.device("cpu")
        super(IgniteCocoMetric, self).__init__(output_transform=output_transform, device=device)

    @reinit__is_reduced
    def reset(self) -> None:
        self.val_targets_all = []
        self.val_outputs_all = []

    @reinit__is_reduced
    def update(self, output: Sequence[Dict]) -> None:
        y_pred, y = output[0], output[1]
        self.val_outputs_all += y_pred
        self.val_targets_all += y

    @sync_all_reduce("val_targets_all", "val_outputs_all")
    def compute(self) -> float:

        self.val_outputs_all = detach_to_numpy(self.val_outputs_all)
        self.val_targets_all = detach_to_numpy(self.val_targets_all)

        results_metric = matching_batch(
            iou_fn=box_utils.box_iou,
            iou_thresholds=self.coco_metric.iou_thresholds,
            pred_boxes=[val_data_i[self.box_key] for val_data_i in self.val_outputs_all],
            pred_classes=[val_data_i[self.label_key] for val_data_i in self.val_outputs_all],
            pred_scores=[val_data_i[self.pred_score_key] for val_data_i in self.val_outputs_all],
            gt_boxes=[val_data_i[self.box_key] for val_data_i in self.val_targets_all],
            gt_classes=[val_data_i[self.label_key] for val_data_i in self.val_targets_all],
        )
        val_epoch_metric_dict = self.coco_metric(results_metric)[0]

        if self.reduce_scalar:
            val_epoch_metric = val_epoch_metric_dict.values()
            val_epoch_metric = sum(val_epoch_metric) / len(val_epoch_metric)
            return val_epoch_metric
        else:
            return val_epoch_metric_dict