Inference

@dataclass(frozen=True, slots=True)
class Answer:
    """The answer to a query.

    Attributes:
        point: The point answer: the probability with every parameter at its credence mean. Because
            the network is multilinear in its independent parameters, this equals the predictive
            ratio ``E[P(A, E)] / E[P(E)]``, where each expectation is taken over the parameters'
            distributions and the ``exclusive`` constraints count as evidence on both sides. It is
            not the average of ``P(A | E)`` over the parameters, which generally differs.
        band: The 5%, 50% and 95% quantiles of the answer over parameter draws, or ``None`` when no
            draws were made because every parameter is a ``Point`` or ``draws`` was 0.
        mean_over_draws: The average of the per-draw answers, reported separately from ``point``
            because it is a different quantity; ``None`` exactly when ``band`` is.
        draws: The number of parameter draws behind ``band`` and ``mean_over_draws``; 0 when none.
    """

    point: float
    band: Band | None = None
    mean_over_draws: float | None = None
    draws: int = 0

    def __post_init__(self) -> None:
        """Check that the band, its mean and the draw count come together or not at all.

        Raises:
            ValidationError: If some but not all of them are given.
        """
        if (self.band is None) != (self.mean_over_draws is None) or (self.band is None) != (self.draws == 0):
            msg = "band, mean_over_draws and a positive draws count must be given together or not at all"
            raise ValidationError(msg)

    def to_dict(self) -> dict[str, object]:
        """Return the answer as JSON-ready data.

        Returns:
            ``{"point": ...}``, plus ``"band"``, ``"mean_over_draws"`` and ``"draws"`` when there is a band.
        """
        if self.band is None:
            return {"point": self.point}
        return {
            "point": self.point,
            "band": self.band.to_dict(),
            "mean_over_draws": self.mean_over_draws,
            "draws": self.draws,
        }

@dataclass(frozen=True, slots=True)
class Band:
    """Quantiles of a query's answer over parameter draws.

    Attributes:
        q05: The 5% quantile.
        q50: The median.
        q95: The 95% quantile.
    """

    q05: float
    q50: float
    q95: float

    def to_dict(self) -> dict[str, float]:
        """Return the band as JSON-ready data.

        Returns:
            ``{"q05": ..., "q50": ..., "q95": ...}``.
        """
        return {"q05": self.q05, "q50": self.q50, "q95": self.q95}

class Engine(ABC):
    """An inference engine.

    An engine only has to compute the probability of a partial assignment. ``query`` turns that into
    conditional and interventional answers, and observes every ``exclusive`` constraint as true.
    """

    @abstractmethod
    def probability(self, network: Network, assignment: Mapping[int, int]) -> float:
        """Compute the probability that the given variables take the given values.

        Args:
            network: The network.
            assignment: Variable indices and values (0 or 1); every other variable is summed out.

        Returns:
            The probability.
        """

    def query(self, network: Network, query: Query) -> float:
        """Answer a query.

        Args:
            network: The network, at the parameter values the answer should use.
            query: The query.

        Returns:
            ``P(target | evidence, constraints)`` in the network after the interventions.

        Raises:
            ValidationError: If the query names a node that is not an inference variable, or the
                interventions set merged nodes to different values.
            ZeroProbabilityError: If the evidence and constraints have probability zero.
        """
        if query.interventions:
            network = network.intervene(query.interventions)
        evidence = _combine(_resolve(network, query.evidence), {i: int(v) for i, v in network.constraints.items()})
        target = _resolve(network, query.target)
        if evidence is None:
            denominator = 0.0
        elif evidence:
            denominator = self.probability(network, evidence)
        else:
            denominator = 1.0
        if denominator <= 0.0 or evidence is None:
            msg = "the evidence, together with the exclusive constraints, has probability zero"
            raise ZeroProbabilityError(msg)
        joint = _combine(evidence, target)
        if joint is None:
            return 0.0
        if len(joint) == len(evidence):
            return 1.0
        answer = self.probability(network, joint) / denominator
        # Rounding can carry the ratio a few ulps past the unit interval.
        return min(1.0, max(0.0, answer))

class Enumeration(Engine):
    """Exhaustive enumeration of every world consistent with the assignment.

    Each world's probability is the product of one entry from every table, and the answer is the sum
    over worlds. It shares no code with variable elimination, which makes it the independent
    reference that engine is tested against. The cost doubles with every variable left free.
    """

    def __init__(self, max_free_variables: int = DEFAULT_MAX_FREE_VARIABLES) -> None:
        """Configure the engine.

        Args:
            max_free_variables: The largest number of summed-out variables allowed; a query needing
                more raises ``ProblemTooLargeError``.
        """
        self.max_free_variables = max_free_variables

    def probability(self, network: Network, assignment: Mapping[int, int]) -> float:
        """Sum the joint probability of every world that agrees with ``assignment``.

        Args:
            network: The network.
            assignment: Variable indices and values (0 or 1).

        Returns:
            The probability.

        Raises:
            ProblemTooLargeError: If more than ``max_free_variables`` variables would be summed out.
        """
        n = len(network)
        free = [index for index in range(n) if index not in assignment]
        if len(free) > self.max_free_variables:
            msg = (
                f"enumeration would sum over {len(free)} free variables, above the limit of "
                f"{self.max_free_variables} (max_free_variables); use variable elimination"
            )
            raise ProblemTooLargeError(msg, len(free), self.max_free_variables)
        worlds = np.zeros((2 ** len(free), n), dtype=np.intp)
        codes = np.arange(2 ** len(free))
        for bit, index in enumerate(free):
            worlds[:, index] = (codes >> bit) & 1
        for index, value in assignment.items():
            worlds[:, index] = value
        weight = np.ones(len(worlds))
        for factor in network.factors:
            weight = weight * factor.table[tuple(worlds[:, index] for index in factor.scope)]
        return float(weight.sum())

class InferenceError(CredenceGraphError):
    """Base class for errors raised while answering a query."""

class ProblemTooLargeError(InferenceError):
    """An exact engine would need more memory than its configured limit allows.

    Exact engines never fall back to an approximation on their own: they raise this error instead.

    Attributes:
        required: The size the query would need, in the unit ``limit`` is expressed in.
        limit: The configured limit.
    """

    def __init__(self, message: str, required: int, limit: int) -> None:
        """Store the sizes alongside the message.

        Args:
            message: Human-readable description naming the limit.
            required: The size the query would need.
            limit: The configured limit.
        """
        super().__init__(message)
        self.required = required
        self.limit = limit

@dataclass(frozen=True, slots=True)
class Query:
    """A probability query: ``P(target | evidence)`` after the interventions are applied.

    Attributes:
        target: Node ids and the values whose joint probability is asked for.
        evidence: Node ids and the observed values conditioned on.
        interventions: Node ids and the values they are set to; relations into them are cut.
    """

    target: Mapping[str, bool]
    evidence: Mapping[str, bool] = field(default_factory=dict)
    interventions: Mapping[str, bool] = field(default_factory=dict)

    def __post_init__(self) -> None:
        """Validate the three assignments and store read-only copies."""
        object.__setattr__(self, "target", _frozen_assignment(self.target, "target"))
        object.__setattr__(self, "evidence", _frozen_assignment(self.evidence, "evidence"))
        object.__setattr__(self, "interventions", _frozen_assignment(self.interventions, "interventions"))

class VariableElimination(Engine):
    """Exact inference by variable elimination with a min-fill order.

    The order depends only on the factor scopes once the assignment is applied, so it is computed
    once for each distinct set of scopes and reused by later queries, including those on networks
    that differ only in their parameter values.

    Before any table is built, the engine works out the scope of every intermediate factor the
    order produces. If the largest would hold more than ``max_factor_size`` entries, the query fails
    with ``ProblemTooLargeError``; the engine never falls back to an approximation.
    """

    def __init__(self, max_factor_size: int = DEFAULT_MAX_FACTOR_SIZE) -> None:
        """Configure the engine.

        Args:
            max_factor_size: The largest intermediate factor allowed, in table entries.
        """
        self.max_factor_size = max_factor_size

    def probability(self, network: Network, assignment: Mapping[int, int]) -> float:
        """Compute the probability of ``assignment`` by summing out every other variable.

        Args:
            network: The network.
            assignment: Variable indices and values (0 or 1).

        Returns:
            The probability.

        Raises:
            ProblemTooLargeError: If an intermediate factor would exceed ``max_factor_size`` entries.
        """
        factors = [_reduce(factor.scope, factor.table, assignment) for factor in network.factors]
        order = _order(tuple(scope for scope, _ in factors))
        self._check_size(factors, order)
        for variable in order:
            factors, new = _eliminate(factors, variable)
            factors.append(new)
        result = 1.0
        for _, table in factors:
            result *= float(table)
        return result

    def _check_size(self, factors: Sequence[_Factor], order: Sequence[int]) -> None:
        """Raise if eliminating in ``order`` would build a factor above the limit.

        Args:
            factors: The factors after the assignment is applied.
            order: The elimination order.

        Raises:
            ProblemTooLargeError: If an intermediate factor would exceed ``max_factor_size`` entries.
        """
        scopes = [set(scope) for scope, _ in factors]
        for variable in order:
            touching = [scope for scope in scopes if variable in scope]
            union = set().union(*touching)
            size = 2 ** len(union)
            if size > self.max_factor_size:
                msg = (
                    f"variable elimination would build an intermediate factor over {len(union)} variables "
                    f"({size} entries), above the limit of {self.max_factor_size} entries (max_factor_size). "
                    "Raise max_factor_size to run the query exactly; exact inference does not fall back to "
                    "an approximation on its own"
                )
                raise ProblemTooLargeError(msg, size, self.max_factor_size)
            scopes = [scope for scope in scopes if variable not in scope]
            scopes.append(union - {variable})

class ZeroProbabilityError(InferenceError, ValueError):
    """The evidence of a query, together with the ``exclusive`` constraints, has probability zero."""