# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

"""Official multi-modal guidance for the LTX model family."""

from __future__ import annotations

import math
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any

import torch
from vllm.logger import init_logger

from vllm_omni.diffusion.distributed.parallel_state import (
    get_cfg_group as get_guidance_parallel_group,
)
from vllm_omni.diffusion.distributed.parallel_state import (
    get_classifier_free_guidance_rank as get_guidance_parallel_rank,
)
from vllm_omni.diffusion.distributed.parallel_state import (
    get_classifier_free_guidance_world_size as get_guidance_parallel_world_size,
)

logger = init_logger(__name__)

if TYPE_CHECKING:
    from .ltx2_denoise import LTXDenoiseContext, LTXForwardContext
    from .ltx2_latents import LTXAVState


@dataclass(frozen=True)
class LTXModalityGuidance:
    """Guidance parameters applied to one LTX output modality."""

    cfg_scale: float = 1.0
    stg_scale: float = 0.0
    modality_scale: float = 1.0
    rescale_scale: float = 0.0
    stg_blocks: tuple[int, ...] = ()

    @property
    def do_cfg(self) -> bool:
        return not math.isclose(self.cfg_scale, 1.0)

    @property
    def do_stg(self) -> bool:
        return not math.isclose(self.stg_scale, 0.0) and bool(self.stg_blocks)

    @property
    def do_modality_guidance(self) -> bool:
        return not math.isclose(self.modality_scale, 1.0)


@dataclass(frozen=True)
class LTXGuidanceSpec:
    """Video and audio guidance requested for one denoise phase."""

    video: LTXModalityGuidance = LTXModalityGuidance()
    audio: LTXModalityGuidance = LTXModalityGuidance()

    @classmethod
    def positive_only(cls) -> LTXGuidanceSpec:
        return cls()

    @property
    def do_cfg(self) -> bool:
        return self.video.do_cfg or self.audio.do_cfg

    @property
    def do_stg(self) -> bool:
        return self.video.do_stg or self.audio.do_stg

    @property
    def do_modality_guidance(self) -> bool:
        return self.video.do_modality_guidance or self.audio.do_modality_guidance

    @property
    def do_rescale(self) -> bool:
        return self.video.rescale_scale != 0.0 or self.audio.rescale_scale != 0.0


@dataclass(frozen=True)
class LTXDenoisePass:
    """One conditioned Transformer evaluation in a guidance batch."""

    name: str
    negative_video_context: bool = False
    negative_audio_context: bool = False


@dataclass(frozen=True)
class LTXGuidancePlan:
    """Concrete Transformer passes derived from a guidance specification."""

    spec: LTXGuidanceSpec
    passes: tuple[LTXDenoisePass, ...]

    @classmethod
    def build(cls, spec: LTXGuidanceSpec) -> LTXGuidancePlan:
        passes = [LTXDenoisePass("cond")]
        if spec.do_cfg:
            passes.append(
                LTXDenoisePass(
                    "uncond",
                    negative_video_context=spec.video.do_cfg,
                    negative_audio_context=spec.audio.do_cfg,
                )
            )
        if spec.do_stg:
            passes.append(LTXDenoisePass("ptb"))
        if spec.do_modality_guidance:
            passes.append(LTXDenoisePass("mod"))
        return cls(spec=spec, passes=tuple(passes))

    @property
    def names(self) -> tuple[str, ...]:
        return tuple(item.name for item in self.passes)


def x0_from_velocity(sample: torch.Tensor, velocity: torch.Tensor, sigma: torch.Tensor) -> torch.Tensor:
    """Convert velocity to x0 using the official fp32 arithmetic order."""
    sigma = sigma.to(torch.float32)
    return (sample.to(torch.float32) - velocity.to(torch.float32) * sigma).to(sample.dtype)


def velocity_from_x0(sample: torch.Tensor, x0: torch.Tensor, sigma: torch.Tensor) -> torch.Tensor:
    """Convert x0 back to velocity using the official fp32 arithmetic order."""
    # Official LTX materializes the scalar schedule value on the host before
    # dividing. CUDA tensor-scalar division can cross bfloat16 rounding bounds
    # even when the scalar has the same float32 value.
    sigma = sigma.to(torch.float32).item()
    return ((sample.to(torch.float32) - x0.to(torch.float32)) / sigma).to(sample.dtype)


def euler_step_from_velocity(
    sample: torch.Tensor,
    velocity: torch.Tensor,
    sigmas: torch.Tensor,
    step_index: int,
) -> torch.Tensor:
    sigma = sigmas[step_index].to(torch.float32)
    sigma_next = sigmas[step_index + 1].to(torch.float32)
    return (sample.float() + velocity.float() * (sigma_next - sigma)).to(sample.dtype)


def combine_guided_x0(
    *,
    cond: torch.Tensor,
    uncond_text: torch.Tensor | float,
    uncond_perturbed: torch.Tensor | float,
    uncond_modality: torch.Tensor | float,
    guidance: LTXModalityGuidance,
    rescale_token_count: int | None = None,
) -> torch.Tensor:
    dtype = cond.dtype
    cond = cond.float()
    uncond_text = uncond_text.float() if isinstance(uncond_text, torch.Tensor) else uncond_text
    uncond_perturbed = uncond_perturbed.float() if isinstance(uncond_perturbed, torch.Tensor) else uncond_perturbed
    uncond_modality = uncond_modality.float() if isinstance(uncond_modality, torch.Tensor) else uncond_modality
    pred = cond + (guidance.cfg_scale - 1) * (cond - uncond_text)
    if guidance.do_stg:
        pred = pred + guidance.stg_scale * (cond - uncond_perturbed)
    pred = pred + (guidance.modality_scale - 1) * (cond - uncond_modality)
    if guidance.rescale_scale != 0.0:
        # The pinned official one-stage entry runs one generated sample and
        # uses its full-tensor standard deviation. In Omni, dimension 0 may
        # instead contain independently scheduled requests. Reduce each item
        # separately so its result cannot depend on unrelated co-batched work.
        stats_cond = cond
        stats_pred = pred
        if rescale_token_count is not None:
            if cond.ndim < 2 or not 0 < rescale_token_count <= cond.shape[1]:
                raise ValueError(
                    f"Guidance rescale token count must be in [1, {cond.shape[1]}], got {rescale_token_count}."
                )
            stats_cond = cond[:, :rescale_token_count]
            stats_pred = pred[:, :rescale_token_count]
        reduce_dims = tuple(range(1, stats_pred.ndim))
        cond_std = stats_cond.std(dim=reduce_dims, keepdim=True)
        pred_std = stats_pred.std(dim=reduce_dims, keepdim=True)
        eps = torch.finfo(pred.dtype).eps
        factor = cond_std / pred_std.clamp_min(eps)
        factor = torch.where(pred_std > eps, factor, torch.ones_like(factor))
        factor = guidance.rescale_scale * factor + (1 - guidance.rescale_scale)
        pred = pred * factor
    return pred.to(dtype)


def _repeat_batch(tensor: torch.Tensor, repeats: int) -> torch.Tensor:
    if repeats == 1:
        return tensor
    return tensor.repeat((repeats,) + (1,) * (tensor.ndim - 1))


def _mask_for_pass(
    pass_count: int,
    batch_size: int,
    pass_index: int,
    reference: torch.Tensor,
) -> torch.Tensor:
    mask = reference.new_ones((pass_count * batch_size, 1, 1))
    start = pass_index * batch_size
    mask[start : start + batch_size] = 0
    return mask


def build_perturbation_kwargs(
    plan: LTXGuidancePlan,
    batch_size: int,
    reference: torch.Tensor,
) -> dict[str, Any]:
    kwargs: dict[str, Any] = {}
    if "ptb" in plan.names:
        mask = _mask_for_pass(len(plan.passes), batch_size, plan.names.index("ptb"), reference)
        if plan.spec.video.do_stg:
            kwargs["video_self_attention_mask"] = mask
            kwargs["video_self_attention_blocks"] = plan.spec.video.stg_blocks
        if plan.spec.audio.do_stg:
            kwargs["audio_self_attention_mask"] = mask
            kwargs["audio_self_attention_blocks"] = plan.spec.audio.stg_blocks
    if "mod" in plan.names:
        mask = _mask_for_pass(len(plan.passes), batch_size, plan.names.index("mod"), reference)
        kwargs["a2v_cross_attention_mask"] = mask
        kwargs["v2a_cross_attention_mask"] = mask
    return kwargs


class LTXGuidanceExecutor:
    """Build and execute official LTX multi-guidance Transformer passes."""

    @staticmethod
    def validate_guidance_world_size(plan: LTXGuidancePlan, guidance_world_size: int) -> None:
        del plan
        if guidance_world_size < 1:
            raise ValueError(f"LTX guidance parallelism requires a positive parallel size, got {guidance_world_size}.")

    @staticmethod
    def _parallel_assignments(pass_count: int, world_size: int) -> list[list[int]]:
        assignments: list[list[int]] = [[] for _ in range(world_size)]
        for pass_index in range(pass_count):
            assignments[pass_index % world_size].append(pass_index)
        return assignments

    @classmethod
    def _model_pass_count(
        cls,
        plan: LTXGuidancePlan,
        guidance_parallel_ready: bool,
        guidance_world_size: int,
    ) -> int:
        if not guidance_parallel_ready:
            return len(plan.passes)
        assignments = cls._parallel_assignments(len(plan.passes), guidance_world_size)
        return max(len(indices) for indices in assignments)

    @classmethod
    def warn_if_imbalanced(
        cls,
        plan: LTXGuidancePlan,
        guidance_world_size: int,
        phase_name: str,
    ) -> None:
        pass_count = len(plan.passes)
        if guidance_world_size <= 1 or pass_count % guidance_world_size == 0 or get_guidance_parallel_rank() != 0:
            return
        slots_per_rank = math.ceil(pass_count / guidance_world_size)
        total_slots = slots_per_rank * guidance_world_size
        wasted_slots = total_slots - pass_count
        utilization = 100.0 * pass_count / total_slots
        logger.warning_once(
            "LTX guidance parallelism is imbalanced for phase %r: %d guidance passes across %d ranks. "
            "%d padded/discarded rank slot(s); expected guidance-slot utilization is %.1f%%.",
            phase_name,
            pass_count,
            guidance_world_size,
            wasted_slots,
            utilization,
        )

    @staticmethod
    def prepare_denoise_context(
        plan: LTXGuidancePlan,
        guidance_parallel_ready: bool,
        guidance_world_size: int,
        denoise_ctx: LTXDenoiseContext,
    ) -> LTXDenoiseContext:
        model_pass_count = LTXGuidanceExecutor._model_pass_count(
            plan,
            guidance_parallel_ready,
            guidance_world_size,
        )
        if model_pass_count > 1:
            denoise_ctx.video_coords = _repeat_batch(denoise_ctx.video_coords, model_pass_count)
            denoise_ctx.audio_coords = _repeat_batch(denoise_ctx.audio_coords, model_pass_count)
            if denoise_ctx.audio_attention_mask is not None:
                denoise_ctx.audio_attention_mask = _repeat_batch(
                    denoise_ctx.audio_attention_mask,
                    model_pass_count,
                )
        return denoise_ctx

    @staticmethod
    def timestep_kwargs(
        ts: torch.Tensor,
        video_token_count: int,
        audio_token_count: int,
        *,
        expand_for_sequence_parallel: bool = False,
    ) -> dict[str, torch.Tensor]:
        timestep = ts
        audio_timestep = ts
        if expand_for_sequence_parallel:
            timestep = ts.reshape(-1, 1).expand(-1, video_token_count)
            audio_timestep = ts.reshape(-1, 1).expand(-1, audio_token_count)
        return {
            "timestep": timestep,
            "audio_timestep": audio_timestep,
            "sigma": ts,
            "audio_sigma": ts,
        }

    @staticmethod
    def combine_cfg_velocity(
        sample: torch.Tensor,
        positive_velocity: torch.Tensor,
        negative_velocity: torch.Tensor,
        sigma: torch.Tensor,
        guidance_scale: float,
        *,
        model_sigma: torch.Tensor | None = None,
    ) -> torch.Tensor:
        guidance = LTXModalityGuidance(cfg_scale=guidance_scale)
        model_sigma = sigma if model_sigma is None else model_sigma
        guided_x0 = combine_guided_x0(
            cond=x0_from_velocity(sample, positive_velocity, model_sigma),
            uncond_text=x0_from_velocity(sample, negative_velocity, model_sigma),
            uncond_perturbed=0.0,
            uncond_modality=0.0,
            guidance=guidance,
        )
        return velocity_from_x0(sample, guided_x0, sigma)

    @staticmethod
    def _guide_modality(
        sample: torch.Tensor,
        splits: dict[str, torch.Tensor],
        sigma: torch.Tensor,
        guidance: LTXModalityGuidance,
        *,
        model_sigma: torch.Tensor | None = None,
        preserve_positive_velocity: bool = False,
        rescale_token_count: int | None = None,
    ) -> torch.Tensor:
        if preserve_positive_velocity and tuple(splits) == ("cond",):
            return splits["cond"]
        model_sigma = sigma if model_sigma is None else model_sigma
        # Official guidance reduces contiguous BSC tensors. The fused
        # transformer output may be channel-major after splitting, which
        # changes fp32 reduction order and can cross bf16 rounding bounds.
        x0 = {name: x0_from_velocity(sample, value, model_sigma).contiguous() for name, value in splits.items()}
        guided = combine_guided_x0(
            cond=x0["cond"],
            uncond_text=x0.get("uncond", 0.0),
            uncond_perturbed=x0.get("ptb", 0.0),
            uncond_modality=x0.get("mod", 0.0),
            guidance=guidance,
            rescale_token_count=rescale_token_count,
        )
        return velocity_from_x0(sample, guided, sigma)

    def predict_parallel_guidance(
        self,
        pipeline: Any,
        plan: LTXGuidancePlan,
        index: int,
        timestep: torch.Tensor,
        state: LTXAVState,
        forward_ctx: LTXForwardContext,
        denoise_ctx: LTXDenoiseContext,
        preserve_positive_velocity: bool = False,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        guidance_world_size = get_guidance_parallel_world_size()
        self.validate_guidance_world_size(plan, guidance_world_size)
        guidance_rank = get_guidance_parallel_rank()
        assignments = self._parallel_assignments(len(plan.passes), guidance_world_size)
        local_pass_indices = assignments[guidance_rank]
        model_pass_count = max(len(indices) for indices in assignments)

        # Every rank executes and gathers the same number of slots. Shorter
        # assignments use a conditional pass whose output is discarded.
        padded_pass_indices: list[int | None] = local_pass_indices + [None] * (
            model_pass_count - len(local_pass_indices)
        )
        local_passes = tuple(
            plan.passes[0] if pass_index is None else plan.passes[pass_index] for pass_index in padded_pass_indices
        )
        local_plan = LTXGuidancePlan(spec=plan.spec, passes=local_passes)

        prompt = forward_ctx.prompt_context
        video_contexts: list[torch.Tensor] = []
        audio_contexts: list[torch.Tensor] = []
        for denoise_pass in local_passes:
            video_context = (
                prompt.negative_connector_prompt_embeds
                if denoise_pass.negative_video_context
                else prompt.positive_connector_prompt_embeds
            )
            audio_context = (
                prompt.negative_connector_audio_prompt_embeds
                if denoise_pass.negative_audio_context
                else prompt.positive_connector_audio_prompt_embeds
            )
            if video_context is None or audio_context is None:
                raise ValueError("Negative prompt context is required when LTX CFG is enabled.")
            video_contexts.append(video_context)
            audio_contexts.append(audio_context)

        video_input = _repeat_batch(state.video, model_pass_count).to(prompt.positive_connector_prompt_embeds.dtype)
        audio_input = _repeat_batch(state.audio, model_pass_count).to(
            prompt.positive_connector_audio_prompt_embeds.dtype
        )
        attention_kwargs = dict(forward_ctx.attention_kwargs or {})
        perturbations = build_perturbation_kwargs(local_plan, state.video.shape[0], video_input)
        if perturbations:
            attention_kwargs["ltx_perturbation_kwargs"] = perturbations
        kwargs = pipeline._build_transformer_kwargs(
            forward_ctx,
            denoise_ctx,
            hidden_states=video_input,
            audio_hidden_states=audio_input,
            encoder_hidden_states=torch.cat(video_contexts),
            audio_encoder_hidden_states=torch.cat(audio_contexts),
            encoder_attention_mask=None,
            audio_encoder_attention_mask=None,
            ts=timestep.expand(video_input.shape[0]),
            attention_kwargs=attention_kwargs,
        )
        with pipeline._transformer_cache_context("guided"):
            local_video, local_audio = pipeline.transformer(**kwargs)
        local_video_slots = local_video.chunk(model_pass_count)
        local_audio_slots = local_audio.chunk(model_pass_count)

        group = get_guidance_parallel_group()
        gathered_video_slots = [group.all_gather(value, separate_tensors=True) for value in local_video_slots]
        gathered_audio_slots = [group.all_gather(value, separate_tensors=True) for value in local_audio_slots]

        video_splits: dict[str, torch.Tensor] = {}
        audio_splits: dict[str, torch.Tensor] = {}
        for pass_index, denoise_pass in enumerate(plan.passes):
            owner_rank = pass_index % guidance_world_size
            slot_index = pass_index // guidance_world_size
            video_splits[denoise_pass.name] = gathered_video_slots[slot_index][owner_rank]
            audio_splits[denoise_pass.name] = gathered_audio_slots[slot_index][owner_rank]

        video_sigma = pipeline.scheduler.sigmas[index]
        return (
            self._guide_modality(
                state.video,
                video_splits,
                video_sigma,
                plan.spec.video,
                model_sigma=pipeline._video_guidance_model_sigma(video_sigma, denoise_ctx),
                preserve_positive_velocity=preserve_positive_velocity,
            ),
            self._guide_modality(
                state.audio,
                audio_splits,
                forward_ctx.audio_scheduler.sigmas[index],
                plan.spec.audio,
                preserve_positive_velocity=preserve_positive_velocity,
                rescale_token_count=forward_ctx.original_audio_num_frames,
            ),
        )

    def predict_noise(
        self,
        pipeline: Any,
        plan: LTXGuidancePlan,
        index: int,
        timestep: torch.Tensor,
        state: LTXAVState,
        forward_ctx: LTXForwardContext,
        denoise_ctx: LTXDenoiseContext,
        preserve_positive_velocity: bool = False,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        if forward_ctx.guidance_parallel_ready:
            return self.predict_parallel_guidance(
                pipeline,
                plan,
                index,
                timestep,
                state,
                forward_ctx,
                denoise_ctx,
                preserve_positive_velocity=preserve_positive_velocity,
            )

        prompt = forward_ctx.prompt_context
        video_contexts: list[torch.Tensor] = []
        audio_contexts: list[torch.Tensor] = []
        for denoise_pass in plan.passes:
            video_context = (
                prompt.negative_connector_prompt_embeds
                if denoise_pass.negative_video_context
                else prompt.positive_connector_prompt_embeds
            )
            audio_context = (
                prompt.negative_connector_audio_prompt_embeds
                if denoise_pass.negative_audio_context
                else prompt.positive_connector_audio_prompt_embeds
            )
            if video_context is None or audio_context is None:
                raise ValueError("Negative prompt context is required when LTX CFG is enabled.")
            video_contexts.append(video_context)
            audio_contexts.append(audio_context)

        pass_count = len(plan.passes)
        video_input = _repeat_batch(state.video, pass_count).to(prompt.positive_connector_prompt_embeds.dtype)
        audio_input = _repeat_batch(state.audio, pass_count).to(prompt.positive_connector_audio_prompt_embeds.dtype)
        attention_kwargs = dict(forward_ctx.attention_kwargs or {})
        perturbations = build_perturbation_kwargs(plan, state.video.shape[0], video_input)
        if perturbations:
            attention_kwargs["ltx_perturbation_kwargs"] = perturbations
        kwargs = pipeline._build_transformer_kwargs(
            forward_ctx,
            denoise_ctx,
            hidden_states=video_input,
            audio_hidden_states=audio_input,
            encoder_hidden_states=torch.cat(video_contexts),
            audio_encoder_hidden_states=torch.cat(audio_contexts),
            encoder_attention_mask=None,
            audio_encoder_attention_mask=None,
            ts=timestep.expand(video_input.shape[0]),
            attention_kwargs=attention_kwargs,
        )
        with pipeline._transformer_cache_context("guided"):
            video_velocity, audio_velocity = pipeline.transformer(**kwargs)
        video_splits = dict(zip(plan.names, video_velocity.chunk(pass_count), strict=True))
        audio_splits = dict(zip(plan.names, audio_velocity.chunk(pass_count), strict=True))

        return (
            self._guide_modality(
                state.video,
                video_splits,
                pipeline.scheduler.sigmas[index],
                plan.spec.video,
                model_sigma=pipeline._video_guidance_model_sigma(
                    pipeline.scheduler.sigmas[index],
                    denoise_ctx,
                ),
                preserve_positive_velocity=preserve_positive_velocity,
            ),
            self._guide_modality(
                state.audio,
                audio_splits,
                forward_ctx.audio_scheduler.sigmas[index],
                plan.spec.audio,
                preserve_positive_velocity=preserve_positive_velocity,
                rescale_token_count=forward_ctx.original_audio_num_frames,
            ),
        )


LTX_GUIDANCE_EXECUTOR = LTXGuidanceExecutor()

# Kept as a narrow compatibility helper for CFG-parallel tests and extensions.
combine_velocity_via_x0 = LTXGuidanceExecutor.combine_cfg_velocity