Source code for pytyche.analysis._functions

"""Polymorphic ``pt.analyze`` — the one-shot fit-and-analyze entry point.

The five primitive function forms (``fit_policy_tree``,
``thompson_allocation``, ``apply_calibration``, ``recommendation_summary``,
``evaluate_against_truth``) are the impl functions themselves, re-exported
through ``pytyche.analysis`` — the posterior methods are lazy delegates to
the same objects, so function form and method form share one code path by
construction.  Only ``analyze`` needs a wrapper, because its function form
accepts MORE types than the method: raw ``ObservedExperimentData`` is fit
first (``pt.fit`` with defaults), then analyzed.

The other function forms are deliberately NOT polymorphic — implicit
re-fit on every call would be a footgun for the sweep-over-alternatives
workflow (spec: "pt.analyze is polymorphic over data type").
"""

from __future__ import annotations

from pytyche.bcf.config import (
    BinaryBCFResult,
    ContinuousBCFResult,
    HurdleBCFResult,
)
from pytyche.contracts import AnalysisResult, ObservedExperimentData




def analyze(
    data: ObservedExperimentData
    | HurdleBCFResult
    | ContinuousBCFResult
    | BinaryBCFResult,
    *,
    max_depth: int = 3,
    min_segment_share: float = 0.10,
    n_bootstrap: int = 50,
    bootstrap_seed: int = 0,
) -> AnalysisResult:
    """Fit-if-needed, then assemble the canonical analysis output.

    Dispatches on the type of *data*:

    - ``ObservedExperimentData`` — the one-shot path: ``pt.fit(data)``
      with fit defaults (customize the fit by calling :func:`~pytyche.bcf.dispatch.fit`
      explicitly), then ``.analyze(**kwargs)`` on the fitted posterior.
    - A posterior result type — delegates to ``data.analyze(**kwargs)``
      directly.

    Keyword arguments are the embedded policy tree's hyperparameters and
    are forwarded to ``.analyze`` in both branches; none reach the fit.

    Args:
        data: Observed experiment data, or one of the three posterior
            result types.
        max_depth: Embedded policy tree depth (forwarded).
        min_segment_share: Minimum per-leaf population share (forwarded).
        n_bootstrap: Stability bootstrap count (forwarded; ``0`` skips
            stability with a ``UserWarning``).
        bootstrap_seed: Stability bootstrap seed (forwarded).

    Returns:
        :class:`~pytyche.contracts.AnalysisResult` — the SUMMARY surface;
        anything needing posterior samples goes through
        ``analysis.posterior``.

    Raises:
        TypeError: When *data* is neither observed data nor a posterior
            result type.
        ValueError: When a posterior *data* carries no observed data.
    """
    if isinstance(data, ObservedExperimentData):
        # Heavy import (jax via the fit stack) — only the one-shot path
        # pays it; `import pytyche.analysis` stays light.
        from pytyche.bcf.dispatch import fit

        data = fit(data)
    elif not isinstance(
        data, (HurdleBCFResult, ContinuousBCFResult, BinaryBCFResult)
    ):
        raise TypeError(
            "pt.analyze accepts ObservedExperimentData (fit-then-analyze) "
            "or a fitted posterior (HurdleBCFResult, ContinuousBCFResult, "
            f"or BinaryBCFResult), got {type(data).__name__}."
        )
    return data.analyze(
        max_depth=max_depth,
        min_segment_share=min_segment_share,
        n_bootstrap=n_bootstrap,
        bootstrap_seed=bootstrap_seed,
    )