# Copyright 2026 Modilify
# SPDX-License-Identifier: LicenseRef-Modilify-Open-Model-1.0
"""Confidence-and-entropy commit policy for inference."""

from __future__ import annotations

from collections.abc import Sequence
from dataclasses import dataclass
import math

import torch

from .latent_deliberation import (
    advance_trajectory_clocks,
    should_force_trajectory_jump,
)

FUSED_EPS = 1e-6


def fused_commit_confidence(
    proposal_confidence: torch.Tensor,
    token_entropy: torch.Tensor,
    *,
    vocab_size: int = 256000,
    eps: float = FUSED_EPS,
) -> torch.Tensor:
    """Fuse proposal confidence with token entropy.

    Effective confidence uses an excess-entropy sigmoid:

        p = clamp(proposal_confidence, eps, 1 - eps)
        h2 = -p * log(p) - (1 - p) * log(1 - p)
        excess = max(token_entropy - h2, 0)
        fused = sigmoid(logit(p) - excess) ** 2

    When token entropy equals the binary entropy of ``p``, fused confidence
    equals ``p ** 2``. Entropy above that binary entropy reduces confidence.

    Args:
        proposal_confidence: Sampled-token probabilities, shape ``[batch, canvas]``.
        token_entropy: Token-level entropy, shape ``[batch, canvas]``.
        vocab_size: Unused; retained so callers can pass the model vocabulary.
        eps: Clamp that keeps logits finite.

    Returns:
        Fused commit confidence in ``(eps, 1 - eps)``.
    """

    del vocab_size
    p = proposal_confidence.float().clamp(min=eps, max=1.0 - eps)
    entropy = token_entropy.float().clamp(min=0.0)
    binary_entropy = -p * torch.log(p) - (1.0 - p) * torch.log1p(-p)
    excess = (entropy - binary_entropy).clamp(min=0.0)
    logit_p = torch.log(p) - torch.log1p(-p)
    fused = torch.sigmoid(logit_p - excess).square()
    return fused.clamp(min=eps, max=1.0 - eps)


def fused_commit_failure_rate(
    proposal_confidence: torch.Tensor,
    token_entropy: torch.Tensor,
    **kwargs: object,
) -> torch.Tensor:
    """Return ``1 - fused_commit_confidence``."""

    return 1.0 - fused_commit_confidence(
        proposal_confidence, token_entropy, **kwargs
    )


@dataclass(frozen=True)
class CommitPolicyDecision:
    """One inference transition from proposal to committed prefix."""

    normal_lengths: torch.LongTensor
    commit_lengths: torch.LongTensor
    commit_token_ids: torch.LongTensor
    jump_rows: torch.BoolTensor
    ponder_steps: torch.IntTensor
    stagnation_steps: torch.IntTensor


def prefix_failure_commit_lengths(
    failure_rate: torch.Tensor,
    *,
    failure_budget: float,
    valid_mask: torch.BoolTensor | None = None,
) -> torch.LongTensor:
    """Return the longest prefix with ``cumsum(failure_rate) < budget``.

    Args:
        failure_rate: Per-token failure rates, shape ``[batch, canvas]``.
        failure_budget: Strict cumulative risk limit.
        valid_mask: Optional canvas mask with the same shape.

    Returns:
        Commit lengths of shape ``[batch]``.
    """

    if failure_rate.ndim != 2:
        raise ValueError("Failure rate must have shape [batch, canvas].")
    if not math.isfinite(failure_budget) or failure_budget <= 0:
        raise ValueError("Commit failure budget must be finite and positive.")
    if valid_mask is None:
        valid_mask = torch.ones_like(failure_rate, dtype=torch.bool)
    if valid_mask.shape != failure_rate.shape:
        raise ValueError("Commit validity mask must match failure rate.")

    risk = failure_rate.float().clamp(0.0, 1.0) * valid_mask.to(torch.float32)
    cumulative_risk = risk.cumsum(dim=-1)
    contiguous_valid = valid_mask.long().cumprod(dim=-1).bool()
    allowed = cumulative_risk.lt(float(failure_budget)) & contiguous_valid
    return allowed.long().cumprod(dim=-1).sum(dim=-1)


def first_committed_token_lengths(
    proposal: torch.LongTensor,
    commit_lengths: torch.LongTensor,
    token_id: int | Sequence[int],
) -> torch.LongTensor:
    """Clip each prefix immediately after its first stop token.

    Args:
        proposal: Token IDs, shape ``[batch, canvas]``.
        commit_lengths: Unclipped prefix lengths, shape ``[batch]``.
        token_id: One stop ID or a sequence of stop IDs.

    Returns:
        Clipped commit lengths of shape ``[batch]``.
    """

    if proposal.ndim != 2 or commit_lengths.shape != proposal.shape[:1]:
        raise ValueError("Proposal and commit lengths must share a batch dimension.")
    positions = torch.arange(proposal.shape[1], device=proposal.device).unsqueeze(0)
    committed = positions.lt(commit_lengths[:, None])
    stop_token_ids = (
        (int(token_id),)
        if isinstance(token_id, int)
        else tuple(dict.fromkeys(int(value) for value in token_id))
    )
    if not stop_token_ids:
        raise ValueError("At least one stop token ID is required.")
    matches = proposal.eq(stop_token_ids[0])
    for value in stop_token_ids[1:]:
        matches |= proposal.eq(value)
    matches &= committed
    sentinel = torch.full_like(positions, proposal.shape[1])
    first = torch.where(matches, positions, sentinel).min(dim=-1).values
    clipped = torch.where(first.lt(proposal.shape[1]), first + 1, commit_lengths)
    return torch.minimum(clipped, commit_lengths)


def bounded_prefix_failure_commit_lengths(
    committed_token_ids: torch.LongTensor,
    failure_rate: torch.Tensor,
    *,
    failure_budget: float,
    remaining_lengths: torch.LongTensor,
    stop_token_id: int | Sequence[int],
    valid_mask: torch.BoolTensor | None = None,
) -> torch.LongTensor:
    """Apply remaining-length and stop-token bounds to the prefix policy."""

    if committed_token_ids.shape != failure_rate.shape:
        raise ValueError("Committed token IDs and failure rate must share [batch, canvas].")
    if remaining_lengths.shape != committed_token_ids.shape[:1]:
        raise ValueError("Remaining lengths must have shape [batch].")
    commit_lengths = prefix_failure_commit_lengths(
        failure_rate,
        failure_budget=failure_budget,
        valid_mask=valid_mask,
    )
    commit_lengths = torch.minimum(commit_lengths, remaining_lengths.clamp_min(0))
    return first_committed_token_lengths(
        committed_token_ids,
        commit_lengths,
        stop_token_id,
    )


def select_commit_lengths(
    sampled_token_ids: torch.LongTensor,
    normal_failure_rate: torch.Tensor,
    previous_failure_rate: torch.Tensor,
    greedy_token_ids: torch.LongTensor,
    jump_failure_rate: torch.Tensor,
    *,
    ponder_steps: torch.Tensor,
    stagnation_steps: torch.Tensor,
    active_rows: torch.BoolTensor,
    remaining_lengths: torch.LongTensor,
    failure_budget: float,
    jump_failure_budget: float,
    stop_token_id: int | Sequence[int],
    max_ponder_steps: int,
    stagnation_threshold: int,
    min_progress: float,
    valid_mask: torch.BoolTensor | None = None,
) -> CommitPolicyDecision:
    """Select sampled commits or a greedy jump after stagnation.

    Progress is the signed change in fused failure rate over the union of the
    previous and current prefixes plus one blocking position.

    Args:
        sampled_token_ids: Temperature-sampled canvas tokens.
        normal_failure_rate: Fused failure rates for the sampled tokens.
        previous_failure_rate: Fused failure rates from the previous step.
        greedy_token_ids: Greedy canvas tokens used for jumps.
        jump_failure_rate: Fused failure rates for the greedy tokens.
        ponder_steps: Per-row useful-ponder clocks.
        stagnation_steps: Per-row stagnation clocks.
        active_rows: Rows that are still generating.
        remaining_lengths: Tokens still allowed on each row.
        failure_budget: Normal commit budget.
        jump_failure_budget: Forced-jump budget.
        stop_token_id: Turn or EOS stop IDs.
        max_ponder_steps: Watchdog on useful pondering.
        stagnation_threshold: Watchdog on true stagnation.
        min_progress: Minimum signed improvement counted as progress.
        valid_mask: Optional canvas mask.

    Returns:
        Commit lengths, token IDs, jump flags, and updated clocks.
    """

    if not (
        sampled_token_ids.shape
        == normal_failure_rate.shape
        == previous_failure_rate.shape
        == greedy_token_ids.shape
        == jump_failure_rate.shape
    ):
        raise ValueError("Sampled and greedy statistics must share [batch, canvas].")

    normal = bounded_prefix_failure_commit_lengths(
        sampled_token_ids,
        normal_failure_rate,
        failure_budget=failure_budget,
        remaining_lengths=remaining_lengths,
        stop_token_id=stop_token_id,
        valid_mask=valid_mask,
    )
    canvas_length = normal_failure_rate.shape[1]
    previous_prefix_length = prefix_failure_commit_lengths(
        previous_failure_rate,
        failure_budget=failure_budget,
        valid_mask=valid_mask,
    )
    frontier_length = torch.maximum(previous_prefix_length, normal) + 1
    valid_lengths = (
        valid_mask.long().sum(dim=-1)
        if valid_mask is not None
        else torch.full_like(frontier_length, canvas_length)
    )
    frontier_length = torch.minimum(frontier_length, valid_lengths)
    positions = torch.arange(canvas_length, device=normal_failure_rate.device)[None, :]
    progress_mask = positions < frontier_length[:, None]
    if valid_mask is not None:
        progress_mask &= valid_mask
    progress_mask &= active_rows[:, None]
    signed_improvement = previous_failure_rate.float() - normal_failure_rate.float()
    weights = progress_mask.float()
    progress = (signed_improvement * weights).sum(dim=-1) / weights.sum(dim=-1).clamp_min(
        1.0
    )
    next_ponder, next_stagnation = advance_trajectory_clocks(
        ponder_steps,
        stagnation_steps,
        commit_lengths=normal,
        active_rows=active_rows,
        progress_scores=progress,
        min_progress=min_progress,
    )
    jump_rows = normal.eq(0) & active_rows & should_force_trajectory_jump(
        next_ponder,
        next_stagnation,
        max_ponder_steps=max_ponder_steps,
        stagnation_threshold=stagnation_threshold,
    )
    jump_commit = bounded_prefix_failure_commit_lengths(
        greedy_token_ids,
        jump_failure_rate,
        failure_budget=jump_failure_budget,
        remaining_lengths=remaining_lengths,
        stop_token_id=stop_token_id,
        valid_mask=valid_mask,
    )
    committed = torch.where(jump_rows, jump_commit, normal)
    commit_token_ids = torch.where(
        jump_rows[:, None],
        greedy_token_ids,
        sampled_token_ids,
    )
    committed = first_committed_token_lengths(
        commit_token_ids,
        committed,
        stop_token_id,
    )
    committed = torch.where(active_rows, committed, 0)
    jump_rows &= committed.gt(0)
    next_ponder = torch.where(committed.gt(0), 0, next_ponder).to(torch.int32)
    next_stagnation = torch.where(committed.gt(0), 0, next_stagnation).to(torch.int32)
    return CommitPolicyDecision(
        normal_lengths=normal,
        commit_lengths=committed,
        commit_token_ids=commit_token_ids,
        jump_rows=jump_rows,
        ponder_steps=next_ponder,
        stagnation_steps=next_stagnation,
    )


__all__ = [
    "CommitPolicyDecision",
    "bounded_prefix_failure_commit_lengths",
    "first_committed_token_lengths",
    "fused_commit_confidence",
    "fused_commit_failure_rate",
    "prefix_failure_commit_lengths",
    "select_commit_lengths",
]