# Uni3C controlnet for Wan 2.1: https://github.com/ewrfcas/Uni3C
# Converted from the original diffusers based implementation.
import torch
import torch.nn as nn

from comfy.ldm.flux.layers import EmbedND
from .model import WanSelfAttention


class Uni3CLayerNormZero(nn.Module):
    def __init__(
        self,
        conditioning_dim,
        embedding_dim,
        eps=1e-5,
        device=None, dtype=None, operations=None
    ):
        super().__init__()
        self.silu = nn.SiLU()
        self.linear = operations.Linear(conditioning_dim, 3 * embedding_dim, device=device, dtype=dtype)
        self.norm = operations.LayerNorm(embedding_dim, eps=eps, elementwise_affine=True, device=device, dtype=dtype)

    def forward(self, x, temb):
        shift, scale, gate = self.linear(self.silu(temb)).chunk(3, dim=1)
        x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :]
        return x, gate[:, None, :]


class Uni3CAttentionBlock(nn.Module):
    def __init__(
        self,
        dim,
        ffn_dim,
        num_heads,
        time_embed_dim=5120,
        eps=1e-6,
        device=None, dtype=None, operations=None
    ):
        super().__init__()
        operation_settings = {"operations": operations, "device": device, "dtype": dtype}
        self.norm1 = Uni3CLayerNormZero(time_embed_dim, dim, device=device, dtype=dtype, operations=operations)
        self.self_attn = WanSelfAttention(dim, num_heads, qk_norm=True, eps=eps, operation_settings=operation_settings)
        self.norm2 = Uni3CLayerNormZero(time_embed_dim, dim, device=device, dtype=dtype, operations=operations)
        self.ffn = nn.Sequential(
            operations.Linear(dim, ffn_dim, device=device, dtype=dtype), nn.GELU(approximate='tanh'),
            operations.Linear(ffn_dim, dim, device=device, dtype=dtype))

    def forward(self, x, temb, freqs):
        norm_x, gate_msa = self.norm1(x, temb)
        x = x + gate_msa * self.self_attn(norm_x, freqs)
        norm_x, gate_ff = self.norm2(x, temb)
        x = x + gate_ff * self.ffn(norm_x)
        return x


class MaskCamEmbed(nn.Module):
    def __init__(
        self,
        add_channels=7,
        mid_channels=256,
        conv_out_dim=5120,
        device=None, dtype=None, operations=None
    ):
        super().__init__()
        self.mask_padding = [0, 0, 0, 0, 3, 0]  # first frame conditioning
        self.mask_proj = nn.Sequential(
            operations.Conv3d(add_channels, mid_channels, kernel_size=(4, 8, 8), stride=(4, 8, 8), device=device, dtype=dtype),
            operations.GroupNorm(mid_channels // 8, mid_channels, device=device, dtype=dtype),
            nn.SiLU())
        self.mask_zero_proj = operations.Conv3d(mid_channels, conv_out_dim, kernel_size=(1, 2, 2), stride=(1, 2, 2), device=device, dtype=dtype)

    def forward(self, add_inputs):
        add_padded = torch.nn.functional.pad(add_inputs, self.mask_padding, mode="constant", value=0)
        add_embeds = self.mask_proj(add_padded)
        add_embeds = self.mask_zero_proj(add_embeds)
        add_embeds = add_embeds.flatten(2).transpose(1, 2)
        return add_embeds


class WanUni3CControlnet(nn.Module):
    def __init__(
        self,
        in_channels=36,
        conv_out_dim=5120,
        dim=1024,
        ffn_dim=8192,
        num_heads=16,
        num_layers=20,
        time_embed_dim=5120,
        out_proj_dim=5120,
        add_channels=7,
        mid_channels=256,
        device=None, dtype=None, operations=None
    ):
        super().__init__()
        patch_size = (1, 2, 2)
        self.num_layers = num_layers

        self.controlnet_patch_embedding = operations.Conv3d(
            in_channels, conv_out_dim, kernel_size=patch_size, stride=patch_size, device=device, dtype=torch.float32)
        self.controlnet_mask_embedding = MaskCamEmbed(add_channels, mid_channels, conv_out_dim, device=device, dtype=dtype, operations=operations)

        if conv_out_dim != dim:
            self.proj_in = operations.Linear(conv_out_dim, dim, device=device, dtype=dtype)
        else:
            self.proj_in = nn.Identity()

        self.controlnet_blocks = nn.ModuleList([
            Uni3CAttentionBlock(dim, ffn_dim, num_heads, time_embed_dim, device=device, dtype=dtype, operations=operations)
            for _ in range(num_layers)])
        self.proj_out = nn.ModuleList([
            operations.Linear(dim, out_proj_dim, device=device, dtype=dtype)
            for _ in range(num_layers)])

        head_dim = dim // num_heads
        self.rope_embedder = EmbedND(dim=head_dim, theta=10000.0, axes_dim=[head_dim - 4 * (head_dim // 6), 2 * (head_dim // 6), 2 * (head_dim // 6)])

    def rope_encode(self, t_len, h_len, w_len, device=None, dtype=None):
        img_ids = torch.zeros((t_len, h_len, w_len, 3), device=device, dtype=dtype)
        img_ids[:, :, :, 0] = img_ids[:, :, :, 0] + torch.arange(t_len, device=device, dtype=dtype).reshape(-1, 1, 1)
        img_ids[:, :, :, 1] = img_ids[:, :, :, 1] + torch.arange(h_len, device=device, dtype=dtype).reshape(1, -1, 1)
        img_ids[:, :, :, 2] = img_ids[:, :, :, 2] + torch.arange(w_len, device=device, dtype=dtype).reshape(1, 1, -1)
        img_ids = img_ids.reshape(1, -1, img_ids.shape[-1])
        freqs = self.rope_embedder(img_ids).movedim(1, 2)
        return freqs

    def process_input(self, control_input, render_mask=None, camera_embedding=None):
        # render_mask/camera_embedding are the checkpoint's extra conditioning path, not wired up yet
        hidden = self.controlnet_patch_embedding(control_input.float()).to(control_input.dtype)
        t_len, h_len, w_len = hidden.shape[2:]
        freqs = self.rope_encode(t_len, h_len, w_len, device=hidden.device, dtype=hidden.dtype)
        hidden = hidden.flatten(2).transpose(1, 2)

        add_inputs = None
        if camera_embedding is not None and render_mask is not None:
            add_inputs = torch.cat([render_mask, camera_embedding], dim=1)
        elif render_mask is not None:
            add_inputs = render_mask

        if add_inputs is not None:
            hidden = hidden + self.controlnet_mask_embedding(add_inputs.to(hidden.dtype))

        hidden = self.proj_in(hidden)
        return hidden, freqs

    def forward_block(self, block_index, hidden, temb, freqs):
        hidden = self.controlnet_blocks[block_index](hidden, temb, freqs)
        residual = self.proj_out[block_index](hidden)
        return hidden, residual
