# Copyright 2025 StepFun Inc. All Rights Reserved.
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
# ==============================================================================

import torch
from einops import rearrange, repeat
from torch import nn

from fastvideo.v1.attention import DistributedAttention, LocalAttention
from fastvideo.v1.configs.models.dits import StepVideoConfig
from fastvideo.v1.distributed.parallel_state import (
    get_sequence_model_parallel_world_size)
from fastvideo.v1.layers.layernorm import LayerNormScaleShift
from fastvideo.v1.layers.linear import ReplicatedLinear
from fastvideo.v1.layers.mlp import MLP
from fastvideo.v1.layers.rotary_embedding import (_apply_rotary_emb,
                                                  get_rotary_pos_embed)
from fastvideo.v1.layers.visual_embedding import TimestepEmbedder
from fastvideo.v1.models.dits.base import BaseDiT
from fastvideo.v1.platforms import _Backend


class PatchEmbed2D(nn.Module):
    """2D Image to Patch Embedding

    Image to Patch Embedding using Conv2d

    A convolution based approach to patchifying a 2D image w/ embedding projection.

    Based on the impl in https://github.com/google-research/vision_transformer

    Hacked together by / Copyright 2020 Ross Wightman

    Remove the _assert function in forward function to be compatible with multi-resolution images.
    """

    def __init__(self,
                 patch_size=16,
                 in_chans=3,
                 embed_dim=768,
                 norm_layer=None,
                 flatten=True,
                 bias=True,
                 dtype=None,
                 prefix: str = ""):
        super().__init__()
        # Convert patch_size to 2-tuple
        if isinstance(patch_size, list | tuple):
            if len(patch_size) == 1:
                patch_size = (patch_size[0], patch_size[0])
        else:
            patch_size = (patch_size, patch_size)

        self.patch_size = patch_size
        self.flatten = flatten

        self.proj = nn.Conv2d(in_chans,
                              embed_dim,
                              kernel_size=patch_size,
                              stride=patch_size,
                              bias=bias,
                              dtype=dtype)
        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()

    def forward(self, x):
        x = self.proj(x)
        if self.flatten:
            x = x.flatten(2).transpose(1, 2)  # BCHW -> BNC
        x = self.norm(x)
        return x


class StepVideoRMSNorm(nn.Module):

    def __init__(
        self,
        dim: int,
        elementwise_affine=True,
        eps: float = 1e-6,
        device=None,
        dtype=None,
    ):
        """
        Initialize the RMSNorm normalization layer.

        Args:
            dim (int): The dimension of the input tensor.
            eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6.

        Attributes:
            eps (float): A small value added to the denominator for numerical stability.
            weight (nn.Parameter): Learnable scaling parameter.

        """
        factory_kwargs = {"device": device, "dtype": dtype}
        super().__init__()
        self.eps = eps
        if elementwise_affine:
            self.weight = nn.Parameter(torch.ones(dim, **factory_kwargs))

    def _norm(self, x) -> torch.Tensor:
        """
        Apply the RMSNorm normalization to the input tensor.

        Args:
            x (torch.Tensor): The input tensor.

        Returns:
            torch.Tensor: The normalized tensor.

        """
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)

    def forward(self, x):
        """
        Forward pass through the RMSNorm layer.

        Args:
            x (torch.Tensor): The input tensor.

        Returns:
            torch.Tensor: The output tensor after applying RMSNorm.

        """
        output = self._norm(x.float()).type_as(x)
        if hasattr(self, "weight"):
            output = output * self.weight
        return output


class SelfAttention(nn.Module):

    def __init__(self,
                 hidden_dim,
                 head_dim,
                 rope_split: tuple[int, int, int] = (64, 32, 32),
                 bias: bool = False,
                 with_rope: bool = True,
                 with_qk_norm: bool = True,
                 attn_type: str = "torch",
                 supported_attention_backends=(_Backend.FLASH_ATTN,
                                               _Backend.TORCH_SDPA)):
        super().__init__()
        self.head_dim = head_dim
        self.hidden_dim = hidden_dim
        self.rope_split = list(rope_split)
        self.n_heads = hidden_dim // head_dim

        self.wqkv = ReplicatedLinear(hidden_dim, hidden_dim * 3, bias=bias)
        self.wo = ReplicatedLinear(hidden_dim, hidden_dim, bias=bias)

        self.with_rope = with_rope
        self.with_qk_norm = with_qk_norm
        if self.with_qk_norm:
            self.q_norm = StepVideoRMSNorm(head_dim, elementwise_affine=True)
            self.k_norm = StepVideoRMSNorm(head_dim, elementwise_affine=True)

        # self.core_attention = self.attn_processor(attn_type=attn_type)
        self.parallel = attn_type == 'parallel'
        self.attn = DistributedAttention(
            num_heads=self.n_heads,
            head_size=head_dim,
            causal=False,
            supported_attention_backends=supported_attention_backends)

    def _apply_rope(self, x: torch.Tensor, cos: torch.Tensor,
                    sin: torch.Tensor):
        """
        x:   [B, S, H, D]
        cos: [S, D/2]  where D = head_dim = sum(self.rope_split)
        sin: [S, D/2]
        returns x with rotary applied exactly as v0 did
        """
        B, S, H, D = x.shape
        # 1) split cos/sin per chunk
        half_splits = [c // 2
                       for c in self.rope_split]  # [32,16,16] for [64,32,32]
        cos_splits = cos.split(half_splits, dim=1)
        sin_splits = sin.split(half_splits, dim=1)

        outs = []
        idx = 0
        for (chunk_size, cos_i, sin_i) in zip(self.rope_split,
                                              cos_splits,
                                              sin_splits,
                                              strict=False):
            # slice the corresponding channels
            x_chunk = x[..., idx:idx + chunk_size]  # [B,S,H,chunk_size]
            idx += chunk_size

            # flatten to [S, B*H, chunk_size]
            x_flat = rearrange(x_chunk, 'b s h d -> s (b h) d')

            # apply rotary on *that* chunk
            out_flat = _apply_rotary_emb(x_flat,
                                         cos_i,
                                         sin_i,
                                         is_neox_style=True)

            # restore [B,S,H,chunk_size]
            out = rearrange(out_flat, 's (b h) d -> b s h d', b=B, h=H)
            outs.append(out)

        # concatenate back to [B,S,H,D]
        return torch.cat(outs, dim=-1)

    def forward(self,
                x,
                cu_seqlens=None,
                max_seqlen=None,
                rope_positions=None,
                cos_sin=None,
                attn_mask=None,
                mask_strategy=None):

        B, S, _ = x.shape
        xqkv, _ = self.wqkv(x)
        xqkv = xqkv.view(*x.shape[:-1], self.n_heads, 3 * self.head_dim)
        q, k, v = torch.split(xqkv, [self.head_dim] * 3, dim=-1)  # [B,S,H,D]

        if self.with_qk_norm:
            q = self.q_norm(q)
            k = self.k_norm(k)

        if self.with_rope:
            if rope_positions is not None:
                F, Ht, W = rope_positions
            assert F * Ht * W == S, "rope_positions mismatches sequence length"

            cos, sin = cos_sin
            cos = cos.to(x.device, dtype=x.dtype)
            sin = sin.to(x.device, dtype=x.dtype)

            q = self._apply_rope(q, cos, sin)
            k = self._apply_rope(k, cos, sin)

        output, _ = self.attn(q, k, v)  # [B,heads,S,D]

        output = rearrange(output, 'b s h d -> b s (h d)')
        output, _ = self.wo(output)

        return output


class CrossAttention(nn.Module):

    def __init__(
        self,
        hidden_dim,
        head_dim,
        bias=False,
        with_qk_norm=True,
        supported_attention_backends=(_Backend.FLASH_ATTN, _Backend.TORCH_SDPA)
    ) -> None:
        super().__init__()
        self.head_dim = head_dim
        self.n_heads = hidden_dim // head_dim

        self.wq = ReplicatedLinear(hidden_dim, hidden_dim, bias=bias)
        self.wkv = ReplicatedLinear(hidden_dim, hidden_dim * 2, bias=bias)
        self.wo = ReplicatedLinear(hidden_dim, hidden_dim, bias=bias)

        self.with_qk_norm = with_qk_norm
        if self.with_qk_norm:
            self.q_norm = StepVideoRMSNorm(head_dim, elementwise_affine=True)
            self.k_norm = StepVideoRMSNorm(head_dim, elementwise_affine=True)

        self.attn = LocalAttention(
            num_heads=self.n_heads,
            head_size=head_dim,
            causal=False,
            supported_attention_backends=supported_attention_backends)

    def forward(self,
                x: torch.Tensor,
                encoder_hidden_states: torch.Tensor,
                attn_mask=None) -> torch.Tensor:

        xq, _ = self.wq(x)
        xq = xq.view(*xq.shape[:-1], self.n_heads, self.head_dim)

        xkv, _ = self.wkv(encoder_hidden_states)
        xkv = xkv.view(*xkv.shape[:-1], self.n_heads, 2 * self.head_dim)

        xk, xv = torch.split(xkv, [self.head_dim] * 2,
                             dim=-1)  ## seq_len, n, dim

        if self.with_qk_norm:
            xq = self.q_norm(xq)
            xk = self.k_norm(xk)

        output = self.attn(xq, xk, xv)

        output = rearrange(output, 'b s h d -> b s (h d)')
        output, _ = self.wo(output)

        return output


class AdaLayerNormSingle(nn.Module):
    r"""
        Norm layer adaptive layer norm single (adaLN-single).

        As proposed in PixArt-Alpha (see: https://arxiv.org/abs/2310.00426; Section 2.3).

        Parameters:
            embedding_dim (`int`): The size of each embedding vector.
            use_additional_conditions (`bool`): To use additional conditions for normalization or not.
    """

    def __init__(self, embedding_dim: int, time_step_rescale=1000):
        super().__init__()

        self.emb = TimestepEmbedder(embedding_dim)

        self.silu = nn.SiLU()
        self.linear = ReplicatedLinear(embedding_dim,
                                       6 * embedding_dim,
                                       bias=True)

        self.time_step_rescale = time_step_rescale  ## timestep usually in [0, 1], we rescale it to [0,1000] for stability

    def forward(
        self,
        timestep: torch.Tensor,
        added_cond_kwargs: dict[str, torch.Tensor] | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        embedded_timestep = self.emb(timestep * self.time_step_rescale)

        out, _ = self.linear(self.silu(embedded_timestep))

        return out, embedded_timestep


class StepVideoTransformerBlock(nn.Module):
    r"""
    A basic Transformer block.

    Parameters:
        dim (`int`): The number of channels in the input and output.
        num_attention_heads (`int`): The number of heads to use for multi-head attention.
        attention_head_dim (`int`): The number of channels in each head.
        dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
        cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention.
        activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward.
        num_embeds_ada_norm (:
            obj: `int`, *optional*): The number of diffusion steps used during training. See `Transformer2DModel`.
        attention_bias (:
            obj: `bool`, *optional*, defaults to `False`): Configure if the attentions should contain a bias parameter.
        only_cross_attention (`bool`, *optional*):
            Whether to use only cross-attention layers. In this case two cross attention layers are used.
        double_self_attention (`bool`, *optional*):
            Whether to use two self-attention layers. In this case no cross attention layers are used.
        upcast_attention (`bool`, *optional*):
            Whether to upcast the attention computation to float32. This is useful for mixed precision training.
        norm_elementwise_affine (`bool`, *optional*, defaults to `True`):
            Whether to use learnable elementwise affine parameters for normalization.
        norm_type (`str`, *optional*, defaults to `"layer_norm"`):
            The normalization layer to use. Can be `"layer_norm"`, `"ada_norm"` or `"ada_norm_zero"`.
        final_dropout (`bool` *optional*, defaults to False):
            Whether to apply a final dropout after the last feed-forward layer.
        positional_embeddings (`str`, *optional*, defaults to `None`):
            The type of positional embeddings to apply to.
        num_positional_embeddings (`int`, *optional*, defaults to `None`):
            The maximum number of positional embeddings to apply.
    """

    def __init__(self,
                 dim: int,
                 attention_head_dim: int,
                 norm_eps: float = 1e-5,
                 ff_inner_dim: int | None = None,
                 ff_bias: bool = False,
                 attention_type: str = 'torch'):
        super().__init__()
        self.dim = dim
        self.norm1 = LayerNormScaleShift(dim,
                                         norm_type="layer",
                                         elementwise_affine=True,
                                         eps=norm_eps)
        self.attn1 = SelfAttention(
            dim,
            attention_head_dim,
            bias=False,
            with_rope=True,
            with_qk_norm=True,
        )

        self.norm2 = LayerNormScaleShift(dim,
                                         norm_type="layer",
                                         elementwise_affine=True,
                                         eps=norm_eps)
        self.attn2 = CrossAttention(dim,
                                    attention_head_dim,
                                    bias=False,
                                    with_qk_norm=True)

        self.ff = MLP(input_dim=dim,
                      mlp_hidden_dim=dim *
                      4 if ff_inner_dim is None else ff_inner_dim,
                      act_type="gelu_pytorch_tanh",
                      bias=ff_bias)

        self.scale_shift_table = nn.Parameter(torch.randn(6, dim) / dim**0.5)

    @torch.no_grad()
    def forward(self,
                q: torch.Tensor,
                kv: torch.Tensor,
                t_expand: torch.LongTensor,
                attn_mask=None,
                rope_positions: list | None = None,
                cos_sin=None,
                mask_strategy=None) -> torch.Tensor:

        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
            torch.clone(chunk)
            for chunk in (self.scale_shift_table[None] +
                          t_expand.reshape(-1, 6, self.dim)).chunk(6, dim=1))

        scale_shift_q = self.norm1(q,
                                   scale=scale_msa.squeeze(1),
                                   shift=shift_msa.squeeze(1))

        attn_q = self.attn1(scale_shift_q,
                            rope_positions=rope_positions,
                            cos_sin=cos_sin,
                            mask_strategy=mask_strategy)

        q = attn_q * gate_msa + q

        attn_q = self.attn2(q, kv, attn_mask)

        q = attn_q + q

        scale_shift_q = self.norm2(q,
                                   scale=scale_mlp.squeeze(1),
                                   shift=shift_mlp.squeeze(1))

        ff_output = self.ff(scale_shift_q)

        q = ff_output * gate_mlp + q

        return q


class StepVideoModel(BaseDiT):
    # (Optional) Keep the same attribute for compatibility with splitting, etc.
    _fsdp_shard_conditions = [
        lambda n, m: "transformer_blocks" in n and n.split(".")[-1].isdigit(),
        # lambda n, m: "pos_embed" in n  # If needed for the patch embedding.
    ]
    _param_names_mapping = StepVideoConfig()._param_names_mapping
    _supported_attention_backends = StepVideoConfig(
    )._supported_attention_backends

    def __init__(self, config: StepVideoConfig) -> None:
        super().__init__(config=config)
        self.num_attention_heads = config.num_attention_heads
        self.attention_head_dim = config.attention_head_dim
        self.in_channels = config.in_channels
        self.out_channels = config.out_channels
        self.num_layers = config.num_layers
        self.dropout = config.dropout
        self.patch_size = config.patch_size
        self.norm_type = config.norm_type
        self.norm_elementwise_affine = config.norm_elementwise_affine
        self.norm_eps = config.norm_eps
        self.use_additional_conditions = config.use_additional_conditions
        self.caption_channels = config.caption_channels
        self.attention_type = config.attention_type
        self.num_channels_latents = config.num_channels_latents
        # Compute inner dimension.
        self.hidden_size = config.hidden_size

        # Image/video patch embedding.
        self.pos_embed = PatchEmbed2D(
            patch_size=self.patch_size,
            in_chans=self.in_channels,
            embed_dim=self.hidden_size,
        )

        self._rope_cache: dict[tuple, tuple[torch.Tensor, torch.Tensor]] = {}
        # Transformer blocks.
        self.transformer_blocks = nn.ModuleList([
            StepVideoTransformerBlock(
                dim=self.hidden_size,
                attention_head_dim=self.attention_head_dim,
                attention_type=self.attention_type)
            for _ in range(self.num_layers)
        ])

        # Output blocks.
        self.norm_out = LayerNormScaleShift(
            self.hidden_size,
            norm_type="layer",
            eps=self.norm_eps,
            elementwise_affine=self.norm_elementwise_affine)
        self.scale_shift_table = nn.Parameter(
            torch.randn(2, self.hidden_size) / (self.hidden_size**0.5))
        self.proj_out = ReplicatedLinear(
            self.hidden_size,
            self.patch_size * self.patch_size * self.out_channels)
        # Time modulation via adaptive layer norm.
        self.adaln_single = AdaLayerNormSingle(self.hidden_size)

        # Set up caption conditioning.
        if isinstance(self.caption_channels, int):
            caption_channel = self.caption_channels
        else:
            caption_channel, clip_channel = self.caption_channels
            self.clip_projection = ReplicatedLinear(clip_channel,
                                                    self.hidden_size)
        self.caption_norm = nn.LayerNorm(
            caption_channel,
            eps=self.norm_eps,
            elementwise_affine=self.norm_elementwise_affine)
        self.caption_projection = MLP(input_dim=caption_channel,
                                      mlp_hidden_dim=self.hidden_size,
                                      act_type="gelu_pytorch_tanh")

        # Flag to indicate if using parallel attention.
        self.parallel = (self.attention_type == "parallel")

        self.__post_init__()

    def patchfy(self, hidden_states) -> torch.Tensor:
        hidden_states = rearrange(hidden_states, 'b f c h w -> (b f) c h w')
        hidden_states = self.pos_embed(hidden_states)
        return hidden_states

    def prepare_attn_mask(self, encoder_attention_mask, encoder_hidden_states,
                          q_seqlen) -> tuple[torch.Tensor, torch.Tensor]:
        kv_seqlens = encoder_attention_mask.sum(dim=1).int()
        mask = torch.zeros([len(kv_seqlens), q_seqlen,
                            max(kv_seqlens)],
                           dtype=torch.bool,
                           device=encoder_attention_mask.device)
        encoder_hidden_states = encoder_hidden_states[:, :max(kv_seqlens)]
        for i, kv_len in enumerate(kv_seqlens):
            mask[i, :, :kv_len] = 1
        return encoder_hidden_states, mask

    def block_forward(self,
                      hidden_states,
                      encoder_hidden_states=None,
                      t_expand=None,
                      rope_positions=None,
                      cos_sin=None,
                      attn_mask=None,
                      parallel=True,
                      mask_strategy=None) -> torch.Tensor:

        for i, block in enumerate(self.transformer_blocks):
            hidden_states = block(hidden_states,
                                  encoder_hidden_states,
                                  t_expand=t_expand,
                                  attn_mask=attn_mask,
                                  rope_positions=rope_positions,
                                  cos_sin=cos_sin,
                                  mask_strategy=mask_strategy[i])

        return hidden_states

    def _get_rope(self, rope_positions: tuple[int, int, int],
                  dtype: torch.dtype, device: torch.device):
        F, Ht, W = rope_positions
        key = (F, Ht, W, dtype)
        if key not in self._rope_cache:
            cos, sin = get_rotary_pos_embed(
                rope_sizes=(F * get_sequence_model_parallel_world_size(), Ht,
                            W),
                hidden_size=self.hidden_size,
                heads_num=self.hidden_size // self.attention_head_dim,
                rope_dim_list=(64, 32, 32),  # same split you used
                rope_theta=1.0e4,
                dtype=torch.float32  # build once in fp32
            )
            # move & cast once
            self._rope_cache[key] = (cos.to(device, dtype=dtype),
                                     sin.to(device, dtype=dtype))
        return self._rope_cache[key]

    @torch.inference_mode()
    def forward(
        self,
        hidden_states: torch.Tensor,
        encoder_hidden_states: torch.Tensor | None = None,
        t_expand: torch.LongTensor | None = None,
        encoder_hidden_states_2: torch.Tensor | None = None,
        added_cond_kwargs: dict[str, torch.Tensor] | None = None,
        encoder_attention_mask: torch.Tensor | None = None,
        fps: torch.Tensor | None = None,
        return_dict: bool = True,
        mask_strategy=None,
        guidance=None,
    ):
        assert hidden_states.ndim == 5
        "hidden_states's shape should be (bsz, f, ch, h ,w)"
        frame = hidden_states.shape[2]
        hidden_states = rearrange(hidden_states,
                                  'b c f h w -> b f c h w',
                                  f=frame)
        if mask_strategy is None:
            mask_strategy = [None, None]
        bsz, frame, _, height, width = hidden_states.shape
        height, width = height // self.patch_size, width // self.patch_size

        hidden_states = self.patchfy(hidden_states)
        len_frame = hidden_states.shape[1]

        t_expand, embedded_timestep = self.adaln_single(t_expand)
        encoder_hidden_states = self.caption_projection(
            self.caption_norm(encoder_hidden_states))

        if encoder_hidden_states_2 is not None and hasattr(
                self, 'clip_projection'):
            clip_embedding, _ = self.clip_projection(encoder_hidden_states_2)
            encoder_hidden_states = torch.cat(
                [clip_embedding, encoder_hidden_states], dim=1)

        hidden_states = rearrange(hidden_states,
                                  '(b f) l d->  b (f l) d',
                                  b=bsz,
                                  f=frame,
                                  l=len_frame).contiguous()
        encoder_hidden_states, attn_mask = self.prepare_attn_mask(
            encoder_attention_mask,
            encoder_hidden_states,
            q_seqlen=frame * len_frame)

        cos_sin = self._get_rope((frame, height, width), hidden_states.dtype,
                                 hidden_states.device)

        hidden_states = self.block_forward(
            hidden_states,
            encoder_hidden_states,
            t_expand=t_expand,
            rope_positions=[frame, height, width],
            cos_sin=cos_sin,
            attn_mask=attn_mask,
            parallel=self.parallel,
            mask_strategy=mask_strategy)

        hidden_states = rearrange(hidden_states,
                                  'b (f l) d -> (b f) l d',
                                  b=bsz,
                                  f=frame,
                                  l=len_frame)

        embedded_timestep = repeat(embedded_timestep, 'b d -> (b f) d',
                                   f=frame).contiguous()

        shift, scale = (self.scale_shift_table[None] +
                        embedded_timestep[:, None]).chunk(2, dim=1)
        hidden_states = self.norm_out(hidden_states,
                                      shift=shift.squeeze(1),
                                      scale=scale.squeeze(1))
        # Modulation
        hidden_states, _ = self.proj_out(hidden_states)

        # unpatchify
        hidden_states = hidden_states.reshape(shape=(-1, height, width,
                                                     self.patch_size,
                                                     self.patch_size,
                                                     self.out_channels))

        hidden_states = rearrange(hidden_states, 'n h w p q c -> n c h p w q')
        output = hidden_states.reshape(shape=(-1, self.out_channels,
                                              height * self.patch_size,
                                              width * self.patch_size))

        output = rearrange(output, '(b f) c h w -> b c f h w', f=frame)
        return output