"""Helpers for in-context conditioning channels (``LatentState.cond_channels``).

Conditioning channels are extra per-token input features (e.g. SCAIL-2 mask
channels) that ride alongside the latent in patchified token space ``(B, T, C)``
and are concatenated onto the latent before the model's first projection. When
present they must stay length-aligned with the token sequence, so any
conditioning item that *appends* tokens must also extend the channels.
"""

from __future__ import annotations

import torch

from ltx_core.types import LatentState


def extend_cond_channels(latent_state: LatentState, num_new_tokens: int) -> torch.Tensor | None:
    """Return ``cond_channels`` extended with ``num_new_tokens`` zero rows, or ``None``.

    Appended tokens (reference / driving / keyframe) carry no in-context signal by default, so they
    are padded with zeros to preserve the ``T``-alignment invariant. Returns ``None`` unchanged when
    the state has no conditioning channels (the standard case), keeping non-SCAIL pipelines untouched.
    """
    cond_channels = latent_state.cond_channels
    if cond_channels is None:
        return None
    zeros = cond_channels.new_zeros(cond_channels.shape[0], num_new_tokens, cond_channels.shape[2])
    return torch.cat([cond_channels, zeros], dim=1)
