"""Soft prefix-attention masks for RTC guided inpainting (paper Sec. 3.2 / App. A).

    Black, Galliker, Levine. "Real-Time Execution of Action Chunking Flow
    Policies." NeurIPS 2025. arXiv:2506.07339.

RECONSTRUCTED MODULE. The original was lost; this restores it from the paper and
from how :mod:`rtc.sampler` calls it (``start=inference_delay``,
``end=horizon-execution_horizon``, ``total=horizon``). The three properties that
make RTC *correct* are asserted by ``test_masking.py``:

  1. ``w[i] == 1`` for ``i < start``  -- the first ``d`` actions are already
     committed to the robot before the new chunk lands, so they are a HARD
     constraint, not a preference. Getting this wrong breaks the guarantee.
  2. ``w`` is monotonically non-increasing -- confidence in the old chunk decays
     with distance into the future.
  3. ``w[i] == 0`` for ``i >= end`` -- past the overlap there is no previous
     chunk to agree with, so the new one is unconstrained.

The exact interior curve between ``start`` and ``end`` is a smoothness
hyperparameter, not a correctness property. ``linear`` reproduces the reference
implementation's documented example exactly::

    get_prefix_weights(start=2, end=6, total=10, schedule="linear")
      -> [1, 1, 4/5, 3/5, 2/5, 1/5, 0, 0, 0, 0]

``exp`` is a reconstruction: same endpoints and monotonicity, but a convex decay
that holds the old chunk more tightly just after the frozen prefix and releases
it faster near ``end``. If you need bit-exact parity with the authors' ``exp``,
compare against Physical-Intelligence/real-time-chunking-kinetix ``src/model.py``
and adjust :func:`_exp_curve` only -- nothing else depends on its shape.

``zeros`` is the naive hard-masking baseline the paper compares against: freeze
the prefix, ignore the overlap entirely.
"""

from __future__ import annotations

from typing import Literal

import numpy as np

PrefixAttentionSchedule = Literal["ones", "zeros", "linear", "exp"]

_SCHEDULES = ("ones", "zeros", "linear", "exp")


def _exp_curve(u: np.ndarray) -> np.ndarray:
    """Convex remap of the linear ramp ``u`` in [0, 1], fixing u=0 -> 0, u=1 -> 1.

    ``u * expm1(u) / (e - 1)``: monotone increasing, so weights stay ordered.
    """
    return u * np.expm1(u) / (np.e - 1.0)


def get_prefix_weights_np(
    start: int,
    end: int,
    total: int,
    schedule: PrefixAttentionSchedule = "exp",
) -> np.ndarray:
    """Soft mask ``W`` of shape ``(total,)``, float64 numpy.

    Args:
        start: ``d``, the inference delay. Indices ``< start`` are frozen (w = 1).
        end: ``H - s``, the prefix attention horizon. Indices ``>= end`` get w = 0.
        total: ``H``, the chunk length.
        schedule: see the module docstring.
    """
    if schedule not in _SCHEDULES:
        raise ValueError(f"unknown schedule {schedule!r}; pick one of {_SCHEDULES}")
    if total <= 0:
        raise ValueError(f"total must be positive, got {total}")
    # A delay longer than the overlap would ask for a frozen prefix the mask
    # cannot express; the reference clamps rather than failing.
    start = int(min(start, end))
    end = int(min(max(end, 0), total))
    i = np.arange(total)

    if schedule == "ones":
        w = np.ones(total, dtype=np.float64)
    elif schedule == "zeros":
        w = (i < start).astype(np.float64)          # hard mask: prefix only
    else:
        span = end - start + 1                      # matches the reference example
        w = np.clip((end - i) / max(span, 1), 0.0, 1.0)
        if schedule == "exp":
            w = _exp_curve(w)
        w = np.where(i < start, 1.0, w)             # property 1: hard prefix

    return np.where(i >= end, 0.0, w)               # property 3


def get_prefix_weights(
    start: int,
    end: int,
    total: int,
    schedule: PrefixAttentionSchedule = "exp",
    device=None,
    dtype=None,
):
    """Torch flavour of :func:`get_prefix_weights_np` (what ``rtc.sampler`` uses)."""
    import torch

    w = get_prefix_weights_np(start, end, total, schedule)
    return torch.as_tensor(w, device=device, dtype=dtype or torch.float32)


def get_prefix_weights_jax(
    start: int,
    end: int,
    total: int,
    schedule: PrefixAttentionSchedule = "exp",
    dtype=None,
):
    """JAX flavour, for the JAX/orbax dual-expert backend (``rtc.sampler_jax``)."""
    import jax.numpy as jnp

    w = get_prefix_weights_np(start, end, total, schedule)
    return jnp.asarray(w, dtype=dtype or jnp.float32)


def guidance_weight(tau: float, beta: float) -> float:
    """PiGDM guidance weight ``min(beta, (1 - tau) / (tau * r2_tau))`` (paper Eqs. 2/4).

    Lives here rather than in a backend sampler so the torch and JAX paths cannot
    drift apart::

        inv_r2 = (tau**2 + (1 - tau)**2) / (1 - tau)**2      # 1 / r2_tau   (Eq. 4)
        c      = (1 - tau) / tau                             # -> inf at tau == 0
        weight = min(c * inv_r2, beta)

    At ``tau == 0`` the unclipped weight is infinite (the score is undefined at pure
    noise), so ``beta`` is what makes it finite -- exactly its role in App. A.2.
    """
    if tau <= 0.0:
        return beta
    inv_r2 = (tau**2 + (1.0 - tau) ** 2) / ((1.0 - tau) ** 2)
    c = (1.0 - tau) / tau
    return min(c * inv_r2, beta)