# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo

# SPDX-License-Identifier: Apache-2.0
from dataclasses import dataclass, field
from typing import Tuple

from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig
from sglang.multimodal_gen.configs.models.fsdp import is_zimage_layer


@dataclass
class ZImageArchConfig(DiTArchConfig):
    all_patch_size: Tuple[int, ...] = (2,)
    all_f_patch_size: Tuple[int, ...] = (1,)
    in_channels: int = 16
    out_channels: int | None = None
    dim: int = 3840
    num_layers: int = 30
    n_refiner_layers: int = 2
    num_attention_heads: int = 30
    n_kv_heads: int = 30
    norm_eps: float = 1e-5
    qk_norm: bool = True
    cap_feat_dim: int = 2560
    rope_theta: float = 256.0
    t_scale: float = 1000.0
    axes_dims: Tuple[int, int, int] = (32, 48, 48)
    axes_lens: Tuple[int, int, int] = (1024, 512, 512)

    _fsdp_shard_conditions: list = field(default_factory=lambda: [is_zimage_layer])

    stacked_params_mapping: list[tuple[str, str, str]] = field(
        default_factory=lambda: [
            # (param_name, shard_name, shard_id)
            (".feed_forward.w13", ".feed_forward.w1", "gate"),
            (".feed_forward.w13", ".feed_forward.w3", "up"),
        ]
    )

    param_names_mapping: dict = field(
        default_factory=lambda: {
            r"(.*)\.attention\.to_q\.weight$": (r"\1.attention.to_qkv.weight", 0, 3),
            r"(.*)\.attention\.to_k\.weight$": (r"\1.attention.to_qkv.weight", 1, 3),
            r"(.*)\.attention\.to_v\.weight$": (r"\1.attention.to_qkv.weight", 2, 3),
            r"(.*)\.attention\.to_q\.weight_scale_inv$": (
                r"\1.attention.to_qkv.weight_scale_inv",
                0,
                3,
            ),
            r"(.*)\.attention\.to_k\.weight_scale_inv$": (
                r"\1.attention.to_qkv.weight_scale_inv",
                1,
                3,
            ),
            r"(.*)\.attention\.to_v\.weight_scale_inv$": (
                r"\1.attention.to_qkv.weight_scale_inv",
                2,
                3,
            ),
            r"(.*)\.attention\.to_q\.(lora_A|lora_B)$": (
                r"\1.attention.to_qkv.\2",
                0,
                3,
            ),
            r"(.*)\.attention\.to_k\.(lora_A|lora_B)$": (
                r"\1.attention.to_qkv.\2",
                1,
                3,
            ),
            r"(.*)\.attention\.to_v\.(lora_A|lora_B)$": (
                r"\1.attention.to_qkv.\2",
                2,
                3,
            ),
            r"(.*)\.feed_forward\.w1\.weight$": (r"\1.feed_forward.w13.weight", 0, 2),
            r"(.*)\.feed_forward\.w3\.weight$": (r"\1.feed_forward.w13.weight", 1, 2),
            r"(.*)\.feed_forward\.w1\.(lora_A|lora_B)$": (
                r"\1.feed_forward.w13.\2",
                0,
                2,
            ),
            r"(.*)\.feed_forward\.w3\.(lora_A|lora_B)$": (
                r"\1.feed_forward.w13.\2",
                1,
                2,
            ),
        }
    )

    def __post_init__(self):
        super().__post_init__()
        self.out_channels = self.out_channels or self.in_channels
        self.num_channels_latents = self.in_channels
        self.hidden_size = self.dim


@dataclass
class ZImageDitConfig(DiTConfig):
    arch_config: ZImageArchConfig = field(default_factory=ZImageArchConfig)

    prefix: str = "zimage"