"""YaRN / RoPE context-extension sanity checks.

YaRN changes how the model *interprets* positions, not the per-token KV size,
so KV cache memory is driven by the target context length the user requests.
These helpers surface warnings when the target context exceeds the model's
training context (YaRN scaling required) and show the effective context as a
sanity check.
"""

from __future__ import annotations


def yarn_effective_context(training_ctx: int, rope_freq_scale: float) -> int:
    """Effective (interpolated) context given a rope_freq_scale.

    effective = training_ctx / rope_freq_scale (rope_freq_scale < 1 extends).
    """
    if rope_freq_scale <= 0:
        return training_ctx
    return int(round(training_ctx / rope_freq_scale))


def yarn_warnings(
    *,
    training_ctx: int,
    target_ctx: int,
    rope_freq_scale: float,
    yarn_ext_factor: float,
    yarn_attn_factor: float,
) -> list[str]:
    """Return human-readable warnings about the YaRN / context config."""
    warns: list[str] = []

    if target_ctx > training_ctx:
        warns.append(
            f"Target context {target_ctx} exceeds training context "
            f"{training_ctx}; YaRN/RoPE scaling is required for coherent "
            f"long-context output."
        )

    eff = yarn_effective_context(training_ctx, rope_freq_scale)
    if target_ctx > training_ctx and eff < target_ctx:
        warns.append(
            f"Effective context from rope_freq_scale={rope_freq_scale} is "
            f"{eff}, which is below the target {target_ctx}. The model may "
            f"not fully cover the requested window; lower rope_freq_scale "
            f"to interpolate more aggressively."
        )

    if target_ctx > training_ctx and yarn_ext_factor < 0.0:
        warns.append("yarn_ext_factor is negative; YaRN extrapolation is off.")

    if yarn_attn_factor < 1.0:
        warns.append(
            "yarn_attn_factor < 1.0 scales attention down; use ~1.0 for "
            "normal long-context YaRN (lower only for high-context YaRN fine-tunes)."
        )

    return warns