from typing import List, Optional

import pandas as pd

from aboba.base import BaseDataSplitter
import numpy as np


class RandomSplitter(BaseDataSplitter):
    """
    Sample random groups from data.

    This splitter randomly selects specified numbers of observations from the data
    to form multiple groups. Each observation can only be selected once across
    all groups.

    Examples:
        ```python
        import pandas as pd
        import numpy as np
        from aboba.splitters.random_splitter import RandomSplitter

        # Create sample data
        np.random.seed(42)
        data = pd.DataFrame({
            'value': np.random.normal(0, 1, 100)
        })

        # Sample 2 groups of 30 observations each
        splitter = RandomSplitter(size=30, groups_n=2)
        samples = splitter.sample(data, {})
        print(f"Number of groups: {len(samples)}")
        print(f"Size of each group: {[len(sample) for sample in samples]}")
        # Check that no observation appears in both groups
        intersection = set(samples[0].index).intersection(set(samples[1].index))
        print(f"Number of shared observations: {len(intersection)}")
        ```
    """

    def __init__(self, size: int, groups_n: int = 2, add_group_column: Optional[str] = None):
        """
        Initialize the RandomSplitter.

        Args:
            size (int): Number of observations to sample for each group.
            groups_n (int): Number of groups to create. Default is 2.
            add_group_column (Optional[str]): If provided, adds a column with this name
                to each resulting DataFrame containing the group index (0-based). Default is None.
        """
        self.size = size
        self.groups_n = groups_n
        self.add_group_column = add_group_column

        assert groups_n >= 1
        assert size >= 1

    def sample(self, data: pd.DataFrame, _: dict = {}) -> List[pd.DataFrame]:
        """
        Sample random groups from data.

        Args:
            data (pd.DataFrame): DataFrame to sample from.
            _: dict: Dictionary to store additional results (not used in this splitter).

        Returns:
            List[pd.DataFrame]: List of DataFrames, one for each group.
        """
        assert len(data.index) >= self.size

        left_index = set(list(data.index))
        result = []

        for _ in range(self.groups_n):
            group = list(
                np.random.choice(list(left_index), size=self.size, replace=False)
            )
            left_index = left_index.difference(group)
            result.append(group)

        groups = [data.loc[group] for group in result]

        if self.add_group_column is not None:
            groups = [
                df.assign(**{self.add_group_column: i})
                for i, df in enumerate(groups)
            ]

        return groups
