# Trimmed port of GenFirst's Prior_only/models/mmdit.py (from Stability-AI/sd3-ref, MIT).
# Only the configuration used by the released MMDiT-XL/2 text-to-image prior is kept:
# LayerNorm, GELU MLP, qkv bias, no qk-norm, no registers, CLIP-L (768-d) context.
# The REPA projector heads are training-only and are dropped.

import math

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F


def attention(q, k, v, heads):
    b, _, dim = q.shape
    dh = dim // heads
    q, k, v = (t.view(b, -1, heads, dh).transpose(1, 2) for t in (q, k, v))
    out = F.scaled_dot_product_attention(q, k, v)
    return out.transpose(1, 2).reshape(b, -1, heads * dh)


def modulate(x, shift, scale):
    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)


def get_2d_sincos_pos_embed(embed_dim, grid_size):
    grid_h = np.arange(grid_size, dtype=np.float32)
    grid_w = np.arange(grid_size, dtype=np.float32)
    grid = np.stack(np.meshgrid(grid_w, grid_h), axis=0).reshape([2, 1, grid_size, grid_size])

    def _1d(d, pos):
        omega = np.arange(d // 2, dtype=np.float64) / (d / 2.0)
        omega = 1.0 / 10000**omega
        out = np.einsum("m,d->md", pos.reshape(-1), omega)
        return np.concatenate([np.sin(out), np.cos(out)], axis=1)

    return np.concatenate([_1d(embed_dim // 2, grid[0]), _1d(embed_dim // 2, grid[1])], axis=1)


class PatchEmbed(nn.Module):
    def __init__(self, img_size, patch_size, in_chans, embed_dim):
        super().__init__()
        self.patch_size = (patch_size, patch_size)
        self.num_patches = (img_size // patch_size) ** 2
        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size, bias=True)

    def forward(self, x):
        return self.proj(x).flatten(2).transpose(1, 2)


class TimestepEmbedder(nn.Module):
    def __init__(self, hidden_size, frequency_embedding_size=256):
        super().__init__()
        self.mlp = nn.Sequential(
            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
            nn.SiLU(),
            nn.Linear(hidden_size, hidden_size, bias=True),
        )
        self.frequency_embedding_size = frequency_embedding_size

    @staticmethod
    def timestep_embedding(t, dim, max_period=10000):
        half = dim // 2
        freqs = torch.exp(
            -math.log(max_period) * torch.arange(0, half, dtype=torch.float32) / half
        ).to(device=t.device)
        args = t[:, None].float() * freqs[None]
        emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
        if torch.is_floating_point(t):
            emb = emb.to(dtype=t.dtype)
        return emb

    def forward(self, t, dtype):
        return self.mlp(self.timestep_embedding(t, self.frequency_embedding_size).to(dtype))


class Mlp(nn.Module):
    def __init__(self, dim, hidden):
        super().__init__()
        self.fc1 = nn.Linear(dim, hidden)
        self.act = nn.GELU(approximate="tanh")
        self.fc2 = nn.Linear(hidden, dim)

    def forward(self, x):
        return self.fc2(self.act(self.fc1(x)))


class SelfAttention(nn.Module):
    def __init__(self, dim, num_heads, pre_only=False):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.qkv = nn.Linear(dim, dim * 3, bias=True)
        if not pre_only:
            self.proj = nn.Linear(dim, dim)

    def pre_attention(self, x):
        qkv = self.qkv(x)
        qkv = qkv.reshape(qkv.shape[0], qkv.shape[1], 3, -1, self.head_dim).movedim(2, 0)
        q, k, v = qkv[0], qkv[1], qkv[2]
        B, L = q.shape[0], q.shape[1]
        return q.reshape(B, L, -1), k.reshape(B, L, -1), v


class DismantledBlock(nn.Module):
    def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, pre_only=False):
        super().__init__()
        self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
        self.attn = SelfAttention(hidden_size, num_heads, pre_only=pre_only)
        if not pre_only:
            self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
            self.mlp = Mlp(hidden_size, int(hidden_size * mlp_ratio))
        n_mods = 2 if pre_only else 6
        self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, n_mods * hidden_size, bias=True))
        self.pre_only = pre_only

    def pre_attention(self, x, c):
        if not self.pre_only:
            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(c).chunk(6, dim=1)
            qkv = self.attn.pre_attention(modulate(self.norm1(x), shift_msa, scale_msa))
            return qkv, (x, gate_msa, shift_mlp, scale_mlp, gate_mlp)
        shift_msa, scale_msa = self.adaLN_modulation(c).chunk(2, dim=1)
        return self.attn.pre_attention(modulate(self.norm1(x), shift_msa, scale_msa)), None

    def post_attention(self, attn, x, gate_msa, shift_mlp, scale_mlp, gate_mlp):
        x = x + gate_msa.unsqueeze(1) * self.attn.proj(attn)
        x = x + gate_mlp.unsqueeze(1) * self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp))
        return x


class JointBlock(nn.Module):
    def __init__(self, hidden_size, num_heads, pre_only):
        super().__init__()
        self.context_block = DismantledBlock(hidden_size, num_heads, pre_only=pre_only)
        self.x_block = DismantledBlock(hidden_size, num_heads, pre_only=False)

    def forward(self, context, x, c):
        cqkv, cint = self.context_block.pre_attention(context, c)
        xqkv, xint = self.x_block.pre_attention(x, c)
        q, k, v = (torch.cat((cqkv[t], xqkv[t]), dim=1) for t in range(3))
        attn = attention(q, k, v, self.x_block.attn.num_heads)
        n_ctx = cqkv[0].shape[1]
        context_attn, x_attn = attn[:, :n_ctx], attn[:, n_ctx:]
        context = None if self.context_block.pre_only else self.context_block.post_attention(context_attn, *cint)
        x = self.x_block.post_attention(x_attn, *xint)
        return context, x


class FinalLayer(nn.Module):
    def __init__(self, hidden_size, patch_size, out_channels):
        super().__init__()
        self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
        self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True)
        self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True))

    def forward(self, x, c):
        shift, scale = self.adaLN_modulation(c).chunk(2, dim=1)
        return self.linear(modulate(self.norm_final(x), shift, scale))


class MMDiT(nn.Module):
    def __init__(self, input_size=64, patch_size=2, in_channels=16, depth=28,
                 hidden_size=1152, num_heads=16, context_dim=768):
        super().__init__()
        self.out_channels = in_channels
        self.x_embedder = PatchEmbed(input_size, patch_size, in_channels, hidden_size)
        self.t_embedder = TimestepEmbedder(hidden_size)
        self.context_embedder = nn.Linear(context_dim, hidden_size)
        self.register_buffer("pos_embed", torch.zeros(1, self.x_embedder.num_patches, hidden_size))
        self.joint_blocks = nn.ModuleList(
            [JointBlock(hidden_size, num_heads, pre_only=(i == depth - 1)) for i in range(depth)]
        )
        self.final_layer = FinalLayer(hidden_size, patch_size, in_channels)
        pe = get_2d_sincos_pos_embed(hidden_size, int(self.x_embedder.num_patches ** 0.5))
        self.pos_embed.data.copy_(torch.from_numpy(pe).float().unsqueeze(0))

    def unpatchify(self, x, h, w):
        c, p = self.out_channels, self.x_embedder.patch_size[0]
        x = x.reshape(x.shape[0], h, w, p, p, c)
        x = torch.einsum("nhwpqc->nchpwq", x)
        return x.reshape(x.shape[0], c, h * p, w * p)

    def forward(self, x, t, context):
        H, W = x.shape[-2:]
        p = self.x_embedder.patch_size[0]
        x = self.x_embedder(x) + self.pos_embed
        c = self.t_embedder(t, dtype=x.dtype)
        context = self.context_embedder(context)
        for block in self.joint_blocks:
            context, x = block(context, x, c=c)
        x = self.final_layer(x, c)
        return self.unpatchify(x, H // p, W // p)