src/tools/analytics/ab_test.py

"""analyze_ab_test — compare a metric across two groups + significance test (§0.9, from comparison).

Rename + reframe of `analyze_comparison` (`comparison.py`): aggregates a value for
two groups of a dimension (group_a is the baseline) and reports the gap — absolute
and percent difference, direction — AND now runs a **Welch two-sample t-test** on the
raw values to say whether the difference is statistically significant. Ships a bar
chart of the two group means with 95% error bars ("comparison + interval").

Absorbs the "significance" analysis discussed on 2026-09-15. Follows
docs/tools/ANALYTICAL_TOOLS_CONVENTIONS.md: pure synchronous compute (a t-test is
fast, not a model fit → no run_fit), Pattern A `data`, errors escape to the
never-throw seam. Deterministic.
"""

from __future__ import annotations

import math
from typing import Any

import pandas as pd
from pydantic import BaseModel, Field

from src.tools.analytics.agg_vocab import normalise_agg
from src.tools.analytics.charts import LAYOUT_PRESET, _axis, envelope
from src.tools.analytics.descriptive import ColumnNotFoundError
from src.tools.spec_builder import build_description

# How to aggregate the value within each group for the headline comparison.
SUPPORTED_AGGS = ("sum", "mean", "count", "min", "max", "median")
# Minimum numeric values per group to run a t-test.
_MIN_TEST_N = 2
# 95% normal error-bar half-width multiplier.
_Z_95 = 1.96


class UnsupportedAggregationError(ValueError):
    """The requested aggregation is not supported (maps to error_code UNSUPPORTED_AGG)."""


class GroupNotFoundError(ValueError):
    """A requested group value does not occur in the dimension column (maps to GROUP_NOT_FOUND)."""


class AbTestInput(BaseModel):
    # `data` FIRST — property order follows field order (planner-prompt readability).
    data: str = Field(..., description="Placeholder ${t<id>} of the upstream table (Pattern A).")
    dimension: str = Field(..., description="Categorical column defining the two groups.")
    value_column: str = Field(..., description="Numeric column to compare between the groups.")
    group_a: Any = Field(..., description="Baseline group value (the reference/'from').")
    group_b: Any = Field(..., description="Comparison group value (the 'to').")
    agg: str = Field("sum", description="How to aggregate the value per group (headline gap).")
    alpha: float = Field(0.05, description="Significance level for the t-test (default 0.05).")


DESCRIPTION = build_description(
    summary="A/B test: compare one metric between TWO groups (group_a is baseline) and test "
    "whether the difference is statistically significant (Welch t-test).",
    use_when="the question pits two named groups against each other, or asks if a difference "
    "is significant — 'vs'/'versus', 'bandingkan', 'A or B', 'selisih/beda antara', "
    "'signifikan', 'A/B test', 'higher/lower than'.",
    dont_use_when=[
        "it aggregates across many groups at once -> analyze_aggregate",
        "it splits a single total into shares -> analyze_pareto",
        "it tracks change over time -> analyze_trend",
    ],
    output="each group's value + absolute/percent difference, a Welch t-test (t, p-value, "
    "significant at alpha), and a bar chart of the two group means with 95% error bars.",
    examples=[
        "compare revenue between Jakarta and Surabaya",
        "is average order value significantly higher for members vs non-members?",
        "produksi bulan 1 vs bulan 2, beda signifikan?",
    ],
)


def _ab_chart(
    group_a: object,
    group_b: object,
    y_a: float | None,
    y_b: float | None,
    err_a: float | None,
    err_b: float | None,
    value_column: str,
) -> dict[str, object]:
    """Bar of the two groups' values with optional 95% error bars."""
    names = [str(group_a), str(group_b)]
    title = f"{value_column}: {names[0]} vs {names[1]}"
    trace: dict[str, object] = {"type": "bar", "x": names, "y": [y_a, y_b], "name": value_column}
    if err_a is not None and err_b is not None:
        trace["error_y"] = {"type": "data", "array": [err_a, err_b], "visible": True}
    layout = {
        **LAYOUT_PRESET,
        "title": {"text": title},
        "showlegend": False,
        "xaxis": _axis(""),
        "yaxis": _axis(value_column),
    }
    return envelope("bar", title, [trace], layout)


def _welch_test(
    a_vals: pd.Series, b_vals: pd.Series, alpha: float
) -> tuple[dict[str, object], bool]:
    """Welch two-sample t-test on the raw values. Returns (test_dict, can_test)."""
    n_a, n_b = len(a_vals), len(b_vals)
    if n_a < _MIN_TEST_N or n_b < _MIN_TEST_N:
        return (
            {
                "method": "welch_t_test",
                "n_a": n_a,
                "n_b": n_b,
                "alpha": alpha,
                "significant": None,
                "reason": f"need >= {_MIN_TEST_N} numeric values per group for a t-test",
            },
            False,
        )

    from scipy import stats  # lazy — keep module import light

    t_stat, p_value = stats.ttest_ind(a_vals, b_vals, equal_var=False)
    t_stat, p_value = float(t_stat), float(p_value)
    significant = None if math.isnan(p_value) else bool(p_value < alpha)
    test = {
        "method": "welch_t_test",
        "n_a": n_a,
        "n_b": n_b,
        "mean_a": float(a_vals.mean()),
        "mean_b": float(b_vals.mean()),
        "std_a": float(a_vals.std(ddof=1)),
        "std_b": float(b_vals.std(ddof=1)),
        "t_statistic": None if math.isnan(t_stat) else t_stat,
        "p_value": None if math.isnan(p_value) else p_value,
        "alpha": alpha,
        "significant": significant,
    }
    return test, True


def analyze_ab_test(
    df: pd.DataFrame,
    dimension: str,
    value_column: str,
    group_a: object,
    group_b: object,
    agg: str = "sum",
    alpha: float = 0.05,
) -> dict[str, object]:
    """Compare a metric between two groups and test whether the gap is significant.

    `data` is materialised upstream (Pattern A) and arrives as `df`.

    Returns a dict with: dimension, value_column, agg; group_a/value_a and
    group_b/value_b (the headline aggregate); diff_abs, diff_pct, comparison;
    `test` (Welch t-test on the raw values); and `charts` (the reserved auto-chart key).

    Raises (wrapped by the invoker's never-throw seam):
        ColumnNotFoundError, UnsupportedAggregationError, GroupNotFoundError.
    """
    missing = [c for c in (dimension, value_column) if c not in df.columns]
    if missing:
        raise ColumnNotFoundError(f"columns not found: {missing}")
    # `avg` is the QueryIR's word for this and planner.md teaches it; the tools
    # speak pandas. Accept both, and report whatever the caller actually wrote.
    # (A23 Bug 2 -- see agg_vocab.)
    requested, agg = agg, normalise_agg(agg)
    if agg not in SUPPORTED_AGGS:
        raise UnsupportedAggregationError(
            f"unsupported aggregation '{requested}'; supported: {list(SUPPORTED_AGGS)}"
        )

    rows_a = df.loc[df[dimension] == group_a, value_column]
    rows_b = df.loc[df[dimension] == group_b, value_column]
    empty = [g for g, rows in ((group_a, rows_a), (group_b, rows_b)) if rows.empty]
    if empty:
        raise GroupNotFoundError(f"no rows for group(s) {empty} in column '{dimension}'")

    value_a = float(rows_a.agg(agg))
    value_b = float(rows_b.agg(agg))
    diff_abs = value_b - value_a
    diff_pct = (diff_abs / value_a) if value_a != 0 else None
    comparison = "higher" if diff_abs > 0 else "lower" if diff_abs < 0 else "equal"

    # Significance test on the raw numeric values (compares MEANS, independent of agg).
    a_vals = pd.to_numeric(rows_a, errors="coerce").dropna()
    b_vals = pd.to_numeric(rows_b, errors="coerce").dropna()
    test, ran = _welch_test(a_vals, b_vals, alpha)

    if ran:
        y_a, y_b = test["mean_a"], test["mean_b"]
        err_a = _Z_95 * float(a_vals.std(ddof=1)) / math.sqrt(len(a_vals))
        err_b = _Z_95 * float(b_vals.std(ddof=1)) / math.sqrt(len(b_vals))
    else:
        y_a, y_b, err_a, err_b = value_a, value_b, None, None

    return {
        "dimension": dimension,
        "value_column": value_column,
        "agg": agg,
        "group_a": group_a,
        "value_a": value_a,
        "group_b": group_b,
        "value_b": value_b,
        "diff_abs": diff_abs,
        "diff_pct": diff_pct,
        "comparison": comparison,
        "test": test,
        # Reserved `quality` block (conventions §3). For an A/B comparison the
        # statistical verdict IS the quality verdict: a gap the test cannot separate
        # from noise is weak evidence no matter how large the headline difference
        # looks. `significant is None` means the test could not run (too few rows,
        # zero variance) — also weak, for the same reason. CK7 surfaces it.
        "quality": {
            "metric": "p_value",
            "value": None if test.get("p_value") is None else round(test["p_value"], 4),
            "floor": alpha,
            "weak": not test["significant"],
            "note": (
                "the difference is not distinguishable from noise at this sample size"
                if not test["significant"]
                else f"significant at alpha={alpha}"
            ),
        },
        "charts": [_ab_chart(group_a, group_b, y_a, y_b, err_a, err_b, value_column)],
    }