import pandas as pd
from typing import Optional, List

from aboba.base import BaseDataSampler
from .group_sampler import GroupSampler
from .random_sampler import RandomSampler


class StratifiedSampler(BaseDataSampler):
    """
    Sample groups with stratification.
    
    This sampler performs stratified sampling by dividing the data into strata
    based on specified columns, then sampling from each stratum to ensure
    proportional representation across groups.
    
    Examples:
        ```python
        import pandas as pd
        import numpy as np
        from aboba.samplers.stratified_sampler import StratifiedSampler

        # Create sample data with strata
        np.random.seed(42)
        data = pd.DataFrame({
            'group': ['A'] * 200 + ['B'] * 200,
            'stratum': ['X'] * 100 + ['Y'] * 100 + ['X'] * 100 + ['Y'] * 100,
            'value': np.random.normal(0, 1, 400)
        })

        # Sample 50 observations from each group, maintaining strata proportions
        sampler = StratifiedSampler(
            group_column='group',
            size=50,
            strata_columns=['stratum']
        )
        sampler.fit(data)  # Fit to calculate strata weights
        samples = sampler.sample(data, {})
        print(f"Number of groups: {len(samples)}")
        print(f"Size of each group: {[len(sample) for sample in samples]}")
        ```
    """

    def __init__(
        self,
        size: int,
        strata_columns: Optional[List[str]] = None,
        sample_time_weights: bool = False,
        group_column: Optional[str] = None,
    ):
        """
        Initialize the StratifiedSampler.

        Strata weights are computed from the data (or refreshed each call when
        ``sample_time_weights=True``), then each stratum is sampled independently.

        When ``group_column`` is None (default) the sampler treats the input as a
        single pool and randomly assigns observations to two groups within every
        stratum using :class:`RandomSampler`.

        When ``group_column`` is provided the sampler preserves the original
        behaviour: it splits each stratum by the pre-existing group label and
        draws ``size`` rows from each label using :class:`GroupSampler`.

        Args:
            size (int): Target number of observations per group per stratum
                (proportionally scaled by stratum weight).
            strata_columns (Optional[List[str]]): Columns to stratify by.
            sample_time_weights (bool): Recompute stratum weights on every
                ``sample()`` call instead of once during ``fit()``.
            group_column (Optional[str]): Pre-existing group label column.
                Leave as None to use random within-stratum assignment.
        """
        self.size = size
        self.group_column = group_column

        self.strata_columns = strata_columns
        if strata_columns is not None and len(strata_columns) == 1:
            self.strata_columns = strata_columns[0]

        self.sample_time_weights = sample_time_weights
        self.strata_weights = None
        self.strata_sizes = None

    def fit(self, data: pd.DataFrame):
        """
        Fit the sampler on the data to calculate strata weights.
        
        Args:
            data (pd.DataFrame): DataFrame to fit on.
        """
        self._assert_strata_columns_in_data(data)
        self._set_strata_weights(data)

    def sample(self, data: pd.DataFrame, artefacts: dict) -> List[pd.DataFrame]:
        """
        Sample data using stratified sampling.
        
        Args:
            data (pd.DataFrame): DataFrame to sample from.
            artefacts (dict): Dictionary to store additional results.
            
        Returns:
            List[pd.DataFrame]: List of DataFrames, one for each group.
        """
        self._assert_strata_columns_in_data(data)

        if self.sample_time_weights:
            self._set_strata_weights(data)

        if self.strata_weights is None or self.strata_sizes is None:
            raise RuntimeError(
                "StratifedGroupSampler must be fitted before sampling without sample_time_weights=True"
            )

        strata_result = []
        for i, (_, strata_df) in enumerate(
            data.groupby(by=self.strata_columns, sort=True)
        ):
            if not self.strata_sizes[i]:
                # TODO warning
                continue
            if self.group_column is None:
                sampler = RandomSampler(size=self.strata_sizes[i], groups_n=2)
            else:
                sampler = GroupSampler(column=self.group_column, size=self.strata_sizes[i])
            strata_result.append(sampler.sample(strata_df, artefacts))

        result = [pd.DataFrame() for _ in range(len(strata_result[0]))]
        for strata in strata_result:
            for i, grouped in enumerate(strata):
                result[i] = pd.concat([result[i], grouped])

        return result

    def _set_strata_weights(self, data: pd.DataFrame):
        strata_weights = data.groupby(by=self.strata_columns).count().iloc[:, 0].values

        self.strata_weights = strata_weights / strata_weights.sum()
        self.strata_sizes = (self.strata_weights * self.size + 0.5).astype(int)

    def _assert_strata_columns_in_data(self, data: pd.DataFrame):
        assert self.strata_columns is not None, "Strata columns must be specified"

        if isinstance(self.strata_columns, str):
            assert (
                self.strata_columns in data.columns
            ), f"Strata column '{self.strata_columns}' not in data"
        else:
            for column in self.strata_columns:
                assert column in data.columns, f"Strata column '{column}' not in data"

