"""T3 per-window covert-channel budget calculator (#000036 §11).

Implements the closed-form bound from
`docs/soft-hash-channel-t3-bound.md` §6:

    I_window ≤ C_B1 · g · W · log₂(SNR_grad + 1)
             + C_B2 · ⌈W / K⌉ · log₂(R)
             + C_B3 · ⌈W / E⌉ · log₂(N_b · σ_grad / ‖∇L_max‖) / 2

Inputs are operator-measurable on a deployment:

    g            T2 gradient-fraction the adversary controls
    ‖∇L_max‖     gradient-norm cap (clipping ceiling)
    σ_grad       per-step gradient stddev (SGD noise floor)
    K            LR-decision interval (steps between LR changes)
    R            LR-grid size
    W            nonce-window length (steps before re-anchor)
    N_b          batches per epoch
    E            steps per epoch (typically equals N_b)

Conservative defaults: C_B1 = C_B2 = C_B3 = 1 (each by data-
processing inequality; see §6 + §10). Operators with deployment-
specific empirical measurements of the constants pass them via
``--c-b1`` / ``--c-b2`` / ``--c-b3`` flags.

**B1 model.** Default ``b1_model="max_envelope"`` (#000036 §3.1,
dav1d review 2026-05-11): ``g`` has two competing interpretations —
*fraction of steerable directions* (each carrying full per-channel
SNR ``G/σ``) and *aggregate amplitude shrinkage* (one effective
channel carrying SNR ``g·G/σ``). The conservative envelope is the
larger of the two:

    B1_fraction_channels = C_B1 · g · W · log₂(1 + G/σ)
    B1_aggregate_bias    = C_B1 ·     W · log₂(1 + g·G/σ)
    B1 (max_envelope)    = max(B1_fraction_channels, B1_aggregate_bias)

Other ``b1_model`` choices: ``"fraction_channels"`` / ``"aggregate_bias"``
(the individual envelope terms) and ``"effective_control_v1"`` (the
older ``C_B1 · g · W · log₂(1 + g·G/σ)`` — an operational risk
score, NOT a worst-case bound; kept for backward comparison).
``certification_status`` is relative to whichever model is selected;
under the default ``max_envelope`` it is a genuine upper-bound
certification.

Pure stdlib. No numpy / scipy dependency — the math is
arithmetic + math.log2.

Usage:

    python -m bench.scripts.t3_bound_calculator \\
        --gradient-fraction 0.05 \\
        --gradient-norm-max 1.0 \\
        --gradient-noise-stddev 0.1 \\
        --lr-decision-interval 100 \\
        --lr-grid-size 8 \\
        --window-length 10000 \\
        --batches-per-epoch 1024 \\
        --steps-per-epoch 1024
"""

from __future__ import annotations

import argparse
import json
import math
import sys


CALCULATOR_VERSION = "t3-bound-v1-bottou-refinement"
# B1 model selection (#000036 §3.1). Default is the conservative
# envelope; effective_control_v1 is the older non-worst-case formula
# kept for backward comparison.
B1_MODELS = (
    "max_envelope",
    "fraction_channels",
    "aggregate_bias",
    "effective_control_v1",
)
B1_MODEL_DEFAULT = "max_envelope"
CERTIFICATION_THRESHOLD_BITS = 256


def _require_finite_float(name: str, value: object) -> float:
    """Reject bools, non-numbers, NaN, and infinities; return float.

    Python's ``isinstance(True, int)`` is True, so a bare ``isinstance``
    check lets ``True`` / ``False`` leak into a numeric calculation —
    unacceptable for a security calculator. Reject bools explicitly.
    """
    if isinstance(value, bool) or not isinstance(value, (int, float)):
        raise ValueError(f"{name} must be a finite number; got {value!r}")
    fv = float(value)
    if not math.isfinite(fv):
        raise ValueError(f"{name} must be finite (no NaN / inf); got {value!r}")
    return fv


def _require_positive_int(name: str, value: object) -> int:
    """Reject bools and non-ints; require value > 0.

    ``type(value) is int`` (not ``isinstance``) rejects ``True`` / ``False``.
    """
    if type(value) is not int or value <= 0:
        raise ValueError(f"{name} must be a positive int; got {value!r}")
    return value


def t3_bound_bits(
    *,
    gradient_fraction: float,        # g
    gradient_norm_max: float,        # ‖∇L_max‖
    gradient_noise_stddev: float,    # σ_grad
    lr_decision_interval: int,       # K
    lr_grid_size: int,               # R
    window_length: int,              # W
    batches_per_epoch: int,          # N_b
    steps_per_epoch: int,            # E
    c_b1: float = 1.0,
    c_b2: float = 1.0,
    c_b3: float = 1.0,
    b1_model: str = B1_MODEL_DEFAULT,
) -> dict:
    """Compute the per-window mutual-information bound + per-bandwidth
    contributions per #000036 §6.

    ``b1_model`` selects the gradient-bias model (see #000036 §3.1):
    ``"max_envelope"`` (default, conservative), ``"fraction_channels"``,
    ``"aggregate_bias"``, or ``"effective_control_v1"`` (older
    non-worst-case formula).

    Returns a dict with the total bound, the per-bandwidth
    contributions (including all three B1 model variants), the
    constants + selected B1 model, a structured ``certification_status``,
    and a sentence of operator guidance based on the bound vs the
    SHA-256 (256-bit) certification threshold.
    """
    # Validation — reject bools / NaN / inf; coerce numeric types.
    gradient_fraction = _require_finite_float("gradient_fraction", gradient_fraction)
    if not 0 <= gradient_fraction <= 1:
        # 0 is allowed: a deployment with no T2 gradient surface still
        # has T3's LR + batch-order channels (B2, B3), so B1=0 is a
        # meaningful component-isolation case.
        raise ValueError(
            f"gradient_fraction must be in [0, 1]; got {gradient_fraction}"
        )
    gradient_norm_max = _require_finite_float("gradient_norm_max", gradient_norm_max)
    if gradient_norm_max <= 0:
        raise ValueError(f"gradient_norm_max must be > 0; got {gradient_norm_max}")
    gradient_noise_stddev = _require_finite_float(
        "gradient_noise_stddev", gradient_noise_stddev
    )
    if gradient_noise_stddev <= 0:
        raise ValueError(
            f"gradient_noise_stddev must be > 0; got {gradient_noise_stddev}"
        )
    lr_decision_interval = _require_positive_int(
        "lr_decision_interval", lr_decision_interval
    )
    lr_grid_size = _require_positive_int("lr_grid_size", lr_grid_size)
    window_length = _require_positive_int("window_length", window_length)
    batches_per_epoch = _require_positive_int("batches_per_epoch", batches_per_epoch)
    steps_per_epoch = _require_positive_int("steps_per_epoch", steps_per_epoch)
    c_b1 = _require_finite_float("c_b1", c_b1)
    c_b2 = _require_finite_float("c_b2", c_b2)
    c_b3 = _require_finite_float("c_b3", c_b3)
    for name, c in [("c_b1", c_b1), ("c_b2", c_b2), ("c_b3", c_b3)]:
        if not 0 <= c <= 1:
            raise ValueError(f"{name} must be in [0, 1]; got {c}")
    if b1_model not in B1_MODELS:
        raise ValueError(
            f"b1_model must be one of {B1_MODELS}; got {b1_model!r}"
        )

    # B1 — gradient bias. Per-step discrete channel-capacity bound on
    # the adversary-controlled parameter shift (NOT Fano's inequality —
    # see #000036 §3): the per-step shift falls in one of ~SNR+1
    # distinguishable buckets, capacity ≤ log₂(SNR+1). Two SNR readings
    # (#000036 §3.1):
    #   per-channel  SNR = G/σ      → fraction_channels model
    #   aggregate    SNR = g·G/σ    → aggregate_bias / effective_control
    # Default b1_model="max_envelope" = max of the two envelope terms —
    # genuinely upper-bounding across both interpretations (dav1d review
    # 2026-05-11).
    snr_grad = gradient_fraction * gradient_norm_max / gradient_noise_stddev   # g·G/σ
    snr_per_channel = gradient_norm_max / gradient_noise_stddev                # G/σ
    b1_fraction_channels = (
        c_b1 * gradient_fraction * window_length * math.log2(snr_per_channel + 1)
    )
    b1_aggregate_bias = c_b1 * window_length * math.log2(snr_grad + 1)
    b1_effective_control = (
        c_b1 * gradient_fraction * window_length * math.log2(snr_grad + 1)
    )
    if b1_model == "max_envelope":
        if b1_fraction_channels >= b1_aggregate_bias:
            b1, b1_selected = b1_fraction_channels, "fraction_channels"
        else:
            b1, b1_selected = b1_aggregate_bias, "aggregate_bias"
    elif b1_model == "fraction_channels":
        b1, b1_selected = b1_fraction_channels, "fraction_channels"
    elif b1_model == "aggregate_bias":
        b1, b1_selected = b1_aggregate_bias, "aggregate_bias"
    else:  # effective_control_v1
        b1, b1_selected = b1_effective_control, "effective_control_v1"

    # B2 — LR selection. R-symbol categorical channel at every K-step
    # decision point (#000036 §4).
    decisions_in_window = -(-window_length // lr_decision_interval)  # ceil
    b2 = c_b2 * decisions_in_window * math.log2(lr_grid_size)

    # B3 — Batch order. Bottou-Bousquet refinement: per-epoch contribution
    # bounded by 0.5 · log₂(N_b · σ_grad / ‖∇L_max‖) (#000036 §5).
    epochs_in_window = -(-window_length // steps_per_epoch)
    b3_per_epoch_factor = batches_per_epoch * gradient_noise_stddev / gradient_norm_max
    if b3_per_epoch_factor <= 1:
        # log₂ of a value ≤ 1 is ≤ 0; the bound becomes trivial (the
        # adversarial order can't move the parameter trajectory beyond
        # the noise floor). Floor at 0 — adversary can't do worse than
        # random shuffle in this regime.
        b3_per_epoch = 0.0
    else:
        b3_per_epoch = math.log2(b3_per_epoch_factor) / 2.0
    b3 = c_b3 * epochs_in_window * b3_per_epoch

    total = b1 + b2 + b3

    # Operator guidance — I_window is an UPPER BOUND. If it exceeds the
    # 256-bit threshold we do NOT know the adversary can steer 256 bits;
    # we know only that this bound is too loose to certify safety. The
    # wording reflects that (per dav1d review 2026-05-11).
    sha256_bits = CERTIFICATION_THRESHOLD_BITS
    if total <= 0:
        certification_status = "CERTIFIED_BY_BOUND"
        recommendation = (
            "I_window upper bound ≤ 0 (degenerate inputs); the bound is "
            "vacuously below the 256-bit certification threshold."
        )
    elif total >= sha256_bits:
        certification_status = "NOT_CERTIFIED_BY_BOUND"
        recommendation = (
            f"I_window upper bound ≈ {total:.1f} bits/window EXCEEDS the "
            f"SHA-256 ({sha256_bits} bit) certification threshold. This "
            f"conservative bound CANNOT CERTIFY M2's single-window "
            f"residual at this W — it does not prove the adversary can "
            f"steer {sha256_bits} bits, only that the bound is too loose "
            f"to certify safety. Reduce W (or reduce g / R / increase K, "
            f"or tighten C_B* empirically) until the certified bound is "
            f"< {sha256_bits} bits/window."
        )
    else:
        certification_status = "CERTIFIED_BY_BOUND"
        windows_to_brute_force = 2 ** (sha256_bits - total)
        recommendation = (
            f"I_window upper bound ≈ {total:.1f} bits/window. Certified "
            f"residual ≥ {sha256_bits - total:.1f} bits per independent "
            f"nonce window under this leakage model; a {sha256_bits}-bit "
            f"specific-target search remains bounded below by ≈ "
            f"2^{sha256_bits - total:.1f} ≈ {windows_to_brute_force:.2e} "
            f"independent windows (assumes fresh nonce per window, no "
            f"structural advantage beyond the bounded channel)."
        )

    return {
        "calculator_version": CALCULATOR_VERSION,
        "b1_model": b1_model,
        "b1_selected": b1_selected,
        "I_window_bits_upper_bound": round(total, 4),
        "B1_contribution": round(b1, 4),
        "B1_fraction_channels": round(b1_fraction_channels, 4),
        "B1_aggregate_bias": round(b1_aggregate_bias, 4),
        "B1_effective_control_v1": round(b1_effective_control, 4),
        "B2_contribution": round(b2, 4),
        "B3_contribution": round(b3, 4),
        "snr_grad": round(snr_grad, 4),
        "snr_per_channel": round(snr_per_channel, 4),
        "decisions_in_window": decisions_in_window,
        "epochs_in_window": epochs_in_window,
        "constants": {"C_B1": c_b1, "C_B2": c_b2, "C_B3": c_b3},
        "certification_status": certification_status,
        "certification_threshold_bits": sha256_bits,
        "model_assumptions": [
            "M2_nonce_per_window",
            "SHA256_random_oracle_baseline",
            f"B1_model_{b1_model}",
            "B3_gradient_noise_refinement",
        ],
        "inputs": {
            "gradient_fraction": gradient_fraction,
            "gradient_norm_max": gradient_norm_max,
            "gradient_noise_stddev": gradient_noise_stddev,
            "lr_decision_interval": lr_decision_interval,
            "lr_grid_size": lr_grid_size,
            "window_length": window_length,
            "batches_per_epoch": batches_per_epoch,
            "steps_per_epoch": steps_per_epoch,
            "c_b1": c_b1,
            "c_b2": c_b2,
            "c_b3": c_b3,
            "b1_model": b1_model,
        },
        "recommendation": recommendation,
    }


def main(argv: list[str] | None = None) -> int:
    p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    p.add_argument(
        "--gradient-fraction", type=float, required=True,
        help="g — gradient-control coefficient; in [0, 1], 0 = no T2 surface. Interpretation depends on --b1-model.",
    )
    p.add_argument(
        "--gradient-norm-max", type=float, required=True,
        help="‖∇L_max‖ — gradient-norm clipping ceiling (deployment param)",
    )
    p.add_argument(
        "--gradient-noise-stddev", type=float, required=True,
        help="σ_grad — per-step gradient noise stddev (SGD noise floor)",
    )
    p.add_argument(
        "--lr-decision-interval", type=int, required=True,
        help="K — steps between LR-schedule decision points",
    )
    p.add_argument(
        "--lr-grid-size", type=int, required=True,
        help="R — number of LR levels in the schedule grid",
    )
    p.add_argument(
        "--window-length", type=int, required=True,
        help="W — steps per nonce-window (M2 anchor re-randomization period)",
    )
    p.add_argument(
        "--batches-per-epoch", type=int, required=True,
        help="N_b — batch count per training epoch",
    )
    p.add_argument(
        "--steps-per-epoch", type=int, required=True,
        help="E — SGD steps per epoch (typically equals batches-per-epoch)",
    )
    p.add_argument(
        "--c-b1", type=float, default=1.0,
        help="C_B1 channel-efficiency constant; default 1.0 (data-processing-inequality ceiling)",
    )
    p.add_argument(
        "--c-b2", type=float, default=1.0,
        help="C_B2 constant; default 1.0",
    )
    p.add_argument(
        "--c-b3", type=float, default=1.0,
        help="C_B3 constant; default 1.0",
    )
    p.add_argument(
        "--b1-model", type=str, default=B1_MODEL_DEFAULT, choices=list(B1_MODELS),
        help=(
            "B1 gradient-bias model (#000036 §3.1). Default 'max_envelope' = "
            "worst-case envelope across the two g-interpretations; "
            "'fraction_channels' / 'aggregate_bias' = the individual terms; "
            "'effective_control_v1' = older non-worst-case formula."
        ),
    )

    args = p.parse_args(argv)
    try:
        result = t3_bound_bits(
            gradient_fraction=args.gradient_fraction,
            gradient_norm_max=args.gradient_norm_max,
            gradient_noise_stddev=args.gradient_noise_stddev,
            lr_decision_interval=args.lr_decision_interval,
            lr_grid_size=args.lr_grid_size,
            window_length=args.window_length,
            batches_per_epoch=args.batches_per_epoch,
            steps_per_epoch=args.steps_per_epoch,
            c_b1=args.c_b1,
            c_b2=args.c_b2,
            c_b3=args.c_b3,
            b1_model=args.b1_model,
        )
    except ValueError as e:
        print(f"input error: {e}", file=sys.stderr)
        return 2

    print(json.dumps(result, indent=2, ensure_ascii=False))
    return 0


if __name__ == "__main__":
    sys.exit(main())
