from typing import List

import numpy as np
import scipy
from deprecated import deprecated
from scipy import special

from ..utils import maskutils
from .coco import CocoDataset

__all__ = ["SemanticCoco", "SemanticCocoDataset"]


class SemanticCoco(CocoDataset):
    """
    An extension of the coco dataset to handle the output of a semantic segmentation model

    """

    def add_annotations(
        self,
        img_id: int,
        logits: np.ndarray,
        min_conf: float,
        start_index=1,
    ) -> List[int]:
        """
        add the annotation from the given logits

        Args:
            img_id (int): the id of the image to associate the annotations
            logits (np.ndarray): an array of shape [NClasses, Height, Width]
            min_conf (float): the minimum confidence used to filter the generated masks
            cats_idxs (List, optional): A list that maps the. Defaults to None.
            start_index (int, optional): the index to start generating the coco polygons.
                Normally, 0 encodes the background. Defaults to 1.

        Raises:
            ValueError: if the shape of masks is different than 2
            ValueError: if the shape of probs is different than 3

        Returns:
            List[int]: [the idx of the annotations added]
        """

        if not isinstance(logits, np.ndarray):
            raise ValueError(
                f"the mask type should be a numpy array not a {type(logits)}"
            )

        if len(logits.shape) != 3:
            raise ValueError("masks.shape should equal to 3")

        probs = special.softmax(logits, axis=0)
        masks = probs.argmax(0)

        annotation_ids = []
        for cat_idx in np.unique(masks)[start_index:]:
            mask = (masks == cat_idx).astype(np.uint8)
            conf = np.round(np.nan_to_num(probs[cat_idx, (mask > 0)].mean()), 2)

            if conf < min_conf:
                continue

            cat_id = int(cat_idx)
            if cat_id not in self.cats:
                raise ValueError(f"cats {cat_id} not in dataset categories")

            groups, n_groups = scipy.ndimage.label(mask)

            # get the groups starting from label 1
            for group_idx in range(1, n_groups + 1):
                group_mask = (groups == group_idx).astype(np.uint8)
                polygons = maskutils.mask_to_polygon(group_mask)
                if len(polygons) == 0:
                    continue

                bbox = maskutils.bbox(polygons, *masks.shape).tolist()
                # an exception is generated when the mask has less than 3 points
                area = int(maskutils.area(group_mask))
                if area == 0:
                    continue
                group_prob_mask = group_mask * probs[cat_idx]
                score = float(np.mean(group_prob_mask[group_prob_mask > 0]))

                annotation_ids.append(
                    self.add_annotation(img_id, cat_id, polygons, area, bbox, 0, score)
                )
        return annotation_ids

    def add_annotations_v2(
        self,
        img_id: int,
        probs: np.ndarray,
        min_conf: float,
        start_index=1,
        largest_group_only: bool = False,
    ) -> List[int]:
        """
        add the annotation from the given logits

        Args:
            img_id (int): the id of the image to associate the annotations
            probs (np.ndarray): an array of shape [NClasses, Height, Width] that contains a vector with probabilities
            min_conf (float): the minimum confidence used to filter the generated masks
            cats_idxs (List, optional): A list that maps the. Defaults to None.
            start_index (int, optional): the index to start generating the coco polygons.
                Normally, 0 encodes the background. Defaults to 1.

        Raises:
            ValueError: if the shape of masks is different than 2
            ValueError: if the shape of probs is different than 3

        Returns:
            List[int]: [the idx of the annotations added]
        """

        if not isinstance(probs, np.ndarray):
            raise ValueError(
                f"the mask type should be a numpy array not a {type(probs)}"
            )

        if len(probs.shape) != 3:
            raise ValueError("masks.shape should equal to 3")

        masks = probs.argmax(0)

        annotation_ids = []
        for cat_idx in np.unique(masks)[start_index:]:
            mask = (masks == cat_idx).astype(np.uint8)

            cat_id = int(cat_idx)
            if cat_id not in self.cats:
                raise ValueError(f"cats {cat_id} not in dataset categories")

            groups, n_groups = scipy.ndimage.label(mask)
            group_to_consider = list(range(1, n_groups + 1))
            if largest_group_only:
                values, counts = np.unique(groups, return_counts=True)
                group_to_consider = [values[np.argmax(counts[1:]) + 1]]

            # get the groups starting from label 1
            for group_idx in group_to_consider:
                group_mask = (groups == group_idx).astype(np.uint8)
                if group_mask.sum() == 0:
                    continue

                polygons = maskutils.mask_to_polygon(group_mask)
                if len(polygons) == 0:
                    continue

                bbox = maskutils.bbox_from_mask(group_mask)
                if bbox[0] is None:
                    continue

                # FIXME can have problems
                # bbox = maskutils.bbox(polygons, *masks.shape).tolist()
                # an exception is generated when the mask has less than 3 points
                area = int(maskutils.area(group_mask))
                if area == 0:
                    continue
                group_prob_mask = group_mask * probs[cat_idx]
                conf = float(np.median(group_prob_mask[group_prob_mask > 0]))
                if conf < min_conf:
                    continue

                annotation_ids.append(
                    self.add_annotation(img_id, cat_id, polygons, area, bbox, 0, conf)
                )
        return annotation_ids


@deprecated(version="0.9.34", reason="you should use SemanticCoco")
class SemanticCocoDataset(SemanticCoco):
    def __init__(self, coco_path: str, image_path: str = None) -> None:
        super().__init__(coco_path, image_path)
