"""Minimum-window suffix budgets; pure contract helpers do not import Torch.

This changes the distribution of retained lengths, not attention visibility or
inference caps. A sampled K always includes the current position within K.
"""
from __future__ import annotations

import math


def validate_minimum(minimum):
    if type(minimum) is not int or minimum < 1 or minimum & (minimum - 1):
        raise ValueError('Minimum retained history must be a positive power of two')


def retained_candidates(available, minimum):
    """Unique ascending candidates, including full available history once."""
    validate_minimum(minimum)
    if type(available) is not int or available < 1:
        raise ValueError('Available history must be a positive integer')
    if available < minimum:
        return (available,)
    values = []
    value = minimum
    while value <= available:
        values.append(value)
        value *= 2
    if values[-1] != available:
        values.append(available)
    return tuple(values)


def retained_from_uniform(available, minimum, uniform):
    """Scalar reference used to verify equal-probability sampling intervals."""
    if not math.isfinite(uniform) or not 0 <= uniform < 1:
        raise ValueError('Uniform variate must lie in [0, 1)')
    choices = retained_candidates(available, minimum)
    return choices[int(uniform * len(choices))]


def minimum_budget_policy(canonical_policy, minimum):
    """Bind the new distribution while preserving legacy min1 serialization."""
    validate_minimum(minimum)
    if minimum == 1:
        return canonical_policy
    if canonical_policy is None:
        raise ValueError('Minimum-window training requires an explicit canonical budget policy')
    return {
        **canonical_policy,
        'retained_length_distribution': {
            'schema': 1,
            'name': 'uniform_unique_dyadic_minimum_plus_full',
            'minimum_history': minimum,
            'short_history': 'dense',
            'full_history_included_once': True,
            'scope': 'independent_sequence_query_shared_heads_and_middle_layers',
        },
    }


def sample_suffix_lengths(batch, length, *, minimum, generator, device):
    """Consume one original-shaped uniform draw; dense short queries still draw.

    CanonicalBudgetPlan calls this separately for each original 16x2048 draw,
    preserving per-example quantiles and generator advancement on world2/4.
    The original min1 trainer continues calling its unchanged sample_prefixes.
    """
    import torch

    validate_minimum(minimum)
    if batch < 1 or length < 1:
        raise ValueError('Sampler dimensions must be positive')
    n = torch.arange(1, length + 1, device=device, dtype=torch.int64)
    powers = torch.tensor(
        [minimum << p for p in range(max(1, length // minimum).bit_length())],
        device=device,
    )
    power_count = torch.bucketize(n, powers, right=True)
    extra_full = ((n & (n - 1)) != 0).long()
    count = torch.where(n < minimum, 1, power_count + extra_full)
    choice = (torch.rand((batch, length), device=device, generator=generator) * count).long()
    retained = torch.where(choice < power_count, minimum << choice, n)
    return torch.where(n < minimum, n, retained)