from typing import List, Optional, Union
import numpy as np
import pandas as pd
import scipy.stats as sps

from aboba.base import BaseTest, TestResult


class StratifiedTTest(BaseTest):
    def __init__(
        self,
        group_column: str,
        group_size: int,
        method: str,
        strata_columns: List[str],
        strata_weights: Union[pd.Series, dict],
        col_name: str = "target",
        alpha: float = 0.05,
    ):
        """
        Performs a stratified t-test on the data.

        This test performs a t-test while accounting for stratification in the data.
        Strata weights must be provided by the caller and represent global (population)
        proportions for each stratum — they are not inferred from the sample.

        Args:
            group_column (str): Name of the column containing group identifiers.
            group_size (int): Size of groups to sample (used externally by samplers).
            method (str): Weighting method. One of 'random', 'stratified', 'post_stratified'.
            strata_columns (List[str]): List of columns to stratify by.
            strata_weights (Union[pd.Series, dict]): Global (population) weights for each
                stratum. Keys/index must match the unique values of the strata columns.
                Will be normalised to sum to 1.
            col_name (str): Name of the column to test.
            alpha (float): Significance level for confidence interval.

        Examples:
            ```python
            import pandas as pd
            import numpy as np
            from aboba.tests.stratified_ttest import StratifiedTTest

            np.random.seed(42)
            data = pd.DataFrame({
                'group': np.repeat(['A', 'B'], 100),
                'strata': np.tile(['X', 'Y'], 100),
                'target': np.concatenate([
                    np.random.normal(10, 2, 100),
                    np.random.normal(12, 2, 100)
                ])
            })

            strata_weights = pd.Series({'X': 0.4, 'Y': 0.6})

            test = StratifiedTTest(
                group_column='group',
                group_size=50,
                method='stratified',
                strata_columns=['strata'],
                strata_weights=strata_weights,
                col_name='target',
            )

            groups = [data[data['group'] == 'A'], data[data['group'] == 'B']]
            result = test.test(groups, {})
            print(f"P-value: {result.pvalue:.4f}")
            print(f"Effect: {result.effect:.4f}")
            ```
        """
        assert method in [
            "random",
            "stratified",
            "post_stratified",
        ], f"Invalid {method = }. Must be one of 'random', 'stratified', 'post_stratified'"
        assert len(strata_columns) > 0, "Must have at least one strata column"

        super().__init__()
        self.col_name = col_name
        self.method = method
        self.strata_columns = strata_columns
        self.alpha = alpha
        self.group_column = group_column
        self.group_size = group_size

        weights = pd.Series(strata_weights) if isinstance(strata_weights, dict) else strata_weights
        self.strata_weights = weights / weights.sum()

    def test(self, groups: List[pd.DataFrame], _artefacts=None) -> TestResult:
        """
        Perform the stratified t-test on the provided groups.

        Args:
            groups (List[pd.DataFrame]): List of two DataFrames representing the groups.
            artefacts: Unused; kept for interface compatibility.

        Returns:
            TestResult: Object containing the p-value and effect size.
        """
        assert len(groups) == 2

        if self.method == "random":
            mean_function = self._simple_mean
            var_function = self._simple_var
        elif self.method == "stratified":
            mean_function = self._weighted_mean
            var_function = self._weighted_var
        elif self.method == "post_stratified":
            mean_function = self._weighted_mean
            var_function = self._weighted_post_var
        else:
            raise RuntimeError(f"method not supported {self.method = }")

        x_mean = mean_function(groups[0])
        y_mean = mean_function(groups[1])
        x_var = var_function(groups[0])
        y_var = var_function(groups[1])

        effect = x_mean - y_mean
        std = np.sqrt(x_var + y_var)
        t_stat = effect / std
        pvalue = 2 * sps.norm.sf(np.abs(t_stat))
        q = sps.norm.ppf(1 - self.alpha / 2)
        left_bound, right_bound = effect - q * std, effect + q * std

        return TestResult(pvalue, effect, effect_interval=(left_bound, right_bound))

    def _weighted_mean(self, data):
        strata_means = data.groupby(by=self.strata_columns)[self.col_name].mean()
        return (strata_means * self.strata_weights).sum()

    def _weighted_var(self, data):
        strata_vars = data.groupby(by=self.strata_columns)[self.col_name].var()
        return (strata_vars * self.strata_weights).sum() / len(data)

    def _weighted_post_var(self, data):
        strata_vars = data.groupby(by=self.strata_columns)[self.col_name].var()
        weighted_var = (strata_vars * self.strata_weights).sum() / len(data)
        post_addition = (strata_vars * (1 - self.strata_weights)).sum() / (
            len(data) ** 2
        )
        return weighted_var + post_addition

    def _simple_mean(self, data):
        return data[self.col_name].mean()

    def _simple_var(self, data):
        return data[self.col_name].var() / len(data)
