from typing import List, Optional
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
    ) -> 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.
        
        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

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

        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 self.weight_column is None:
                    mean_control = a_group[name].mean()
                else:
                    assert self.weight_column in a_group.columns, (
                        f"Weight column '{self.weight_column}' not found in control group"
                    )
                    w = a_group[self.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 self.weight_column is None:
            model = ols(formula, data=data).fit(cov_type="HC3")
        else:
            assert self.weight_column in data.columns, (
                f"Weight column '{self.weight_column}' not found in data"
            )
            model = wls(
                formula,
                data=data,
                weights=data[self.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)
