from itertools import combinations

import numpy as np
import numpy.typing as npt


class PowersetConverter:
    def __init__(self, num_classes: int, max_set_size: int) -> None:
        self.num_classes = num_classes
        self.max_set_size = max_set_size
        self._mapping = self._build_mapping()

    def _build_mapping(self) -> npt.NDArray[np.float32]:
        powerset_classes = self._get_powerset_classes()
        num_powerset = len(powerset_classes)

        mapping = np.zeros((num_powerset, self.num_classes), dtype=np.float32)

        for powerset_idx, current_set in enumerate(powerset_classes):
            for class_idx in current_set:
                mapping[powerset_idx, class_idx] = 1.0

        return mapping

    def _get_powerset_classes(self) -> list[set[int]]:
        powerset_classes = []
        for set_size in range(self.max_set_size + 1):
            for current_set in combinations(range(self.num_classes), set_size):
                powerset_classes.append(set(current_set))  # noqa: PERF401
        return powerset_classes

    @property
    def num_powerset_classes(self) -> int:
        return len(self._get_powerset_classes())

    def to_multilabel(self, powerset_logprobs: npt.NDArray[np.float32], soft: bool = False) -> npt.NDArray[np.float32]:
        if soft:
            powerset_probs = np.exp(powerset_logprobs)
        else:
            num_powerset = self.num_powerset_classes
            powerset_probs = np.eye(num_powerset, dtype=np.float32)[np.argmax(powerset_logprobs, axis=-1)]

        return np.matmul(powerset_probs, self._mapping)
