"""A fail-closed numerical stability harness for small logit vectors.

The stable path validates a content-bound precision contract, rounds inputs to
their declared IEEE interchange format, accumulates in Python binary64, and
rejects silent overflow or underflow. Unsafe helpers are explicitly demo-only.
"""

from __future__ import annotations

from dataclasses import asdict, dataclass
import hashlib
import json
import math
import re
import struct
from typing import Any


_IDENTIFIER = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._:/@-]{0,127}$")
_STRUCT_FORMAT = {
    "binary16": ">e",
    "binary32": ">f",
    "binary64": ">d",
}


def _require_identifier(name: str, value: str) -> None:
    if not isinstance(value, str) or not _IDENTIFIER.fullmatch(value):
        raise ValueError(f"{name} must be a non-empty stable identifier")


def _require_owner(name: str, value: str) -> None:
    _require_identifier(name, value)
    if not value.startswith("team:"):
        raise ValueError(f"{name} must name an accountable team: owner")


def _require_finite(name: str, value: float) -> None:
    if isinstance(value, bool) or not isinstance(value, (int, float)):
        raise TypeError(f"{name} must be a real number")
    if not math.isfinite(value):
        raise ValueError(f"{name} must be finite")


def _canonical(value: Any) -> Any:
    if isinstance(value, dict):
        return {key: _canonical(item) for key, item in value.items()}
    if isinstance(value, tuple):
        return [_canonical(item) for item in value]
    return value


def _content_id(prefix: str, payload: dict[str, Any]) -> str:
    encoded = json.dumps(
        _canonical(payload), sort_keys=True, separators=(",", ":"), allow_nan=False
    ).encode("utf-8")
    return f"{prefix}@sha256:{hashlib.sha256(encoded).hexdigest()}"


def _round_to_format(value: float, dtype: str, *, reject_underflow: bool) -> float:
    if dtype not in _STRUCT_FORMAT:
        raise ValueError(f"unsupported dtype: {dtype}")
    _require_finite("value", value)
    try:
        rounded = struct.unpack(
            _STRUCT_FORMAT[dtype], struct.pack(_STRUCT_FORMAT[dtype], value)
        )[0]
    except (OverflowError, struct.error) as error:
        raise FloatingPointError(f"{dtype} representation overflow") from error
    if not math.isfinite(rounded):
        raise FloatingPointError(f"{dtype} representation overflow")
    if reject_underflow and value != 0.0 and rounded == 0.0:
        raise FloatingPointError(f"{dtype} representation underflow")
    return rounded


@dataclass(frozen=True)
class PrecisionContract:
    contract_version: str
    input_dtype: str
    accumulation_dtype: str
    output_dtype: str
    max_vector_length: int
    probability_sum_tolerance: float
    underflow_policy: str
    runtime_contract: str
    numerical_owner: str
    semantic_id: str
    allowed_source_id: str

    def __post_init__(self) -> None:
        _require_identifier("contract_version", self.contract_version)
        if self.input_dtype not in _STRUCT_FORMAT:
            raise ValueError("input_dtype must be binary16, binary32, or binary64")
        if self.accumulation_dtype != "binary64":
            raise ValueError(
                "this implementation requires accumulation_dtype=binary64"
            )
        if self.output_dtype != "binary64":
            raise ValueError("this implementation requires output_dtype=binary64")
        if (
            isinstance(self.max_vector_length, bool)
            or not isinstance(self.max_vector_length, int)
            or self.max_vector_length <= 0
        ):
            raise ValueError("max_vector_length must be a positive integer")
        _require_finite(
            "probability_sum_tolerance", self.probability_sum_tolerance
        )
        if not 0.0 < self.probability_sum_tolerance <= 1e-6:
            raise ValueError(
                "probability_sum_tolerance must be positive and at most 1e-6"
            )
        if self.underflow_policy != "reject":
            raise ValueError("underflow_policy must be reject")
        if self.runtime_contract != "python-float-binary64":
            raise ValueError("unsupported runtime_contract")
        _require_owner("numerical_owner", self.numerical_owner)
        _require_identifier("semantic_id", self.semantic_id)
        _require_identifier("allowed_source_id", self.allowed_source_id)

    @property
    def contract_id(self) -> str:
        return _content_id(self.contract_version, asdict(self))


@dataclass(frozen=True)
class NumericVector:
    contract_id: str
    vector_id: str
    source_id: str
    dtype: str
    values: tuple[float, ...]


def validate_vector(
    contract: PrecisionContract, vector: NumericVector
) -> tuple[float, ...]:
    if type(contract) is not PrecisionContract:
        raise TypeError("contract must be a concrete PrecisionContract")
    if type(vector) is not NumericVector:
        raise TypeError("vector must be a concrete NumericVector")
    if vector.contract_id != contract.contract_id:
        raise ValueError("vector contract_id does not match the precision contract")
    _require_identifier("vector_id", vector.vector_id)
    _require_identifier("source_id", vector.source_id)
    if vector.source_id != contract.allowed_source_id:
        raise ValueError("vector source_id does not match contract allowed_source_id")
    if vector.dtype != contract.input_dtype:
        raise ValueError("vector dtype does not match contract input_dtype")
    if not isinstance(vector.values, tuple):
        raise TypeError("values must be an immutable tuple")
    if not vector.values:
        raise ValueError("values must not be empty")
    if len(vector.values) > contract.max_vector_length:
        raise ValueError("values exceed max_vector_length")
    rounded: list[float] = []
    for index, value in enumerate(vector.values):
        try:
            rounded.append(
                _round_to_format(value, contract.input_dtype, reject_underflow=True)
            )
        except (TypeError, ValueError, FloatingPointError) as error:
            raise type(error)(f"values[{index}]: {error}") from error
    return tuple(rounded)


def _shifted_exponentials(values: tuple[float, ...]) -> tuple[float, tuple[float, ...]]:
    shift = max(values)
    exponentials: list[float] = []
    for index, value in enumerate(values):
        shifted = value - shift
        if not math.isfinite(shifted):
            raise FloatingPointError(f"shift overflow at index {index}")
        term = math.exp(shifted)
        if not math.isfinite(term):
            raise FloatingPointError(f"exponential overflow at index {index}")
        if shifted != 0.0 and term == 0.0:
            raise FloatingPointError(f"exponential underflow at index {index}")
        exponentials.append(term)
    return shift, tuple(exponentials)


def stable_logsumexp(contract: PrecisionContract, vector: NumericVector) -> float:
    values = validate_vector(contract, vector)
    shift, exponentials = _shifted_exponentials(values)
    denominator = math.fsum(exponentials)
    if not math.isfinite(denominator) or denominator <= 0.0:
        raise FloatingPointError("log-sum-exp denominator is not positive and finite")
    result = shift + math.log(denominator)
    if not math.isfinite(result):
        raise FloatingPointError("log-sum-exp result overflow")
    return result


def stable_softmax(
    contract: PrecisionContract, vector: NumericVector
) -> tuple[float, ...]:
    values = validate_vector(contract, vector)
    _, exponentials = _shifted_exponentials(values)
    denominator = math.fsum(exponentials)
    if not math.isfinite(denominator) or denominator <= 0.0:
        raise FloatingPointError("softmax denominator is not positive and finite")
    probabilities = tuple(term / denominator for term in exponentials)
    for index, probability in enumerate(probabilities):
        if not math.isfinite(probability):
            raise FloatingPointError(f"softmax output is non-finite at index {index}")
        if probability == 0.0 and exponentials[index] != 0.0:
            raise FloatingPointError(f"softmax division underflow at index {index}")
    total = math.fsum(probabilities)
    if abs(total - 1.0) > contract.probability_sum_tolerance:
        raise FloatingPointError("softmax probabilities do not sum to one")
    return probabilities


def negative_log_likelihood_from_logits(
    contract: PrecisionContract,
    vector: NumericVector,
    target_index: int,
) -> float:
    values = validate_vector(contract, vector)
    if isinstance(target_index, bool) or not isinstance(target_index, int):
        raise TypeError("target_index must be an integer")
    if not 0 <= target_index < len(values):
        raise ValueError("target_index is outside the logit vector")
    shift, exponentials = _shifted_exponentials(values)
    denominator = math.fsum(exponentials)
    if not math.isfinite(denominator) or denominator <= 0.0:
        raise FloatingPointError("likelihood denominator is not positive and finite")
    # This form avoids subtracting two large, nearly equal values.
    result = math.log(denominator) - (values[target_index] - shift)
    if not math.isfinite(result):
        raise FloatingPointError("negative log likelihood is non-finite")
    if result < -contract.probability_sum_tolerance:
        raise FloatingPointError("negative log likelihood violated its lower bound")
    return max(0.0, result)


def stable_sum(contract: PrecisionContract, vector: NumericVector) -> float:
    values = validate_vector(contract, vector)
    result = math.fsum(values)
    if not math.isfinite(result):
        raise FloatingPointError("reduction overflow")
    return result


def unsafe_softmax_for_demo(values: tuple[float, ...]) -> tuple[float, ...]:
    """Textbook expression retained only to reproduce an overflow failure."""

    exponentials = tuple(math.exp(value) for value in values)
    denominator = sum(exponentials)
    return tuple(value / denominator for value in exponentials)


def naive_quantized_sum_for_demo(values: tuple[float, ...], dtype: str) -> float:
    """Sequentially round the accumulator; this intentionally demonstrates loss."""

    if not isinstance(values, tuple) or not values:
        raise ValueError("demo values must be a non-empty tuple")
    accumulator = 0.0
    for value in values:
        rounded = _round_to_format(value, dtype, reject_underflow=False)
        accumulator = _round_to_format(
            accumulator + rounded, dtype, reject_underflow=False
        )
    return accumulator


ILLUSTRATIVE_CONTRACT = PrecisionContract(
    contract_version="numerical-stability-v1",
    input_dtype="binary64",
    accumulation_dtype="binary64",
    output_dtype="binary64",
    max_vector_length=8192,
    probability_sum_tolerance=1e-12,
    underflow_policy="reject",
    runtime_contract="python-float-binary64",
    numerical_owner="team:numerical-reliability",
    semantic_id="illustrative-classification-logits-v1",
    allowed_source_id="illustrative-fixture-v1",
)

ILLUSTRATIVE_LOGITS = NumericVector(
    contract_id=ILLUSTRATIVE_CONTRACT.contract_id,
    vector_id="illustrative-logits-001",
    source_id="illustrative-fixture-v1",
    dtype="binary64",
    values=(1000.0, 1001.0, 1002.0),
)


def format_example() -> str:
    try:
        unsafe_softmax_for_demo(ILLUSTRATIVE_LOGITS.values)
    except OverflowError:
        unsafe_result = "OVERFLOW"
    else:
        unsafe_result = "unexpected_success"

    probabilities = stable_softmax(ILLUSTRATIVE_CONTRACT, ILLUSTRATIVE_LOGITS)
    logsumexp = stable_logsumexp(ILLUSTRATIVE_CONTRACT, ILLUSTRATIVE_LOGITS)
    nll = negative_log_likelihood_from_logits(
        ILLUSTRATIVE_CONTRACT, ILLUSTRATIVE_LOGITS, 2
    )
    reduction_vector = NumericVector(
        contract_id=ILLUSTRATIVE_CONTRACT.contract_id,
        vector_id="illustrative-reduction-001",
        source_id="illustrative-fixture-v1",
        dtype="binary64",
        values=(1.0,) * 4096,
    )
    naive = naive_quantized_sum_for_demo(reduction_vector.values, "binary16")
    repaired = stable_sum(ILLUSTRATIVE_CONTRACT, reduction_vector)
    return "\n".join(
        [
            "example=illustrative_only",
            f"contract_version={ILLUSTRATIVE_CONTRACT.contract_version}",
            f"unsafe_softmax={unsafe_result}",
            "stable_softmax=" + ",".join(f"{value:.6f}" for value in probabilities),
            f"logsumexp={logsumexp:.6f}",
            f"nll_target_2={nll:.6f}",
            f"naive_binary16_sum={naive:.1f}",
            f"stable_binary64_sum={repaired:.1f}",
        ]
    )


if __name__ == "__main__":
    print(format_example())