"""Decoder template from the H3 VAE config (no H3 VAE weights needed).

The checkpoint holds every decoder parameter: the template is built on meta, then the
checkpoint is strictly assigned. QK normalization and rotary embedding follow the Diffusers
H3 source (MiniMaxH3VideoAttnProcessor).
"""
from types import SimpleNamespace

import torch
from torch import nn
from torch.nn import functional as F


def _apply_qk_norm(norm, x):
    return norm(x.float()).to(x.dtype)


def _apply_rotary_pos_emb_impl(x, emb):
    cos, sin = (t.to(x.dtype) for t in emb)
    dim = cos.shape[-1]
    a, tail = x[..., :dim], x[..., dim:]
    first, second = a.chunk(2, dim=-1)
    rotated = torch.cat((-second, first), dim=-1)
    return torch.cat((a * cos + rotated * sin, tail), dim=-1)


class SwiGLU(nn.Module):
    def __init__(self, dim, hidden):
        super().__init__()
        self.w1 = nn.Linear(dim, 2 * hidden)
        self.w2 = nn.Linear(hidden, dim)

    def forward(self, x):
        gate, up = self.w1(x).chunk(2, dim=-1)
        return self.w2(F.silu(gate) * up)


class RGBProcessor:
    def revert_tensor(self, x):
        # H3 returns ImageNet-normalized RGB, not pixels in [-1, 1].
        mean = x.new_tensor((0.485, 0.456, 0.406)).view(1, 3, 1, 1, 1)
        std = x.new_tensor((0.229, 0.224, 0.225)).view(1, 3, 1, 1, 1)
        return x * std + mean


def make_template(config):
    from diffusers.models.autoencoders.autoencoder_kl_minimax_h3 import MiniMaxH3VideoRotaryPosEmbed

    heads = config['decoder_num_attention_heads']
    dh = config['decoder_attention_head_dim']
    dim = heads * dh
    eps = config['decoder_norm_eps']
    # These nonpersistent buffers must be real CPU tensors, not meta tensors.
    rope = MiniMaxH3VideoRotaryPosEmbed(
        int(dh * config['decoder_rope_dim_ratio']), theta=config['decoder_rope_theta'])
    mean = torch.tensor(config['latents_mean'], dtype=torch.float32)
    std = torch.tensor(config['latents_std'], dtype=torch.float32)
    with torch.device('meta'):
        dec = nn.Module()
        dec.x_embedder = nn.Linear(24, dim)
        dec.register_tokens = nn.Parameter(torch.empty(1, config['decoder_num_register_tokens'], dim))
        dec.transformer_blocks = nn.ModuleList()
        for _ in range(config['decoder_num_layers']):
            block = nn.Module()
            block.norm1 = nn.RMSNorm(dim, eps=eps)
            block.norm2 = nn.RMSNorm(dim, eps=eps)
            block.scale1 = nn.Parameter(torch.empty(dim))
            block.scale2 = nn.Parameter(torch.empty(dim))
            block.attn = nn.Module()
            block.attn.heads, block.attn.dim_head = heads, dh
            block.attn.to_qkv = nn.Linear(dim, 3 * dim)
            block.attn.to_out = nn.Linear(dim, dim)
            block.attn.norm_q = nn.RMSNorm(dh, eps=eps, elementwise_affine=False)
            block.attn.norm_k = nn.RMSNorm(dh, eps=eps, elementwise_affine=False)
            block.ff = SwiGLU(dim, dim * config['decoder_ffn_mult'])
            dec.transformer_blocks.append(block)
        dec.norm_out = nn.LayerNorm(dim, eps=eps)
        dec.proj_out = nn.Linear(dim, 3 * 4 * 16 * 16)
        post_quant_conv = nn.Conv3d(24, 24, 1)
    dec.pos_embed = rope
    return SimpleNamespace(decoder=dec, post_quant_conv=post_quant_conv,
                           latents_mean=mean, latents_std=std, processor=RGBProcessor())