"""The oracle — the neural analogue of Layer 1's "reality's no".

An oracle maps an observation to the mode it belongs to. For the fully-synthetic sandbox
the oracle is **exact** (it decodes the lossless identity segment), so the measured mode
distribution ``p_hat`` carries zero measurement noise — this is what lets a trained model's
collapse be read directly against the known ``p*``. (The MNIST tier will add a
``ClassifierOracle`` wrapping a frozen network plus its confusion matrix; that arrives with
Stage C and is confirmation-only.)
"""

from __future__ import annotations

from typing import Protocol, runtime_checkable

import numpy as np

from .config import SyntheticCfg


@runtime_checkable
class Oracle(Protocol):
    """Adjudicates which mode an observation belongs to."""

    def classify(self, X: np.ndarray) -> np.ndarray:
        """Return the length-``n`` mode index for each row of ``X``."""
        ...


class ExactOracle:
    """Zero-error oracle for the synthetic sandbox: decodes the identity segment.

    Args:
        cfg (SyntheticCfg): The synthetic configuration whose grammar produced ``X``.
    """

    def __init__(self, cfg: SyntheticCfg) -> None:
        self.cfg = cfg
        self.id_len = cfg.id_len
        self.base = cfg.id_base
        self.K = cfg.K
        # Positional weights for base-`base` decoding, most-significant digit first.
        self._weights = self.base ** np.arange(self.id_len - 1, -1, -1, dtype=np.int64)

    def classify(self, X: np.ndarray) -> np.ndarray:
        """Decode mode indices from the identity segment of each observation.

        Args:
            X (np.ndarray): Token sequences of shape ``(n, seq_len)``.

        Returns:
            np.ndarray: Length-``n`` decoded indices. Grammar-valid data always decodes into
            ``[0, K)``; a *neural* model may emit an invalid codeword that decodes to
            ``>= K`` — such samples are dropped by :func:`measure_distribution` rather than
            being clipped onto a real mode (which would bias ``p_hat``).
        """
        X = np.asarray(X, dtype=np.int64)
        ident = X[:, : self.id_len]
        return ident @ self._weights


def measure_distribution(X: np.ndarray, oracle: Oracle, K: int) -> np.ndarray:
    """Measure the empirical mode distribution ``p_hat`` of a sample.

    This is the neural readout of ``p_t``: classify every observation and normalise the
    mode histogram. Modes absent from ``X`` receive zero mass (support shrinks exactly as
    in Layer 1). Invalid codewords (decoded index ``>= K``, only producible by a neural
    model) are dropped, so ``p_hat`` is renormalised over grammar-valid samples.

    Args:
        X (np.ndarray): Token sequences of shape ``(n, seq_len)``.
        oracle (Oracle): The mode adjudicator.
        K (int): Number of modes.

    Returns:
        np.ndarray: Length-``K`` probability vector summing to 1.
    """
    modes = oracle.classify(X)
    counts = np.bincount(modes, minlength=K)[:K].astype(float)
    total = counts.sum()
    if total <= 0:
        raise ValueError("measure_distribution received an empty sample")
    return counts / total