Source code for pytyche.metrics

"""v2 metric dispatch registry.

Maps metric_id to model structure and extraction config. Single source of
truth for metric → model routing in the v2 pipeline. No v1 dependency.
"""

from __future__ import annotations

import dataclasses
from typing import Literal

from pytyche.contracts import MetricFamily




@dataclasses.dataclass(frozen=True)
class MetricConfig:
    """Maps a metric to its model, posterior extraction, and lift semantics.

    Fields:
        metric_id: Canonical metric name (e.g. ``"conversion_rate"``).
        metric_family: Abstract metric family from the contract enum.
        model_id: Which PyMC model to fit.
        posterior_prefix: Variable name prefix in the posterior
            (e.g. ``"p"`` for ``p_arm0``).
        lift_unit: Unit for lift samples (``"pct"`` for relative,
            ``"dollar"`` for absolute).
        component_prefixes: Prefixes of decomposition effect deterministics
            in the posterior (e.g. ``("frequency", "severity")`` maps to
            ``frequency_effect_arm1_vs_arm0``). Empty tuple for metrics
            without decomposition.
    """

    metric_id: str
    metric_family: MetricFamily
    model_id: Literal["binary", "hurdle_lognormal"]
    posterior_prefix: str
    lift_unit: str
    component_prefixes: tuple[str, ...]



_REGISTRY: dict[str, MetricConfig] = {
    "conversion_rate": MetricConfig(
        metric_id="conversion_rate",
        metric_family=MetricFamily.BINARY,
        model_id="binary",
        posterior_prefix="p",
        lift_unit="pct",
        component_prefixes=(),
    ),
    "revenue_per_visitor": MetricConfig(
        metric_id="revenue_per_visitor",
        metric_family=MetricFamily.HURDLE_REAL,
        model_id="hurdle_lognormal",
        posterior_prefix="rpv",
        lift_unit="dollar",
        component_prefixes=("frequency", "severity"),
    ),
}




def get_metric_config(metric: str) -> MetricConfig:
    """Look up metric config by canonical metric name.

    Raises:
        ValueError: If *metric* is not in the registry.
    """
    try:
        return _REGISTRY[metric]
    except KeyError:
        known = ", ".join(sorted(_REGISTRY))
        raise ValueError(
            f"Unknown metric {metric!r}. Known metrics: {known}"
        ) from None