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

"""
Base pipeline class for Diffusion models with shared CFG functionality.
"""

from abc import ABCMeta
from typing import Any

import torch
from vllm.logger import init_logger
from vllm.sequence import IntermediateTensors

from vllm_omni.diffusion.distributed.parallel_state import (
    get_cfg_group,
    get_classifier_free_guidance_rank,
    get_classifier_free_guidance_world_size,
    is_cfg_group_initialized,
)

logger = init_logger(__name__)


def _get_cfg_world_size_or_one() -> int:
    """Return the CFG world size, defaulting only when no group exists.

    Diffusion workers initialize the CFG process group even when its world
    size is one.  Pipelines are also invoked directly by unit tests and some
    offline integrations, though, where no distributed groups exist. Treat
    that explicit case like a one-rank CFG group. Errors from an existing CFG
    coordinator must propagate instead of silently selecting sequential CFG.
    """
    if not is_cfg_group_initialized():
        return 1
    return get_classifier_free_guidance_world_size()


def _wrap(pred: torch.Tensor | tuple[torch.Tensor, ...]) -> tuple[torch.Tensor, ...]:
    """Normalize prediction to tuple form."""
    return pred if isinstance(pred, tuple) else (pred,)


def _unwrap(pred: tuple[torch.Tensor, ...]) -> torch.Tensor | tuple[torch.Tensor, ...]:
    """Unwrap single-element tuple to plain tensor; keep multi-element as tuple."""
    return pred[0] if len(pred) == 1 else pred


def _slice_pred(pred: tuple[torch.Tensor, ...], output_slice: int) -> tuple[torch.Tensor, ...]:
    """Slice each element along dim 1."""
    return tuple(p[:, :output_slice] for p in pred)


def _dispatch_branches(n_branches: int, n_ranks: int) -> list[list[int]]:
    """
    Round-robin dispatch N branches to M ranks.

    Rule: branch i → rank (i % n_ranks).

    Examples:
        _dispatch_branches(3, 2) -> [[0, 2], [1]]
        _dispatch_branches(3, 3) -> [[0], [1], [2]]
        _dispatch_branches(4, 2) -> [[0, 2], [1, 3]]
        _dispatch_branches(4, 4) -> [[0], [1], [2], [3]]
    """
    assignments: list[list[int]] = [[] for _ in range(n_ranks)]
    for i in range(n_branches):
        assignments[i % n_ranks].append(i)
    return assignments


class CFGParallelMixin(metaclass=ABCMeta):
    """
    Base Mixin class for Diffusion pipelines providing shared CFG methods.

    All pipelines should inherit from this class to reuse
    classifier-free guidance logic.

    CFG Parallel Architecture:
        When cfg_world_size > 1, each rank computes one branch (positive or
        negative), then all_gather exchanges results. All ranks then compute
        the CFG combine and scheduler step locally — no broadcast needed
        because the operations are deterministic.

    Multi-output models:
        Models that return tuple from predict_noise() (e.g., video + audio)
        should override combine_cfg_noise() to define per-element combine logic,
        and set self.scheduler to a composite scheduler that handles tuples.
    """

    def predict_noise_maybe_with_cfg(
        self,
        do_true_cfg: bool,
        true_cfg_scale: float,
        positive_kwargs: dict[str, Any],
        negative_kwargs: dict[str, Any] | None,
        cfg_normalize: bool = True,
        output_slice: int | None = None,
        kwargs: dict[str, Any] | None = None,
    ) -> torch.Tensor | tuple[torch.Tensor, ...]:
        """
        Predict noise with optional classifier-free guidance.

        Args:
            do_true_cfg: Whether to apply CFG
            true_cfg_scale: CFG scale factor
            positive_kwargs: Kwargs for positive/conditional prediction
            negative_kwargs: Kwargs for negative/unconditional prediction
            cfg_normalize: Whether to normalize CFG output (default: True)
            output_slice: If set, slice each output to [:, :output_slice] for image editing
            kwargs: Optional extra context for custom combine implementations.

        Returns:
            Predicted noise tensor or tuple of tensors.
            In CFG parallel mode, result is valid on ALL ranks (not just rank 0).

        Note:
            For multi-output models (e.g., video + audio where predict_noise
            returns a tuple), override combine_cfg_noise() for per-element CFG
            logic and set self.scheduler to a composite scheduler.
        """
        if do_true_cfg:
            # Automatically detect CFG parallel configuration
            cfg_parallel_ready = _get_cfg_world_size_or_one() > 1

            if cfg_parallel_ready:
                cfg_group = get_cfg_group()
                cfg_rank = get_classifier_free_guidance_rank()

                # Each rank computes one branch
                if cfg_rank == 0:
                    logger.debug("CFG Parallel: Rank 0 computing positive branch")
                    local_pred = _wrap(self.predict_noise(**positive_kwargs))
                else:
                    logger.debug("CFG Parallel: Rank %d computing negative branch", cfg_rank)
                    local_pred = _wrap(self.predict_noise(**negative_kwargs))

                if output_slice is not None:
                    local_pred = _slice_pred(local_pred, output_slice)

                # All-gather each element, reconstruct positive/negative tuples
                gathered = [cfg_group.all_gather(p, separate_tensors=True) for p in local_pred]
                positive_noise_pred = tuple(g[0] for g in gathered)
                negative_noise_pred = tuple(g[1] for g in gathered)

                # All ranks compute combine (deterministic, same result)
                return self.combine_cfg_noise(
                    positive_noise_pred,
                    negative_noise_pred,
                    true_cfg_scale,
                    cfg_normalize,
                    **({} if kwargs is None else {"kwargs": kwargs}),
                )
            else:
                # Sequential CFG: compute both positive and negative
                positive_noise_pred = _wrap(self.predict_noise(**positive_kwargs))
                negative_noise_pred = _wrap(self.predict_noise(**negative_kwargs))

                if output_slice is not None:
                    positive_noise_pred = _slice_pred(positive_noise_pred, output_slice)
                    negative_noise_pred = _slice_pred(negative_noise_pred, output_slice)

                return self.combine_cfg_noise(
                    positive_noise_pred,
                    negative_noise_pred,
                    true_cfg_scale,
                    cfg_normalize,
                    **({} if kwargs is None else {"kwargs": kwargs}),
                )
        else:
            # No CFG: only compute positive/conditional prediction
            pred = self.predict_noise(**positive_kwargs)
            if output_slice is not None:
                pred = _unwrap(_slice_pred(_wrap(pred), output_slice))
            return pred

    def cfg_normalize_function(self, noise_pred: torch.Tensor, comb_pred: torch.Tensor) -> torch.Tensor:
        """
        Normalize the combined noise prediction.

        Args:
            noise_pred: positive noise prediction
            comb_pred: combined noise prediction after CFG

        Returns:
            Normalized noise prediction tensor
        """
        cond_norm = torch.norm(noise_pred, dim=-1, keepdim=True)
        noise_norm = torch.norm(comb_pred, dim=-1, keepdim=True)
        noise_pred = comb_pred * (cond_norm / noise_norm)
        return noise_pred

    def combine_cfg_noise(
        self,
        positive_noise_pred: torch.Tensor | tuple[torch.Tensor, ...],
        negative_noise_pred: torch.Tensor | tuple[torch.Tensor, ...],
        true_cfg_scale: float,
        cfg_normalize: bool = False,
        kwargs: dict[str, Any] | None = None,
    ) -> torch.Tensor | tuple[torch.Tensor, ...]:
        """
        Combine conditional and unconditional noise predictions with CFG.

        Accepts both plain tensors (backward-compatible, used by LTX2 etc.)
        and tuples (multi-output models). Default implementation applies the
        standard CFG formula to every element.

        Multi-output models can override this to apply different logic per element.

        Example override for a model returning (video_pred, audio_pred)::

            def combine_cfg_noise(self, positive_noise_pred, negative_noise_pred, scale, normalize):
                (video_pos, audio_pos) = positive_noise_pred
                (video_neg, audio_neg) = negative_noise_pred
                video_combined = super().combine_cfg_noise(video_pos, video_neg, scale, normalize)
                return (video_combined, audio_pos)  # audio: positive only, no CFG

        Args:
            positive_noise_pred: Positive/conditional prediction(s) — Tensor or tuple
            negative_noise_pred: Negative/unconditional prediction(s) — Tensor or tuple
            true_cfg_scale: CFG scale factor
            cfg_normalize: Whether to normalize the combined prediction (default: False)
            kwargs: Optional extra context for custom combine implementations.

        Returns:
            Combined noise prediction(s) — same type as inputs
        """
        pos_t = _wrap(positive_noise_pred)
        neg_t = _wrap(negative_noise_pred)

        results = []
        for p, n in zip(pos_t, neg_t):
            comb = n + true_cfg_scale * (p - n)
            if cfg_normalize:
                comb = self.cfg_normalize_function(p, comb)
            results.append(comb)
        return _unwrap(tuple(results))

    # ── N-branch CFG interface (for 3+ branch models) ──

    def predict_noise_with_multi_branch_cfg(
        self,
        do_true_cfg: bool,
        true_cfg_scale: float | dict[str, float],
        branches_kwargs: list[dict[str, Any]],
        cfg_normalize: bool = False,
        output_slice: int | None = None,
    ) -> torch.Tensor | tuple[torch.Tensor, ...]:
        """
        Predict noise with N-branch CFG dispatch across M GPUs.

        This is the multi-branch counterpart of predict_noise_maybe_with_cfg().
        Use this for models with 3 or more CFG branches (e.g., OmniGen2, Bagel).
        Existing 2-branch models should continue using
        predict_noise_maybe_with_cfg().

        Args:
            do_true_cfg: Whether to apply CFG.
            true_cfg_scale: CFG scale factor (passed to combine_multi_branch_cfg_noise).
            branches_kwargs: List of N dicts, each containing kwargs for one
                predict_noise() call. branches_kwargs[0] is always the
                positive/conditional branch.
            cfg_normalize: Whether to normalize (passed to combine_multi_branch_cfg_noise).
            output_slice: If set, slice each output to [:, :output_slice].

        Returns:
            Combined noise prediction, identical on all ranks in CFG parallel.
        """
        if do_true_cfg:
            n_branches = len(branches_kwargs)
            cfg_world_size = _get_cfg_world_size_or_one()
            cfg_parallel_ready = cfg_world_size > 1

            if cfg_parallel_ready:
                return self._predict_multi_branch_parallel(
                    branches_kwargs,
                    n_branches,
                    cfg_world_size,
                    true_cfg_scale,
                    cfg_normalize,
                    output_slice,
                )
            else:
                # Sequential: run all N branches on single device
                preds: list[torch.Tensor | tuple[torch.Tensor, ...]] = []
                for kw in branches_kwargs:
                    pred = _wrap(self.predict_noise(**kw))
                    if output_slice is not None:
                        pred = _slice_pred(pred, output_slice)
                    preds.append(_unwrap(pred))
                return self.combine_multi_branch_cfg_noise(preds, true_cfg_scale, cfg_normalize)
        else:
            # No CFG: only compute positive/conditional prediction
            pred = self.predict_noise(**branches_kwargs[0])
            if output_slice is not None:
                pred = _unwrap(_slice_pred(_wrap(pred), output_slice))
            return pred

    def _predict_multi_branch_parallel(
        self,
        branches_kwargs: list[dict[str, Any]],
        n_branches: int,
        cfg_world_size: int,
        true_cfg_scale: float | dict[str, float],
        cfg_normalize: bool,
        output_slice: int | None,
    ) -> torch.Tensor | tuple[torch.Tensor, ...]:
        """Dispatch N branches across M ranks, all_gather, then combine."""
        cfg_group = get_cfg_group()
        cfg_rank = get_classifier_free_guidance_rank()

        if cfg_world_size > n_branches:
            logger.warning_once(
                "cfg_parallel_size=%d > n_branches=%d, %d GPU(s) will be idle for CFG",
                cfg_world_size,
                n_branches,
                cfg_world_size - n_branches,
            )

        # Assign branches to ranks via round-robin
        assignments = _dispatch_branches(n_branches, cfg_world_size)
        my_branch_ids = assignments[cfg_rank]
        max_per_rank = max(len(a) for a in assignments)

        # Run assigned branches
        my_preds: list[tuple[torch.Tensor, ...]] = []
        for bid in my_branch_ids:
            pred = _wrap(self.predict_noise(**branches_kwargs[bid]))
            if output_slice is not None:
                pred = _slice_pred(pred, output_slice)
            my_preds.append(pred)

        # Idle ranks (cfg_world_size > n_branches) run a forward pass to get the output shape for all_gather.
        # Output shape cannot be inferred from kwargs — may be tuple, sliced, etc.
        if not my_preds:
            pred = _wrap(self.predict_noise(**branches_kwargs[0]))
            if output_slice is not None:
                pred = _slice_pred(pred, output_slice)
            my_preds.append(pred)

        # Pad to max_per_rank with zeros so all ranks have same size
        ref_pred = my_preds[0]
        while len(my_preds) < max_per_rank:
            my_preds.append(tuple(torch.zeros_like(t) for t in ref_pred))

        # All-gather each output element separately (like predict_noise_maybe_with_cfg)
        # For each slot, gather across ranks; then pick valid results by owner_rank
        # all_slots[slot][elem_idx] = [rank0_tensor, rank1_tensor, ...]
        all_slots: list[list[list[torch.Tensor]]] = []
        for slot in range(max_per_rank):
            slot_results: list[list[torch.Tensor]] = []
            for p in my_preds[slot]:
                gathered = cfg_group.all_gather(p, separate_tensors=True)
                slot_results.append(gathered)
            all_slots.append(slot_results)

        # Reconstruct final_preds in branch order
        final_preds: list[torch.Tensor | tuple[torch.Tensor, ...]] = []
        for bid in range(n_branches):
            owner_rank = bid % cfg_world_size
            slot_idx = bid // cfg_world_size
            elements = tuple(all_slots[slot_idx][elem_idx][owner_rank] for elem_idx in range(len(ref_pred)))
            final_preds.append(_unwrap(elements))

        return self.combine_multi_branch_cfg_noise(final_preds, true_cfg_scale, cfg_normalize)

    def combine_multi_branch_cfg_noise(
        self,
        predictions: list[torch.Tensor | tuple[torch.Tensor, ...]],
        true_cfg_scale: float | dict[str, float],
        cfg_normalize: bool = False,
    ) -> torch.Tensor | tuple[torch.Tensor, ...]:
        """
        Combine N branch predictions. Default: standard 2-branch CFG formula.

        Override this method for custom multi-branch combine logic.

        Args:
            predictions: List of N predictions, where predictions[0] is always
                the positive/conditional branch.
            true_cfg_scale: CFG scale factor (float for 2-branch, dict for multi-branch).
            cfg_normalize: Whether to normalize the combined prediction.

        Returns:
            Combined noise prediction.
        """
        positive = _wrap(predictions[0])
        negative = _wrap(predictions[1])

        results = []
        for p, n in zip(positive, negative):
            comb = n + true_cfg_scale * (p - n)
            if cfg_normalize:
                comb = self.cfg_normalize_function(p, comb)
            results.append(comb)
        return _unwrap(tuple(results))

    def predict_noise(self, *args: Any, **kwargs: Any) -> torch.Tensor | tuple[torch.Tensor, ...] | IntermediateTensors:
        """
        Forward pass through transformer to predict noise.

        Subclasses should override this if they need custom behavior,
        but the default implementation calls self.transformer.

        Returns:
            Single Tensor for standard models, or tuple of Tensors for
            multi-output models (e.g., video + audio). Multi-output models
            must also override combine_cfg_noise() and set self.scheduler
            to a composite scheduler that handles tuples.
            Non-last Pipeline Parallel stages return `IntermediateTensors`
            instead of final noise tensors wrapped in a tuple.
        """
        result = self.transformer(*args, **kwargs)
        return result if isinstance(result, IntermediateTensors) else result[0]

    def diffuse(
        self,
        *args: Any,
        **kwargs: Any,
    ) -> Any:
        """
        Diffusion loop with optional classifier-free guidance.

        Subclasses MUST implement this method to define the complete
        diffusion/denoising loop for their specific model.

        Typical implementation pattern (single output):
        ```python
        def diffuse(self, latents, timesteps, prompt_embeds, negative_embeds, ...):
            for t in timesteps:
                positive_kwargs = {...}
                negative_kwargs = {...}

                noise_pred = self.predict_noise_maybe_with_cfg(
                    do_true_cfg=do_true_cfg,
                    true_cfg_scale=self.guidance_scale,
                    positive_kwargs=positive_kwargs,
                    negative_kwargs=negative_kwargs,
                )

                latents = self.scheduler_step_maybe_with_cfg(
                    noise_pred, t, latents, do_true_cfg=do_true_cfg
                )

            return latents
        ```

        Multi-output models (e.g., video + audio) should:
        1. Override ``predict_noise()`` to return a tuple
        2. Override ``combine_cfg_noise()`` for per-element CFG logic
        3. Set ``self.scheduler`` to a composite scheduler that handles tuples

        ```python
        def diffuse(self, video_latents, audio_latents, timesteps_video, timesteps_audio, ...):
            for t_v, t_a in zip(timesteps_video, timesteps_audio):
                positive_kwargs = {...}
                negative_kwargs = {...}

                # Returns tuple: (video_pred, audio_pred)
                video_pred, audio_pred = self.predict_noise_maybe_with_cfg(
                    do_true_cfg=do_true_cfg,
                    true_cfg_scale=self.guidance_scale,
                    positive_kwargs=positive_kwargs,
                    negative_kwargs=negative_kwargs,
                )

                # self.scheduler = VideoAudioScheduler(video_sched, audio_sched)
                # which accepts and returns tuples
                video_latents, audio_latents = self.scheduler_step_maybe_with_cfg(
                    (video_pred, audio_pred),
                    (t_v, t_a),
                    (video_latents, audio_latents),
                    do_true_cfg=do_true_cfg,
                )

            return video_latents, audio_latents
        ```
        """
        raise NotImplementedError("Subclasses must implement diffuse")

    def check_cfg_parallel_validity(self, true_cfg_scale: float, has_neg_prompt: bool) -> bool:
        """
        Check if CFG parallel configuration is valid.

        Args:
            true_cfg_scale: The classifier-free guidance scale value
            has_neg_prompt: Whether a negative prompt is provided

        Returns:
            True if CFG parallel configuration is valid, False otherwise
        """
        if _get_cfg_world_size_or_one() == 1:
            return True

        if true_cfg_scale <= 1:
            logger.warning("CFG parallel is NOT working correctly when true_cfg_scale <= 1.")
            return False

        if not has_neg_prompt:
            logger.warning(
                "CFG parallel is NOT working correctly when there is no negative prompt or negative prompt embeddings."
            )
            return False
        return True

    def scheduler_step(
        self,
        noise_pred: torch.Tensor | tuple[torch.Tensor, ...],
        t: torch.Tensor | tuple[torch.Tensor, ...],
        latents: torch.Tensor | tuple[torch.Tensor, ...],
        per_request_scheduler: Any | None = None,
        generator: torch.Generator | None = None,
    ) -> torch.Tensor | tuple[torch.Tensor, ...]:
        """
        Step the scheduler.

        Default implementation passes inputs directly to ``self.scheduler.step()``.
        For multi-output models, set ``self.scheduler`` to a composite scheduler
        that handles tuples (e.g., ``VideoAudioScheduler``).

        Args:
            noise_pred: Predicted noise (Tensor or tuple for multi-output)
            t: Current timestep (Tensor or tuple when schedulers differ per output)
            latents: Current latents (Tensor or tuple for multi-output)
            per_request_scheduler: Optional request-scoped scheduler that
                overrides ``self.scheduler`` for this call. This is
                primarily used by step-wise execution, where each request
                may keep scheduler state in its own runner-managed state
                object. Request-level execution should usually leave this
                as ``None`` and continue using ``self.scheduler``.
            generator: Optional torch Generator for reproducible sampling.
                When using CFG parallel, both ranks should receive generators
                initialized with the same seed so that non-deterministic
                schedulers (e.g., DDPM) produce identical results.

        Returns:
            Updated latents after scheduler step
        """
        sched = per_request_scheduler if per_request_scheduler is not None else getattr(self, "scheduler", None)
        if sched is None:
            raise ValueError("No scheduler is available. Set self.scheduler or pass per_request_scheduler.")
        if not callable(getattr(sched, "step", None)):
            raise TypeError("per_request_scheduler must provide a callable step(...) method.")
        step_kwargs = dict(return_dict=False)
        if generator is not None:
            step_kwargs["generator"] = generator
        return sched.step(noise_pred, t, latents, **step_kwargs)[0]

    def scheduler_step_maybe_with_cfg(
        self,
        noise_pred: torch.Tensor | tuple[torch.Tensor, ...],
        t: torch.Tensor | tuple[torch.Tensor, ...],
        latents: torch.Tensor | tuple[torch.Tensor, ...],
        do_true_cfg: bool,
        per_request_scheduler: Any | None = None,
        generator: torch.Generator | None = None,
    ) -> torch.Tensor | tuple[torch.Tensor, ...]:
        """
        Step the scheduler with automatic CFG parallel handling.

        All ranks compute the scheduler step locally — no broadcast needed
        because predict_noise_maybe_with_cfg already ensures all ranks have
        identical noise_pred after all_gather + local combine.

        Args:
            noise_pred: Predicted noise (Tensor or tuple, valid on all ranks)
            t: Current timestep (Tensor or tuple when schedulers differ per output)
            latents: Current latents (Tensor or tuple)
            do_true_cfg: Whether CFG is enabled
            per_request_scheduler: Optional request-scoped scheduler that
                overrides ``self.scheduler`` for this call. This is mainly
                needed by step-wise execution, where scheduler state may be
                stored per request. Request-level execution should normally
                leave this as ``None``.
            generator: Optional torch Generator for reproducible sampling.
                When using CFG parallel, both ranks should receive generators
                initialized with the same seed so that non-deterministic
                schedulers (e.g., DDPM) produce identical results.

        Returns:
            Updated latents (identical across all CFG ranks)
        """
        return self.scheduler_step(
            noise_pred,
            t,
            latents,
            per_request_scheduler=per_request_scheduler,
            generator=generator,
        )