from typing import Any, Dict, List, Optional, Tuple, Union
import pandas as pd

from aboba.base import BaseDataProcessor
from aboba.utils import ratio_metric


def _merge_user_split_arms(
    parts: List[pd.DataFrame],
    group_column: str,
    group_control: Any,
    group_test: Any,
) -> pd.DataFrame:
    """
    Vertical concat of two arm-level frames (e.g. UserSplitter output) with explicit group labels.
    """
    if len(parts) != 2:
        raise ValueError(
            "LinearizationProcessor expects a list of exactly two DataFrames "
            "(control arm, test arm) when used after a user-level split."
        )
    merged = []
    labels = (group_control, group_test)
    for arm_df, label in zip(parts, labels):
        chunk = arm_df.copy()
        chunk[group_column] = label
        merged.append(chunk)
    return pd.concat(merged, ignore_index=True)


def _linearize_period(
    df: pd.DataFrame,
    numerator_name: str,
    denominator_name: str,
    user_name: str,
    group_column: str,
    group_control: Any,
    group_test: Any,
    result_name: str,
    eta: float,
) -> Tuple[pd.DataFrame, float]:
    control = df[df[group_column] == group_control]
    test = df[df[group_column] == group_test]
    if control.empty or test.empty:
        raise ValueError("Both control and test groups must be present")

    kappa = (1.0 - eta) * ratio_metric(
        control, numerator_name, denominator_name
    ) + eta * ratio_metric(test, numerator_name, denominator_name)

    grouped = (
        df.groupby([group_column, user_name], as_index=False)[
            [numerator_name, denominator_name]
        ]
        .sum()
    )
    grouped[result_name] = (
        grouped[numerator_name] - kappa * grouped[denominator_name]
    )
    return grouped, float(kappa)


class LinearizationProcessor(BaseDataProcessor):
    """
    Perform linearization for ratio metrics.

    Single-period mode (default): one kappa on all rows, user-level aggregation.

    Dual-period mode: when ``period_column`` is set, computes kappa separately for
    pilot and prepilot slices, merges on ``(group_column, user_name)`` so CUPED
    can use pilot linearized outcome and prepilot linearized covariate on the same rows.

    **List input:** if ``transform`` receives a list of two DataFrames (e.g. from
    ``UserSplitter``), they are merged with ``group_control`` / ``group_test`` labels by
    list order, linearized, then split again into two DataFrames ``[control, test]`` in
    the same order (so downstream tests like ``CupedLinearRegressionTTest`` need no
    ``GroupSplitter``).

    Notes:
    - If denominator_name is not provided, temporary
      denominator column filled with 1.0 is created, and removed from output after transformation.
    """

    def __init__(
        self,
        numerator_name: str = "session_lengths",
        denominator_name: Optional[str] = None,
        user_name: str = "user_id",
        group_column: str = "group",
        group_control: Any = 0,
        group_test: Any = 1,
        result_name: str = "linearized_values",
        eta: float = 0.0,
        period_column: Optional[str] = None,
        period_pilot: Any = None,
        period_prepilot: Any = None,
        prepilot_result_name: str = "linearized_prepilot",
        merge_how: str = "inner",
    ):
        super().__init__()
        self.numerator_name = numerator_name
        self.denominator_name = denominator_name
        self.user_name = user_name
        self.group_column = group_column
        self.group_control = group_control
        self.group_test = group_test
        self.result_name = result_name
        self.eta = float(eta)
        self.period_column = period_column
        self.period_pilot = period_pilot
        self.period_prepilot = period_prepilot
        self.prepilot_result_name = prepilot_result_name
        if merge_how not in ("inner", "outer"):
            raise ValueError('merge_how must be "inner" or "outer"')
        self.merge_how = merge_how
        if period_column is not None and result_name == prepilot_result_name:
            raise ValueError("result_name and prepilot_result_name must differ in dual-period mode")

    def fit(self, data: Union[pd.DataFrame, List[pd.DataFrame]]):
        """
        Fit method for pipeline compatibility.
        """
        return self

    def _package_output(
        self,
        out: pd.DataFrame,
        artifacts: Dict[str, float],
        split_input: bool,
    ) -> Tuple[Union[pd.DataFrame, List[pd.DataFrame]], Dict[str, float]]:
        if not split_input:
            return out, artifacts
        control_df = out[out[self.group_column] == self.group_control].copy()
        test_df = out[out[self.group_column] == self.group_test].copy()
        if control_df.empty or test_df.empty:
            raise ValueError("After linearization, both control and test splits must be non-empty")
        return [control_df, test_df], artifacts

    def transform(
        self, data: Union[pd.DataFrame, List[pd.DataFrame]]
    ) -> Tuple[Union[pd.DataFrame, List[pd.DataFrame]], Dict[str, float]]:
        """
        Transform raw data into user-level linearized metric values.

        If ``data`` is a list of two DataFrames, returns ``[control_df, test_df]``;
        otherwise returns a single DataFrame.

        Returns:
            Tuple[Union[pd.DataFrame, List[pd.DataFrame]], Dict]:
                - one DataFrame or two (when input was two arms)
                - artifacts: ``kappa`` or dual-period keys as documented
        """
        split_input = isinstance(data, list)
        if split_input:
            df = _merge_user_split_arms(
                data,
                self.group_column,
                self.group_control,
                self.group_test,
            )
        else:
            df = data.copy()

        eff_denom = self.denominator_name
        if eff_denom is None:
            eff_denom = "ratio_denominator"
            self.denominator_name = eff_denom

        created_temp_denominator = False
        if eff_denom not in df.columns:
            df[eff_denom] = 1.0
            created_temp_denominator = True

        required = [
            self.group_column,
            self.user_name,
            self.numerator_name,
            eff_denom,
        ]
        if self.period_column is not None:
            required.append(self.period_column)
            if self.period_pilot is None or self.period_prepilot is None:
                raise ValueError(
                    "period_pilot and period_prepilot must be set when period_column is given"
                )
            if self.period_pilot == self.period_prepilot:
                raise ValueError("period_pilot and period_prepilot must differ")

        missing = [c for c in required if c not in df.columns]
        if missing:
            raise ValueError(f"Missing required columns: {missing}")

        if self.period_column is None:
            grouped, kappa = _linearize_period(
                df,
                self.numerator_name,
                eff_denom,
                self.user_name,
                self.group_column,
                self.group_control,
                self.group_test,
                self.result_name,
                self.eta,
            )
            if created_temp_denominator:
                grouped = grouped.drop(columns=[eff_denom])
            return self._package_output(grouped, {"kappa": kappa}, split_input)

        pilot_df = df[df[self.period_column] == self.period_pilot]
        prepilot_df = df[df[self.period_column] == self.period_prepilot]
        if pilot_df.empty or prepilot_df.empty:
            raise ValueError("Both pilot and prepilot period slices must be non-empty")

        pilot_g, kappa_pilot = _linearize_period(
            pilot_df,
            self.numerator_name,
            eff_denom,
            self.user_name,
            self.group_column,
            self.group_control,
            self.group_test,
            self.result_name,
            self.eta,
        )
        prepilot_g, kappa_prepilot = _linearize_period(
            prepilot_df,
            self.numerator_name,
            eff_denom,
            self.user_name,
            self.group_column,
            self.group_control,
            self.group_test,
            self.result_name,
            self.eta,
        )
        prepilot_g = prepilot_g.rename(
            columns={self.result_name: self.prepilot_result_name}
        )

        merge_keys = [self.group_column, self.user_name]
        merged = pilot_g.merge(
            prepilot_g,
            on=merge_keys,
            how=self.merge_how,
            suffixes=("_pilot", "_prepilot"),
        )

        if created_temp_denominator:
            drop_cols = [
                f"{eff_denom}_pilot",
                f"{eff_denom}_prepilot",
            ]
            merged = merged.drop(
                columns=[c for c in drop_cols if c in merged.columns],
                errors="ignore",
            )

        artifacts: Dict[str, float] = {
            "kappa_pilot": kappa_pilot,
            "kappa_prepilot": kappa_prepilot,
            "kappa": kappa_pilot,
        }
        return self._package_output(merged, artifacts, split_input)
