"""Padding and inference helpers for caller-batch-independent verification.

Padding masks and no-drop routing settings address semantic batch dependence.
They do not guarantee bitwise equality: kernels can round differently for
changed shapes. Use ``strict_deterministic_inference`` when fixed per-example
shapes matter, with an eval model whose kernels and routing are deterministic.
"""
from __future__ import annotations

import math
from numbers import Real

import torch
from torch import nn


def _validate_sequences(sequences: list[torch.Tensor], pad_token_id: int) -> None:
    if not isinstance(sequences, list):
        raise TypeError("sequences must be a list of tensors")
    if not sequences:
        raise ValueError("sequences must contain at least one nonempty tensor")
    if isinstance(pad_token_id, bool) or not isinstance(pad_token_id, int):
        raise TypeError("pad_token_id must be an integer")
    for sequence in sequences:
        if not isinstance(sequence, torch.Tensor):
            raise TypeError("each sequence must be a tensor")
        if sequence.ndim != 1 or sequence.numel() == 0:
            raise ValueError("each sequence must be a nonempty 1D tensor")
        if sequence.dtype not in (torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64):
            raise TypeError("token sequences must have an integer dtype")
        if sequence.dtype != sequences[0].dtype or sequence.device != sequences[0].device:
            raise ValueError("all sequences must have the same dtype and device")
    limits = torch.iinfo(sequences[0].dtype)
    if not limits.min <= pad_token_id <= limits.max:
        raise ValueError("pad_token_id must be representable in the sequence dtype")


def prepare_batch_inputs(
    sequences: list[torch.Tensor], pad_token_id: int = 0
) -> tuple[torch.Tensor, torch.Tensor]:
    """Left-pad nonempty 1D integer sequences; return IDs and an int64 mask.

    Input dtype/device are preserved in IDs, and masks share their device.
    Masks are length-derived, so genuine tokens equal to ``pad_token_id`` remain
    visible. For positional models derive positions as
    ``(attention_mask.cumsum(-1) - 1).clamp_min(0)`` and pass them explicitly.
    """
    _validate_sequences(sequences, pad_token_id)
    width = max(sequence.numel() for sequence in sequences)
    ids = sequences[0].new_full((len(sequences), width), pad_token_id)
    mask = torch.zeros(ids.shape, dtype=torch.long, device=ids.device)
    for row, sequence in enumerate(sequences):
        ids[row, -sequence.numel():] = sequence
        mask[row, -sequence.numel():] = 1
    return ids, mask


def enforce_moe_deterministic_routing(model_or_config: object) -> tuple[str, ...]:
    """Configure exposed no-drop controls, returning the touched attribute paths.

    Supports a config object or an ``nn.Module`` with ``drop_tokens`` on itself,
    its submodules, or their ``config`` objects. Sets boolean ``drop_tokens`` to
    False. On that *same object* only, raises numeric ``capacity_factor``,
    ``train_capacity_factor``, and ``eval_capacity_factor`` to at least 100;
    zero is deliberately never used as a purported no-limit sentinel.

    This is a best-effort configuration adapter, not proof of deterministic
    routing or unlimited capacity. A large finite factor is not an infinite
    capacity, and setting config attributes may not update instantiated routers.
    Verify the actual model implementation honors these controls. Unrecognized
    routers and exposed capacity controls without a local boolean drop flag
    raise ValueError before any planned changes. Other architectures need an
    architecture-specific adapter. No floating-point equality is promised.
    """
    candidates: list[tuple[str, object]] = [("root", model_or_config)]
    if isinstance(model_or_config, nn.Module):
        candidates.extend((f"root.{name}", module) for name, module in model_or_config.named_modules() if name)
    candidates += [(f"{path}.config", obj.config) for path, obj in candidates if hasattr(obj, "config")]
    factors = ("capacity_factor", "train_capacity_factor", "eval_capacity_factor")
    # Absolute capacities lack a portable no-limit value, even with drop_tokens.
    absolute_capacities = ("expert_capacity", "token_capacity", "max_tokens_per_expert")
    updates: list[tuple[str, object, str, object]] = []
    seen: set[int] = set()
    for path, obj in candidates:
        if id(obj) in seen:
            continue
        seen.add(id(obj))
        local_factors = [name for name in factors if hasattr(obj, name)]
        if any(hasattr(obj, name) for name in absolute_capacities):
            raise ValueError(f"{path}: absolute capacity requires an architecture-specific adapter")
        has_drop = hasattr(obj, "drop_tokens")
        if not has_drop and local_factors:
            raise ValueError(f"{path}: capacity controls require a local boolean drop_tokens flag")
        if not has_drop:
            continue
        if not isinstance(obj.drop_tokens, bool):
            raise ValueError(f"{path}.drop_tokens must be boolean")
        updates.append((path, obj, "drop_tokens", False))
        for name in local_factors:
            value = getattr(obj, name)
            if isinstance(value, bool) or not isinstance(value, Real) or not math.isfinite(value):
                raise ValueError(f"{path}.{name} must be a finite numeric capacity factor")
            updates.append((path, obj, name, max(100, value)))
    if not updates:
        raise ValueError("unsupported router: no boolean drop_tokens control found")
    for _, obj, name, value in updates:
        setattr(obj, name, value)
    return tuple(f"{path}.{name}" for path, _, name, _ in updates)


def strict_deterministic_inference(
    model: nn.Module, sequences: list[torch.Tensor], pad_token_id: int = 0
) -> list[torch.Tensor]:
    """Return per-example logits using one unpadded forward per example.

    Removes the singleton batch axis: classifier logits ``(1, classes)`` become
    ``(classes,)`` and token logits ``(1, length, classes)`` become
    ``(length, classes)``. The model must accept
    ``input_ids``, ``attention_mask``, and ``position_ids`` keyword arguments and
    return a Tensor or an object with a Tensor ``logits`` attribute. All modules
    must already be in eval mode. No training state or global flags are changed.
    Models needing extra sequence-aligned inputs (such as segment masks) need
    a wrapper that supplies those inputs for each unpadded example; this helper
    intentionally does not infer them from token values.

    Serial execution trades batching throughput for identical per-example
    shapes regardless of caller batch composition/order. Exact repeatability
    still requires deterministic kernels/routing, fixed parameters, no mutable
    cross-call state or stochastic eval behavior, and the same device/runtime.
    This helper cannot make a nondeterministic model deterministic.
    """
    _validate_sequences(sequences, pad_token_id)
    if any(module.training for module in model.modules()):
        raise ValueError("strict inference requires every module to be in eval mode")
    results = []
    with torch.inference_mode():
        for sequence in sequences:
            ids, mask = prepare_batch_inputs([sequence], pad_token_id)
            positions = (mask.cumsum(-1) - 1).clamp_min(0)
            output = model(input_ids=ids, attention_mask=mask, position_ids=positions)
            logits = output if isinstance(output, torch.Tensor) else getattr(output, "logits", None)
            if not isinstance(logits, torch.Tensor) or logits.ndim < 2 or logits.shape[0] != 1:
                raise ValueError("model must return logits with a singleton leading batch dimension")
            results.append(logits[0].clone())
    return results