mp_regular_sampler.py

"""Self-contained start-point sampler builder for regular ISD algorithms."""

from __future__ import annotations

import warnings

import mpmath as mp

from mp_regular_isd import get_regular_algorithm
from mp_utils import (
    FLOAT_PRECISION,
    Float,
    as_float_array,
    binomH,
    is_finite_number,
    next_down,
    next_up,
    random_unit,
)
from mp_validator import build_state_resolver, safe_constraint_value

LAMBDA_RATIOS = (Float("0"), Float("1"), Float("0.5"), Float("0.25"), Float("0.75"), Float("0.125"), Float("0.875"))
PI_RATIOS = LAMBDA_RATIOS
OMEGA_V_RATIOS = (Float("0.125"), Float("0.875"), Float("0.5"), Float("0.25"), Float("0.75"))
REP_OMEGA_V_RATIOS = (Float("1"), Float("0.125"), Float("0.875"), Float("0.5"), Float("0.25"), Float("0.75"))
REP_RATIO_GRID = LAMBDA_RATIOS


def _vector_objective(spec, state_resolver):
    """Bind one registered regular formula to a vector-only objective callback.

    Args:
        spec: Algorithm descriptor from the regular registry.
        kappa: Fixed original code rate ``k / n`` of the current point.
        omega: Fixed RSD target weight ``w / n`` of the current point.
    """
    if spec["direct"]:
        return None

    def objective(x):
        _, state = state_resolver(x)
        if hasattr(state, "omega_v") and Float(state.omega_v) <= Float("0"):
            return Float(spec["penalty"])
        args = [getattr(state, name) for name in spec["field_names"]]
        with warnings.catch_warnings():
            warnings.simplefilter("ignore", RuntimeWarning)
            value = Float(spec["objective"](*args))
        return value if is_finite_number(value) else Float(spec["penalty"])

    return objective


def _wrap_constraints(spec, state_resolver):
    """Wrap namedtuple-style regular constraints into SciPy vector callbacks.

    Args:
        kappa: Fixed original code rate ``k / n`` of the current point.
        omega: Fixed RSD target weight ``w / n`` of the current point.
        spec: Algorithm descriptor containing the original constraints.
    """
    wrapped = []
    for constraint in spec["constraints"]:
        ctype = constraint["type"]
        fun = constraint["fun"]
        wrapped.append(
            {
                "type": ctype,
                "fun": lambda x, _fun=fun: safe_constraint_value(_fun, state_resolver(x)[1]),
            }
        )
    return wrapped


def effective_kappa(kappa, omega):
    """Return the post-regularity effective code rate ``kappa - omega``.

    Args:
        kappa: Original code rate ``k / n`` before regularity encoding.
        omega: RSD target weight ``w / n`` subtracted by regularity encoding.
    """
    return Float(kappa) - Float(omega)


def interpolate_interval(lower, upper, ratio):
    """Map a unit ratio to a point inside a closed interval.

    Args:
        lower: Left endpoint of the interval.
        upper: Right endpoint of the interval.
        ratio: Relative position inside the interval.
    """
    lower = Float(lower)
    upper = Float(upper)
    ratio = Float(ratio)
    if upper <= lower:
        return lower
    return lower + (upper - lower) * ratio


def strict_interval_for_sampling(lower, upper):
    """Move an interval slightly inward when a strict interior exists.

    Args:
        lower: Left endpoint of the interval.
        upper: Right endpoint of the interval.
    """
    lower = Float(lower)
    upper = Float(upper)
    if upper <= lower:
        return lower, upper
    span = upper - lower
    margin = max(FLOAT_PRECISION * Float("1024"), span / Float("64"))
    if margin * Float("2") < span:
        return lower + margin, upper - margin
    strict_lo = next_up(lower)
    strict_hi = next_down(upper)
    if strict_lo < strict_hi:
        return strict_lo, strict_hi
    return lower, upper


def strict_upper_cap(lower, upper, cap):
    """Shrink an upper endpoint slightly when it sits on a hard boundary.

    Args:
        lower: Left endpoint of the interval.
        upper: Right endpoint of the interval.
        cap: Hard upper cap that should be avoided when possible.
    """
    lower = Float(lower)
    upper = Float(upper)
    cap = Float(cap)
    if upper < lower:
        return lower
    strict_cap = next_down(cap)
    if strict_cap <= lower:
        return upper
    if abs(upper - cap) <= FLOAT_PRECISION * Float("64"):
        return min(upper, strict_cap)
    return upper


def regular_enum_lambda_interval(kappa, omega):
    """Return the feasible global ``ell / n`` interval for regular ENUM."""
    del omega
    return Float("0"), max(Float("0"), min(Float("1"), Float("1") - Float(kappa)))


def regular_rep_lambda_interval(kappa, omega):
    """Return the feasible global ``ell / n`` interval for regular REP variants."""
    del omega
    return Float("0"), max(Float("0"), min(Float("1"), Float("1") - Float(kappa)))


def regular_enum_pi_interval(kappa, omega, lambda_):
    """Return the feasible ``p / n`` slice for fixed ``ell / n`` in regular ENUM."""
    kappa = Float(kappa)
    omega = Float(omega)
    lambda_ = Float(lambda_)
    kappa_p = effective_kappa(kappa, omega)
    lower = max(Float("0"), kappa + lambda_ - Float("1"))
    upper = min(Float("1"), omega, kappa_p + lambda_)
    return lower, max(lower, upper)


def regular_rep_pi_interval(kappa, omega, lambda_):
    """Return the feasible ``p / n`` slice for fixed ``ell / n`` in regular REP."""
    kappa = Float(kappa)
    omega = Float(omega)
    lambda_ = Float(lambda_)
    kappa_p = effective_kappa(kappa, omega)
    lower = max(Float("0"), kappa + lambda_ - Float("1"))
    upper = min(Float("1"), omega, kappa_p + lambda_, lambda_)
    return lower, max(lower, upper)


def regular_omega_v_interval(kappa, omega, lambda_, pi, *, strict_lower=True):
    """Return the feasible ``N / n`` slice for fixed ``ell / n`` and ``p / n``.

    Args:
        kappa: Original code rate ``k / n`` before regularity encoding.
        omega: RSD target weight ``w / n``.
        lambda_: Current extra-equation parameter ``ell / n``.
        pi: Current selected error-block parameter ``p / n``.
        strict_lower: Whether to push the lower bound slightly inward.
    """
    omega = Float(omega)
    lambda_ = Float(lambda_)
    pi = Float(pi)
    numerator = effective_kappa(kappa, omega) + lambda_
    lower = max(Float("0"), pi)
    if omega < Float("1"):
        lower_from_width = numerator * omega / (Float("1") - omega)
        if strict_lower:
            lower_from_width = next_up(lower_from_width)
        lower = max(lower, lower_from_width)
    upper = min(Float("1"), omega, numerator)
    return lower, max(lower, upper)


def regular_enum_opt_bounds(kappa, omega):
    """Return conservative box bounds for regular ENUM free variables."""
    lambda_lo, lambda_hi = regular_enum_lambda_interval(kappa, omega)
    omega_cap = min(Float("1"), max(Float("0"), Float(omega)))
    omega_floor = next_up(Float("0"))
    return [
        (Float("0"), omega_cap),
        (max(Float("0"), lambda_lo), min(Float("1"), lambda_hi)),
        (omega_floor, omega_cap),
    ]


def regular_rep_opt_bounds(kappa, omega):
    """Return conservative box bounds for regular REP free variables."""
    lambda_lo, lambda_hi = regular_rep_lambda_interval(kappa, omega)
    omega_cap = min(Float("1"), max(Float("0"), Float(omega)))
    omega_floor = next_up(Float("0"))
    return [
        (Float("0"), omega_cap),
        (max(Float("0"), lambda_lo), min(Float("1"), lambda_hi)),
        (Float("0"), omega_cap),
        (Float("0"), omega_cap),
        (omega_floor, omega_cap),
    ]


def regular_v_width(kappa, omega, lambda_, omega_v):
    """Return the per-block width ``v = (kappa - omega + ell) / N``."""
    omega_v = Float(omega_v)
    if omega_v <= 0:
        return Float("0")
    return (effective_kappa(kappa, omega) + Float(lambda_)) / omega_v


def regular_representation_cost(base_weight, window, epsilon, v):
    """Return one regular representation count exponent.

    Args:
        base_weight: Target block mass before adding cancelling blocks.
        window: Remaining selected-block mass that can host cancelling blocks.
        epsilon: Normalized cancelling-block mass introduced by the split.
        v: Per-block width inside the extended information set.
    """
    base_weight = Float(base_weight)
    window = max(Float("0"), Float(window))
    epsilon = Float(epsilon)
    v = Float(v)
    if window <= 0 or epsilon <= 0:
        return base_weight
    return base_weight + binomH(window, epsilon) + mp.log(v, 2) * epsilon


def regular_representation_peak(window, v):
    """Return the maximizing cancelling-block mass on the regular branch."""
    window = max(Float("0"), Float(window))
    v = Float(v)
    if window <= 0 or v <= 0:
        return Float("0")
    return window * v / (Float("1") + v)


def solve_regular_epsilon_for_target(base_weight, window, target_cost, v):
    """Invert the regular representation count on its increasing branch."""
    base_weight = Float(base_weight)
    window = max(Float("0"), Float(window))
    target_cost = Float(target_cost)
    v = Float(v)
    if window <= 0 or v <= 0 or target_cost <= base_weight + FLOAT_PRECISION:
        return Float("0")
    upper = min(window, regular_representation_peak(window, v))
    upper_cost = regular_representation_cost(base_weight, window, upper, v)
    if target_cost >= upper_cost:
        return upper
    lower = Float("0")
    for _ in range(100):
        middle = (lower + upper) / Float("2")
        if regular_representation_cost(base_weight, window, middle, v) < target_cost:
            lower = middle
        else:
            upper = middle
    return (lower + upper) / Float("2")


def regular_enum_opt_sample_func(kappa, omega):
    """Build the conditional sampler for regular ENUM.

    Args:
        kappa: Fixed original code rate ``k / n``.
        omega: Fixed RSD target weight ``w / n``.
    """
    kappa = Float(kappa)
    omega = Float(omega)
    lambda_lo, lambda_hi = regular_enum_lambda_interval(kappa, omega)
    fixed_grid_size = len(LAMBDA_RATIOS) * len(PI_RATIOS) * len(OMEGA_V_RATIOS)
    floor = next_up(Float("0"))

    def sample(rng, iteration):
        lambda_sample_lo, lambda_sample_hi = strict_interval_for_sampling(lambda_lo, lambda_hi)
        if lambda_sample_hi <= lambda_sample_lo:
            lambda_ = lambda_sample_lo
            pi_seed_index = iteration
            omega_v_seed_index = iteration
        elif iteration < fixed_grid_size:
            lambda_idx = iteration % len(LAMBDA_RATIOS)
            rem = iteration // len(LAMBDA_RATIOS)
            pi_seed_index = rem % len(PI_RATIOS)
            omega_v_seed_index = (rem // len(PI_RATIOS)) % len(OMEGA_V_RATIOS)
            lambda_ = interpolate_interval(lambda_sample_lo, lambda_sample_hi, LAMBDA_RATIOS[lambda_idx])
        else:
            lambda_ = interpolate_interval(lambda_sample_lo, lambda_sample_hi, random_unit(rng))
            pi_seed_index = iteration
            omega_v_seed_index = iteration

        pi_lo, pi_hi = regular_enum_pi_interval(kappa, omega, lambda_)
        pi_lo, pi_hi = strict_interval_for_sampling(pi_lo, pi_hi)
        pi = (
            pi_lo
            if pi_hi <= pi_lo
            else interpolate_interval(
                pi_lo,
                pi_hi,
                PI_RATIOS[pi_seed_index % len(PI_RATIOS)] if iteration < fixed_grid_size else random_unit(rng),
            )
        )

        omega_v_lo, omega_v_hi = regular_omega_v_interval(kappa, omega, lambda_, pi, strict_lower=True)
        omega_v_lo = max(omega_v_lo, floor)
        omega_v_hi = max(omega_v_lo, strict_upper_cap(omega_v_lo, omega_v_hi, omega))
        omega_v_lo, omega_v_hi = strict_interval_for_sampling(omega_v_lo, omega_v_hi)
        omega_v = (
            omega_v_lo
            if omega_v_hi <= omega_v_lo
            else interpolate_interval(
                omega_v_lo,
                omega_v_hi,
                OMEGA_V_RATIOS[omega_v_seed_index % len(OMEGA_V_RATIOS)] if iteration < fixed_grid_size else random_unit(rng),
            )
        )
        return [pi, lambda_, omega_v]

    return sample


def regular_rep_opt_sample_func(kappa, omega, objective):
    """Build the conditional sampler for regular REP.

    Args:
        kappa: Fixed original code rate ``k / n``.
        omega: Fixed RSD target weight ``w / n``.
        objective: Bound objective callback used to reject obvious penalty points.
    """
    kappa = Float(kappa)
    omega = Float(omega)
    lambda_lo, lambda_hi = regular_rep_lambda_interval(kappa, omega)
    fixed_grid_size = len(LAMBDA_RATIOS) * len(PI_RATIOS) * len(REP_OMEGA_V_RATIOS)
    omega_floor = next_up(Float("0"))

    def sample(rng, iteration):
        lambda_sample_lo, lambda_sample_hi = strict_interval_for_sampling(lambda_lo, lambda_hi)
        if lambda_sample_hi <= lambda_sample_lo:
            lambda_ = lambda_sample_lo
            pi_seed_index = iteration
            omega_v_seed_index = iteration
        elif iteration < fixed_grid_size:
            lambda_idx = iteration % len(LAMBDA_RATIOS)
            rem = iteration // len(LAMBDA_RATIOS)
            pi_seed_index = rem % len(PI_RATIOS)
            omega_v_seed_index = (rem // len(PI_RATIOS)) % len(REP_OMEGA_V_RATIOS)
            lambda_ = interpolate_interval(lambda_sample_lo, lambda_sample_hi, LAMBDA_RATIOS[lambda_idx])
        else:
            lambda_ = interpolate_interval(lambda_sample_lo, lambda_sample_hi, random_unit(rng))
            pi_seed_index = iteration
            omega_v_seed_index = iteration

        pi_lo, pi_hi = regular_rep_pi_interval(kappa, omega, lambda_)
        pi_lo, pi_hi = strict_interval_for_sampling(pi_lo, pi_hi)
        pi = (
            pi_lo
            if pi_hi <= pi_lo
            else interpolate_interval(
                pi_lo,
                pi_hi,
                PI_RATIOS[pi_seed_index % len(PI_RATIOS)] if iteration < fixed_grid_size else random_unit(rng),
            )
        )

        omega_v_lo, omega_v_hi = regular_omega_v_interval(kappa, omega, lambda_, pi, strict_lower=True)
        omega_v_lo = max(omega_v_floor := omega_floor, omega_v_lo)
        omega_v_hi = max(omega_v_lo, strict_upper_cap(omega_v_lo, omega_v_hi, omega))
        omega_v_lo, omega_v_hi = strict_interval_for_sampling(omega_v_lo, omega_v_hi)
        omega_v = (
            omega_v_lo
            if omega_v_hi <= omega_v_lo
            else interpolate_interval(
                omega_v_lo,
                omega_v_hi,
                REP_OMEGA_V_RATIOS[omega_v_seed_index % len(REP_OMEGA_V_RATIOS)]
                if iteration < fixed_grid_size
                else random_unit(rng),
            )
        )

        v = regular_v_width(kappa, omega, lambda_, omega_v)
        window_x = max(Float("0"), omega_v - pi)
        epsilon_x_cap = min(window_x, regular_representation_peak(window_x, v))
        rx_floor = pi
        rx_cap = min(lambda_, regular_representation_cost(pi, window_x, epsilon_x_cap, v))
        if rx_cap <= rx_floor + FLOAT_PRECISION:
            epsilon_x = Float("0")
            pi_x = pi / Float("2")
            rx = rx_floor
        else:
            rx_target = rx_floor + (rx_cap - rx_floor) * (
                REP_RATIO_GRID[(iteration // len(PI_RATIOS)) % len(REP_RATIO_GRID)] if iteration < fixed_grid_size else random_unit(rng)
            )
            epsilon_x = solve_regular_epsilon_for_target(pi, window_x, rx_target, v)
            pi_x = pi / Float("2") + epsilon_x
            rx = regular_representation_cost(pi, window_x, epsilon_x, v)

        window_y = max(Float("0"), omega_v - pi_x)
        epsilon_y_cap = min(window_y, regular_representation_peak(window_y, v))
        ry_floor = pi_x
        ry_cap = min(rx, regular_representation_cost(pi_x, window_y, epsilon_y_cap, v))
        if ry_cap <= ry_floor + FLOAT_PRECISION:
            epsilon_y = Float("0")
        else:
            ry_target = ry_floor + (ry_cap - ry_floor) * (
                REP_RATIO_GRID[(iteration // (len(PI_RATIOS) * 2)) % len(REP_RATIO_GRID)]
                if iteration < fixed_grid_size
                else random_unit(rng)
            )
            epsilon_y = solve_regular_epsilon_for_target(pi_x, window_y, ry_target, v)

        candidate = [pi, lambda_, epsilon_x, epsilon_y, omega_v]
        if objective is None:
            return candidate
        candidate_obj = Float(objective(candidate))
        if is_finite_number(candidate_obj) and candidate_obj < Float("1e3"):
            return candidate
        omega_v_safe = interpolate_interval(omega_v_lo, omega_v_hi, Float("0.6"))
        pi_safe = min(pi, max(pi_lo, min(pi_hi, lambda_ - FLOAT_PRECISION * Float("64"), omega_v_safe - FLOAT_PRECISION * Float("64"))))
        return [Float(pi_safe), Float(lambda_), Float("0"), Float("0"), Float(omega_v_safe)]

    return sample


def build_regular_sampler_config(algorithm, kappa, omega):
    """Build all optimization ingredients for one regular algorithm.

    Args:
        algorithm: Algorithm name from the regular registry.
        kappa: Fixed original code rate ``k / n`` of the current point.
        omega: Fixed RSD target weight ``w / n`` of the current point.
    """
    spec = get_regular_algorithm(algorithm)
    config = dict(spec)
    config["family"] = "regular"
    config["kappa"] = Float(kappa)
    config["omega"] = Float(omega)
    state_resolver = build_state_resolver(kappa, omega, var_names=spec["free_var_names"], vars_type=spec["vars_type"])
    config["objective_vector"] = _vector_objective(spec, state_resolver)

    if spec["direct"]:
        config["bounds"] = []
        config["sample_func"] = None
        config["wrapped_constraints"] = []
        return config

    if spec["name"] == "enum":
        config["bounds"] = regular_enum_opt_bounds(kappa, omega)
        config["sample_func"] = regular_enum_opt_sample_func(kappa, omega)
    elif spec["name"] == "rep":
        config["bounds"] = regular_rep_opt_bounds(kappa, omega)
        config["sample_func"] = regular_rep_opt_sample_func(kappa, omega, config["objective_vector"])
    elif spec["name"] == "rep_mo":
        config["bounds"] = regular_rep_opt_bounds(kappa, omega)
        config["sample_func"] = regular_rep_opt_sample_func(kappa, omega, config["objective_vector"])
    else:
        raise ValueError(f"Unsupported regular sampler type `{algorithm}`.")

    config["wrapped_constraints"] = _wrap_constraints(spec, state_resolver)
    return config