# 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
FUSED_ENTROPY_WEIGHT = 0.5


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

    Effective confidence is defined using a multiplicative entropy penalty:

        fused = confidence * exp(-entropy_weight * token_entropy)

    This scales proposal confidence by the entropy discount factor e^(-alpha * H),
    penalizing token predictions with high distribution disorder.
    """

    if vocab_size <= 1:
        raise ValueError("`vocab_size` must be greater than one.")
    if not math.isfinite(entropy_weight) or entropy_weight < 0.0:
        raise ValueError("`entropy_weight` must be finite and non-negative.")
    confidence = proposal_confidence.float().clamp(min=eps, max=1.0)
    entropy = token_entropy.float().clamp(0.0, math.log(vocab_size))
    entropy_penalty = torch.exp(-entropy_weight * entropy)
    fused = confidence * entropy_penalty
    return fused.clamp(min=eps, max=1.0)


def fused_commit_failure_rate(
    proposal_confidence: torch.Tensor,
    token_entropy: torch.Tensor,
    **kwargs: object,
) -> torch.Tensor:
    """Return ``1 - effective_confidence`` from the shared fusion helper."""

    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 valid prefix satisfying ``cumsum(failure_rate) < budget``."""

    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 committed prefix immediately after its first matching stop token."""

    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 length and stop-token bounds to the shared failure-rate 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:
    """Use normal sampled commits and a fixed-budget greedy JUMP.

    Progress is measured from the signed change in fused failure rate over the
    frontier region (the union of the previous and current commit prefixes plus
    one blocking position), not from raw confidence/entropy deltas.
    """

    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",
    "FUSED_ENTROPY_WEIGHT",
    "bounded_prefix_failure_commit_lengths",
    "first_committed_token_lengths",
    "fused_commit_confidence",
    "fused_commit_failure_rate",
    "prefix_failure_commit_lengths",
    "select_commit_lengths",
]