src/tools/analytics/segmentation.py

"""analyze_segmentation — bucket rows into rule-based segments (§0.9, from analyze_segment).

Activation + rename of `analyze_segment`: bins a NUMERIC column into segments
(explicit cut edges, or equal-frequency quantile buckets) and reports how rows
distribute across them (count, and optionally an aggregate of another column per
segment). Now ships a bar chart as a package.

Rule-based on purpose — the caller supplies the boundaries/bucket count. That is
the deliberate contrast with analyze_cluster, which discovers the groups itself.

Follows docs/tools/ANALYTICAL_TOOLS_CONVENTIONS.md: pure synchronous compute (no
model fit → stays sync), Pattern A `data`, errors escape to the invoker's
never-throw seam. Deterministic.
"""

from __future__ import annotations

import math

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

# Binning strategies.
SUPPORTED_METHODS = ("edges", "quantile")

# How to aggregate the value column within each segment.
SUPPORTED_AGGS = ("sum", "mean", "count", "min", "max", "median")


class InvalidMethodError(ValueError):
    """The requested binning method is not supported (maps to INVALID_METHOD)."""


class NonNumericColumnError(ValueError):
    """The column to segment on is not numeric (maps to NON_NUMERIC_COLUMN)."""


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


def _clean(value: object) -> object:
    """Convert numpy scalars to plain Python; NaN -> None for JSON-clean output."""
    if value is None:
        return None
    if hasattr(value, "item"):
        value = value.item()
    if isinstance(value, float) and math.isnan(value):
        return None
    return value


class SegmentationInput(BaseModel):
    # `data` FIRST — property order follows field order (planner-prompt readability).
    data: str = Field(..., description="Placeholder ${t<id>} of the upstream table (Pattern A).")
    column: str = Field(..., description="Numeric column to bin into segments.")
    bins: list[float] | int = Field(
        ..., description="Cut edges list (method 'edges', e.g. [0,18,35,60]) OR bucket count "
        "(method 'quantile', e.g. 4 for quartiles)."
    )
    method: str = Field("edges", description="'edges' (explicit boundaries) or 'quantile'.")
    labels: list[str] | None = Field(None, description="Optional segment names, in bin order.")
    value_column: str | None = Field(None, description="Optional column to aggregate per segment.")
    agg: str = Field("sum", description="How to aggregate value_column (sum/mean/count/...).")


DESCRIPTION = build_description(
    summary="Bin a numeric column into rule-based segments and count how rows distribute "
    "(optionally aggregating another column per segment).",
    use_when="the user gives the boundaries/bucket count to bracket a continuous number "
    "into ranges — 'segmen', 'kelompokkan ke rentang', 'kelompok umur/tingkatan', "
    "'quartile/decile', 'bins/tiers'.",
    dont_use_when=[
        "the groups are UNKNOWN and should be discovered -> analyze_cluster",
        "the category already exists (no binning) -> analyze_pareto",
        "it aggregates by an existing key -> analyze_aggregate",
    ],
    output="a table of segments with row count (and optional per-segment aggregate) plus a "
    "bar chart across the segments.",
    examples=[
        "split customers into age brackets 0-18, 18-35, 35-60",
        "bucket orders into value tiers",
        "divide users into spending quartiles",
    ],
)


def _segment_chart(
    segments: list[dict[str, object]], column: str, value_column: str | None, agg: str
) -> dict[str, object]:
    """Bar chart across segments: the per-segment aggregate when given, else row count."""
    names = [s["segment"] for s in segments]
    if value_column is not None:
        ys = [s.get("value") for s in segments]
        y_label, title = f"{agg}({value_column})", f"{agg} of {value_column} by segment"
    else:
        ys = [s["count"] for s in segments]
        y_label, title = "count", f"Row count by segment of {column}"
    layout = {
        **LAYOUT_PRESET,
        "title": {"text": title},
        "showlegend": False,
        "xaxis": _axis("segment"),
        "yaxis": _axis(y_label),
    }
    return envelope("bar", title, [{"type": "bar", "x": names, "y": ys, "name": y_label}], layout)


def analyze_segmentation(
    df: pd.DataFrame,
    column: str,
    bins: list[float] | int,
    method: str = "edges",
    labels: list[str] | None = None,
    value_column: str | None = None,
    agg: str = "sum",
) -> dict[str, object]:
    """Segment rows by binning a numeric column (rule-based).

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

    Returns a dict with: column, method, `segments` ([{segment, count, (value)}] in
    bin order), agg (only when value_column is given), and `charts` (the reserved
    auto-chart key).

    Raises (wrapped by the invoker's never-throw seam):
        ColumnNotFoundError, NonNumericColumnError, InvalidMethodError,
        UnsupportedAggregationError.
    """
    referenced = [column] + ([value_column] if value_column else [])
    missing = [c for c in referenced if c not in df.columns]
    if missing:
        raise ColumnNotFoundError(f"columns not found: {missing}")
    if not pd.api.types.is_numeric_dtype(df[column]):
        raise NonNumericColumnError(f"column '{column}' is not numeric")
    if method not in SUPPORTED_METHODS:
        raise InvalidMethodError(f"unknown method '{method}'; supported: {list(SUPPORTED_METHODS)}")
    # See agg_vocab: `avg` (QueryIR's word) is accepted as `mean`. Normalised
    # unconditionally so the value used downstream matches the one validated.
    requested, agg = agg, normalise_agg(agg)
    if value_column is not None and agg not in SUPPORTED_AGGS:
        raise UnsupportedAggregationError(
            f"unsupported aggregation '{requested}'; supported: {list(SUPPORTED_AGGS)}"
        )

    if method == "edges":
        cats = pd.cut(df[column], bins=bins, labels=labels, include_lowest=True)
    else:  # quantile
        cats = pd.qcut(df[column], q=bins, labels=labels, duplicates="drop")

    grouped = df.groupby(cats, observed=False)
    counts = grouped.size()

    segments: list[dict[str, object]] = []
    for seg in counts.index:
        row: dict[str, object] = {"segment": str(seg), "count": int(counts[seg])}
        if value_column is not None:
            row["value"] = _clean(grouped[value_column].agg(agg).get(seg))
        segments.append(row)

    out: dict[str, object] = {
        "column": column,
        "method": method,
        "segments": segments,
        "charts": [_segment_chart(segments, column, value_column, agg)],
    }
    if value_column is not None:
        out["agg"] = agg
    return out