"""Explicit evidence policies for soft symbolic distributions."""

from typing import Literal

import torch
import torch.nn.functional as F

EvidenceMode = Literal["observed", "predicted", "none"]


def clamp_evidence(
    logits: torch.Tensor,
    mask: torch.Tensor | None,
    *,
    mode: EvidenceMode,
    observed_symbols: torch.Tensor | None = None,
) -> torch.Tensor:
    """Produce soft symbols, optionally anchoring observed positions.

    Observed symbols are an explicit training input, never inferred from labels
    hidden in a batch. The caller selects the policy for its instantiation.
    """
    probabilities = F.softmax(logits, dim=-1)
    if mode == "none":
        return probabilities
    if mode not in ("observed", "predicted"):
        raise ValueError(f"Unknown evidence mode: {mode}")
    if mask is None or mask.shape != logits.shape[:-1]:
        raise ValueError("Evidence mask must match the symbolic positions")
    if mode == "observed":
        if observed_symbols is None or observed_symbols.shape != mask.shape:
            raise ValueError("Observed evidence requires symbols matching the mask")
        symbols = observed_symbols
    else:
        symbols = logits.argmax(dim=-1)
    weights = mask.unsqueeze(-1).float()
    anchors = F.one_hot(symbols, num_classes=logits.shape[-1]).float()
    return weights * anchors + (1.0 - weights) * probabilities


def restore_evidence(
    post_logits: torch.Tensor, pre_logits: torch.Tensor, mask: torch.Tensor
) -> torch.Tensor:
    """Optional token-head policy; not appropriate for a pooled task head."""
    if post_logits.shape != pre_logits.shape or mask.shape != pre_logits.shape[:-1]:
        raise ValueError("Evidence restoration requires matching token logits and mask")
    weights = mask.unsqueeze(-1).float()
    return weights * pre_logits + (1.0 - weights) * post_logits