import pandas as pd
import numpy as np
from typing import Optional, List

from aboba.base import BaseDataSampler


class UserSampler(BaseDataSampler):
    """
    A sampler that creates experimental groups by randomly selecting unique users.

    When ``group_column`` is None (default) the sampler treats the input as a single
    pool, randomly splits the unique users into two groups, and returns all rows
    belonging to each group.

    When ``group_column`` is provided the sampler preserves the original behaviour:
    it splits users by the pre-existing group label and draws ``size`` unique users
    from each label independently.

    Attributes:
        group_column (Optional[str]): Column that defines the pre-existing grouping.
            Leave as None to use random user-level assignment.
        user_column (str): Column containing unique user identifiers. Default 'user_id'.
        size (Optional[int]): Number of unique users to sample per group.
            If None, all available users are used.
    """

    def __init__(
        self,
        size: Optional[int] = None,
        user_column: str = "user_id",
        group_column: Optional[str] = None,
    ):
        self.group_column = group_column
        self.user_column = user_column
        self.size = size

    def sample(self, data: pd.DataFrame, artefacts: dict) -> List[pd.DataFrame]:
        if self.group_column is None:
            return self._sample_random(data)
        return self._sample_by_group(data)

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _sample_random(self, data: pd.DataFrame) -> List[pd.DataFrame]:
        """Randomly assign unique users to two groups."""
        all_users = data[self.user_column].unique()

        if self.size is not None:
            assert len(all_users) >= 2 * self.size, (
                f"Not enough unique users ({len(all_users)}) to sample "
                f"2 × {self.size} without overlap"
            )
            chosen = np.random.choice(all_users, size=2 * self.size, replace=False)
            users_a, users_b = chosen[: self.size], chosen[self.size :]
        else:
            shuffled = np.random.permutation(all_users)
            mid = len(shuffled) // 2
            users_a, users_b = shuffled[:mid], shuffled[mid:]

        return [
            data[data[self.user_column].isin(users_a)],
            data[data[self.user_column].isin(users_b)],
        ]

    def _sample_by_group(self, data: pd.DataFrame) -> List[pd.DataFrame]:
        """Sample users from each pre-labelled group independently."""
        assert self.group_column in data.columns, (
            f"group_column '{self.group_column}' not found in data"
        )
        grouped = list(data.groupby(by=self.group_column, sort=True))

        if self.size is None:
            return [df for _, df in grouped]

        return [
            group_df[
                group_df[self.user_column].isin(
                    np.random.choice(
                        group_df[self.user_column].unique(),
                        size=self.size,
                        replace=False,
                    )
                )
            ]
            for _, group_df in grouped
        ]
