__all__ = [
    'SegmentationEM',
]
import crowdkit.aggregation.base
import pandas
import typing


class SegmentationEM(crowdkit.aggregation.base.BaseImageSegmentationAggregator):
    """The EM algorithm for the image segmentation task.

    This method performs a categorical aggregation task for each pixel: should
    it be included to the resulting aggregate or no. This task is solved by
    the single coin Dawid-Skene algorithm. Each worker has a latent parameter
    "skill" that shows the probability of this worker to answer correctly.
    Skills and true pixels' labels are optimized by the Expectation-Maximization
    algorithm.


    Doris Jung-Lin Lee. 2018.
    Quality Evaluation Methods for Crowdsourced Image Segmentation
    <https://ilpubs.stanford.edu:8090/1161/1/main.pdf>

    Args:
        n_iter: A number of EM iterations.

    Examples:
        >>> import numpy as np
        >>> import pandas as pd
        >>> from crowdkit.aggregation import SegmentationEM
        >>> df = pd.DataFrame(
        >>>     [
        >>>         ['t1', 'p1', np.array([[1, 0], [1, 1]])],
        >>>         ['t1', 'p2', np.array([[0, 1], [1, 1]])],
        >>>         ['t1', 'p3', np.array([[0, 1], [1, 1]])]
        >>>     ],
        >>>     columns=['task', 'worker', 'segmentation']
        >>> )
        >>> result = SegmentationEM().fit_predict(df)
    Attributes:
        segmentations_ (Series): Tasks' segmentations.
            A pandas.Series indexed by `task` such that `labels.loc[task]`
            is the tasks's aggregated segmentation.
    """

    def fit(self, data: pandas.DataFrame) -> 'SegmentationEM':
        """Fit the model.
        Args:
            data (DataFrame): Workers' segmentations.
                A pandas.DataFrame containing `worker`, `task` and `segmentation` columns'.

        Returns:
            SegmentationEM: self.
        """
        ...

    def fit_predict(self, data: pandas.DataFrame) -> pandas.Series:
        """Fit the model and return the aggregated segmentations.
        Args:
            data (DataFrame): Workers' segmentations.
                A pandas.DataFrame containing `worker`, `task` and `segmentation` columns'.

        Returns:
            Series: Tasks' segmentations.
                A pandas.Series indexed by `task` such that `labels.loc[task]`
                is the tasks's aggregated segmentation.
        """
        ...

    def __init__(
        self,
        n_iter: int = 10,
        tol: float = 1e-05
    ) -> None:
        """Method generated by attrs for class SegmentationEM.
        """
        ...

    segmentations_: pandas.Series
    n_iter: int
    tol: float
    eps: float
    loss_history_: typing.List[float]
