from typing import Dict, List, Optional, Union
import pandas as pd
import scipy.stats as sps
from statsmodels.formula.api import ols, wls
import numpy as np

from aboba.base import BaseTest, TestResult


class CupedLinearRegressionTTest(BaseTest):
    def __init__(
        self,
        covariate_names: Optional[List[str]] = None,
        group_column: str = "group",
        value_column: str = "target",
        alpha: float = 0.05,
        center_on_control: bool = True,
        weight_column: Optional[str] = None,
        include_extra: bool = False,
        strata_column: Optional[str] = None,
        strata_weights: Optional[Union[Dict, pd.Series]] = None,
    ) -> None:
        """
        CUPED (Controlled-experiment Using Pre-Experiment Data) via linear regression.
        
        This test uses linear regression to adjust for pre-experiment covariates, reducing
        variance and increasing statistical power. The method centers covariates on the
        control group mean and estimates the treatment effect using OLS or WLS regression
        with heteroscedasticity-robust standard errors (HC3).
        
        Args:
            covariate_names (List[str], optional): List of pre-experiment covariates to adjust for.
                These should be variables measured before the experiment that correlate with
                the outcome metric.
            group_column (str): Name of column containing group assignment (A/B). Default "group".
            value_column (str): Name of column containing metric values to test. Default "target".
            alpha (float): Significance level for confidence interval. Default 0.05.
            center_on_control (bool): If True, covariates are centered by their mean in control
                group. This is recommended for variance reduction. Default True.
            weight_column (Optional[str]): Column with observation weights for weighted least
                squares regression. If None, uses ordinary least squares.
            include_extra (bool): If True, includes additional regression artifacts (parameters,
                design matrix, residuals) in TestResult.extra. Default False.
            strata_column (Optional[str]): Column containing stratum identifiers. Used together
                with strata_weights to derive per-observation WLS weights when weight_column
                is not provided.
            strata_weights (Optional[Union[Dict, pd.Series]]): Global (population) weight for
                each stratum. When strata_column is set and weight_column is None, the
                per-observation weight is computed as:
                    w_i = strata_weight[stratum_i] / count(stratum_i in group_i)
        
        Examples:
            ```python
            import pandas as pd
            import numpy as np
            from aboba.tests.cuped_lreg import CupedLinearRegressionTTest

            # Create sample data with pre-experiment covariate
            np.random.seed(42)
            n = 200
            pre_metric = np.random.normal(100, 15, n)
            
            # Control group
            group_a = pd.DataFrame({
                'target': pre_metric[:100] + np.random.normal(0, 10, 100),
                'pre_metric': pre_metric[:100],
                'group': 0
            })
            
            # Treatment group with effect
            group_b = pd.DataFrame({
                'target': pre_metric[100:] + np.random.normal(5, 10, 100),
                'pre_metric': pre_metric[100:],
                'group': 1
            })

            # Perform CUPED test
            test = CupedLinearRegressionTTest(
                covariate_names=['pre_metric'],
                value_column='target',
                group_column='group'
            )
            result = test.test([group_a, group_b], {})
            print(f"P-value: {result.pvalue:.4f}")
            print(f"Effect: {result.effect:.4f}")
            print(f"CI: [{result.effect_interval[0]:.4f}, {result.effect_interval[1]:.4f}]")
            ```
        """
        super().__init__()
        self.value_column = value_column
        self.group_column = group_column
        self.covariate_names = covariate_names or []
        self.alpha = alpha
        self.center_on_control = center_on_control
        self.weight_column = weight_column
        self.include_extra = include_extra
        self.strata_column = strata_column
        self.strata_weights = (
            pd.Series(strata_weights) if isinstance(strata_weights, dict) else strata_weights
        )

    def test(self, groups: List[pd.DataFrame]) -> TestResult:
        assert len(groups) == 2, "CupedLinearRegressionTTest expects exactly two groups"

        a_group, b_group = groups[0].copy(), groups[1].copy()

        # Derive per-observation WLS weights from strata when weight_column is absent.
        effective_weight_column = self.weight_column
        if (
            self.strata_column is not None
            and self.strata_weights is not None
            and self.weight_column is None
        ):
            for group in (a_group, b_group):
                stratum_counts = group[self.strata_column].value_counts()
                group["_strata_weight"] = (
                    group[self.strata_column].map(self.strata_weights)
                    / group[self.strata_column].map(stratum_counts)
                )
            effective_weight_column = "_strata_weight"

        data = pd.concat([a_group, b_group], ignore_index=True)

        feature_names: List[str] = [self.group_column]
        for name in self.covariate_names:
            if self.center_on_control:
                if effective_weight_column is None:
                    mean_control = a_group[name].mean()
                else:
                    assert effective_weight_column in a_group.columns, (
                        f"Weight column '{effective_weight_column}' not found in control group"
                    )
                    w = a_group[effective_weight_column].to_numpy(float)
                    x = a_group[name].to_numpy(float)
                    mean_control = float(np.average(x, weights=w))
                cname = f"{name}_c"
                data[cname] = data[name] - mean_control
                feature_names.append(cname)
            else:
                feature_names.append(name)

        formula = f"{self.value_column} ~ " + " + ".join(feature_names)

        if effective_weight_column is None:
            model = ols(formula, data=data).fit(cov_type="HC3")
        else:
            assert effective_weight_column in data.columns, (
                f"Weight column '{effective_weight_column}' not found in data"
            )
            model = wls(
                formula,
                data=data,
                weights=data[effective_weight_column],
            ).fit(cov_type="HC3")

        ef = model.params[self.group_column]
        se = model.bse[self.group_column]
        pvalue = model.pvalues[self.group_column]

        q = sps.norm.ppf(1 - self.alpha / 2)
        left_bound, right_bound  = ef - q * se, ef + q * se  
        extra = None
        if self.include_extra:
            extra = {
                "params": model.params,
                "design_matrix": model.model.exog,
                "resid": model.resid
            }

        return TestResult(pvalue=pvalue, effect=ef, effect_interval=(left_bound, right_bound),extra=extra if self.include_extra else None)
