"""Module for the Difference operation."""

from .base_operation import BaseOperation
from ..label_tensor import LabelTensor
from ..utils import check_consistency


class Difference(BaseOperation):
    r"""
    Implementation of the difference operation defined on a list of domains.

    Given two sets :math:`A` and :math:`B`, define their difference as:

    .. math::

        A \setminus B = \{x \mid x \in A \land x \not\in B\}

    For multiple sets :math:`A_1, A_2, \ldots, A_n`, define their difference as
    the set of points that belong to the first set but not to any of the
    remaining sets.

    No check is performed to ensure that the resulting domain is non-empty.

    :Example:

        >>> cartesian1 = CartesianDomain({'x': [0, 1], 'y': [0, 1]})
        >>> cartesian2 = CartesianDomain({'x': [0, 1], 'y': [0.5, 1.5]})
        >>> difference = Difference([cartesian1, cartesian2])
    """

    def is_inside(self, point, check_border=False):
        """
        Check if a point is inside the difference of the domains.

        :param LabelTensor point: The point to check.
        :param bool check_border: If ``True``, the boundary is considered inside
            the domain. Default is ``False``.
        :raises ValueError: If ``point`` is not a :class:`LabelTensor`.
        :raises ValueError: If the labels of ``point`` differ from the variables
            of the domain.
        :return: Whether the point is inside the domain or not.
        :rtype: bool
        """
        # Checks on point
        check_consistency(point, LabelTensor)
        if set(self.variables) != set(point.labels):
            raise ValueError(
                "Point labels differ from domain's dictionary labels. "
                f"Got {sorted(point.labels)}, expected {self.variables}."
            )

        # Check if the point is inside the first geometry and not in any other
        inside_first = self.geometries[0].is_inside(point, check_border)
        inside_others = any(
            g.is_inside(point, check_border) for g in self.geometries[1:]
        )

        return inside_first and not inside_others

    def sample(self, n, mode="random", variables="all"):
        """
        The sampling routine.

        .. note::

            This sampling method relies on rejection sampling. Points are drawn
            from the individual geometries, and only those that lie exclusively
            within one geometry are kept. When the exclusion domain is small
            relative to the combined area of the input domains, the method may
            become highly inefficient.

        :param int n: The number of samples to generate.
        :param str mode: The sampling method. Default is ``random``.
        :param variables: The list of variables to sample. If ``all``, all
            variables are sampled. Default is ``all``.
        :type variables: list[str] | str
        :raises AssertionError: If ``n`` is not a positive integer.
        :raises ValueError: If the sampling mode is invalid.
        :raises ValueError: If ``variables`` is neither ``all``, a string, nor a
            list/tuple of strings.
        :raises ValueError: If any of the specified variables is unknown.
        :return: The sampled points.
        :rtype: LabelTensor
        """
        # Validate sampling settings
        variables = self._validate_sampling(n, mode, variables)

        # Allocate list for samples
        samples = []

        # Sample until we have enough points
        while len(samples) < n:

            # Sample a sufficiently large number of points
            batch_size = 2 * (n - len(samples))
            pts = self.geometries[0].sample(batch_size, mode)

            # Filter points inside the intersection
            for p in pts:
                p = p.reshape(1, -1)
                p.labels = pts.labels
                if self.is_inside(p):
                    samples.append(p[variables])
                    if len(samples) >= n:
                        break

        return LabelTensor.cat(samples, dim=0)

    def partial(self):
        """
        Return the boundary of the domain resulting from the operation.

        :raises NotImplementedError: The :meth:`partial` method is not
            implemented for difference domains. Please operate on the individual
            domains instead.
        """
        raise NotImplementedError(
            "The partial method is not implemented for difference domains. "
            "Please operate on the individual domains instead."
        )
