import pandas as pd
import numpy as np

from typing import Optional, Any, List, Union

from aboba.base import BaseDataProcessor


class CupedProcessor(BaseDataProcessor):
    """
    Performs CUPED (Controlled-experiment Using Pre-Experiment Data) transformation.

    CUPED is a variance reduction technique that uses pre-experiment data (covariates)
    to improve the sensitivity of A/B tests. It adjusts the target metric using
    information from covariates that are correlated with the target metric.

    The transformation fits a linear regression Y = a + b^T X (without a treatment
    variable) and produces Z = Y - b^T X. Note that the intercept is *not* subtracted,
    so the mean of Z equals the intercept, preserving interpretability.

    Examples:
        ```python
        import pandas as pd
        import numpy as np
        from aboba.processing.cuped_processor import CupedProcessor

        # Create sample data with pre-experiment data
        np.random.seed(42)
        n = 1000
        # Pre-experiment metric (covariate)
        pre_metric = np.random.normal(100, 10, n)
        # Post-experiment metric (target) - correlated with pre-metric
        post_metric = 2 * pre_metric + np.random.normal(0, 5, n)
        # Add treatment effect for second half
        post_metric[n//2:] += 5

        data = pd.DataFrame({
            'group': ['control'] * (n//2) + ['test'] * (n//2),
            'pre_metric': pre_metric,
            'post_metric': post_metric
        })

        # Apply CUPED transformation (single covariate)
        processor = CupedProcessor(
            value_column='post_metric',
            covariate_columns='pre_metric',
            result_column='cuped_metric',
        )
        processed_data, _ = processor.transform(data)
        print(f"Original variance: {data['post_metric'].var():.2f}")
        print(f"CUPED variance: {processed_data['cuped_metric'].var():.2f}")
        ```
    """

    def __init__(
        self,
        value_column: str,
        covariate_columns: Union[str, List[str]],
        result_column: str,
        group_column: Optional[str] = None,
        group_test: Any = None,
        group_control: Any = None,
    ):
        """
        Initialize the CUPED processor.

        Args:
            value_column (str): Name of the column with the target metric values.
            covariate_columns (Union[str, List[str]]): Name(s) of the column(s) with
                covariate values. Accepts a single column name or a list of names.
            result_column (str): Name of the column to store the CUPED-transformed values.
            group_column (Optional[str]): Unused; kept for backward compatibility.
            group_test (Any): Unused; kept for backward compatibility.
            group_control (Any): Unused; kept for backward compatibility.
        """

        super().__init__()
        self.value_column = value_column
        if isinstance(covariate_columns, str):
            self.covariate_columns = [covariate_columns]
        else:
            self.covariate_columns = list(covariate_columns)
        self.result_column = result_column
        self.group_column = group_column
        self.group_test = group_test
        self.group_control = group_control

    def transform(self, data: pd.DataFrame):
        """
        Apply CUPED transformation to the data.

        Fits OLS regression Y = a + b^T X on all rows (no treatment variable),
        then computes Z = Y - b^T X (subtracts only the slope contribution).

        Args:
            data (pd.DataFrame): DataFrame to transform.

        Returns:
            Tuple[pd.DataFrame, Dict]: Transformed DataFrame and empty artifacts dict.
        """
        X = data[self.covariate_columns].to_numpy(float)
        Y = data[self.value_column].to_numpy(float)

        # Fit OLS with intercept: Y = a + b^T X
        X_with_intercept = np.column_stack([np.ones(len(X)), X])
        b, _, _, _ = np.linalg.lstsq(X_with_intercept, Y, rcond=None)
        # b[0] = intercept, b[1:] = slopes

        # Z = Y - b^T X  (intercept is not subtracted)
        Z = Y - X @ b[1:]

        data = data.copy()
        data[self.result_column] = Z

        return data, dict()
