"""Audio reference conditioning items."""

from __future__ import annotations

import torch

from ltx_core.conditioning.cond_channels import extend_cond_channels
from ltx_core.conditioning.mask_utils import update_attention_mask
from ltx_core.tools import LatentTools
from ltx_core.types import LatentState


class AudioConditionByReferenceLatent:
    """Append patchified reference audio tokens after the target audio sequence.
    Mirrors :class:`ltx_core.conditioning.types.reference_video_cond.VideoConditionByReferenceLatent`
    but for audio. The reference tokens are appended so the target audio tokens stay
    in the first ``num_noisy_tokens`` positions and can be kept by
    :meth:`ltx_core.tools.LatentTools.clear_conditioning`.
    Args:
        patchified: Patchified reference latent ``[B, T_ref, C]``.
        positions: RoPE positions for reference tokens, ``[B, 1, T_ref, 2]``.
        strength: 1.0 keeps reference clean; 0.0 would fully denoise it.
    """

    def __init__(
        self,
        patchified: torch.Tensor,
        positions: torch.Tensor,
        strength: float = 1.0,
    ) -> None:
        self.patchified = patchified
        self.positions = positions.to(dtype=torch.float32)
        self.strength = strength

    def apply_to(self, latent_state: LatentState, latent_tools: LatentTools) -> LatentState:
        tokens = self.patchified
        denoise_mask = torch.full(
            size=(*tokens.shape[:2], 1),
            fill_value=1.0 - self.strength,
            device=tokens.device,
            dtype=tokens.dtype,
        )

        new_attention_mask = update_attention_mask(
            latent_state=latent_state,
            attention_mask=None,
            num_noisy_tokens=latent_tools.patchifier.get_token_count(latent_tools.target_shape),
            num_new_tokens=tokens.shape[1],
            batch_size=tokens.shape[0],
            device=tokens.device,
            dtype=tokens.dtype,
        )

        return LatentState(
            latent=torch.cat([latent_state.latent, torch.zeros_like(tokens)], dim=1),
            denoise_mask=torch.cat([latent_state.denoise_mask, denoise_mask], dim=1),
            positions=torch.cat([latent_state.positions, self.positions], dim=2),
            clean_latent=torch.cat([latent_state.clean_latent, tokens], dim=1),
            attention_mask=new_attention_mask,
            cond_channels=extend_cond_channels(latent_state, tokens.shape[1]),
        )
