from __future__ import annotations

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.checkpoint

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

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

def rope_freqs(positions, dim, base=10000.0):
    inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
    return torch.outer(positions.float(), inv_freq)

def rope_cos_sin(freqs):
    emb = torch.cat([freqs, freqs], dim=-1)
    return emb.cos(), emb.sin()

def rotate_half(x):
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat([-x2, x1], dim=-1)

def apply_rope(x, cos, sin):
    return x * cos + rotate_half(x) * sin

def apply_rope_2d(x, row_cos, row_sin, col_cos, col_sin):
    x1, x2 = x.chunk(2, dim=-1)
    x1 = apply_rope(x1, row_cos, row_sin)
    x2 = apply_rope(x2, col_cos, col_sin)
    return torch.cat([x1, x2], dim=-1)

class RMSNormHead(nn.Module):
    def __init__(self, head_dim, eps=1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(head_dim))
        self.eps = eps

    def forward(self, x):
        n = x.pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()
        return x * n * self.weight

class SwiGLU(nn.Module):
    def __init__(self, dim, hidden):
        super().__init__()
        self.gate = nn.Linear(dim, hidden)
        self.up = nn.Linear(dim, hidden)
        self.down = nn.Linear(hidden, dim)

    def forward(self, x):
        return self.down(F.silu(self.gate(x)) * self.up(x))

class JointBlock(nn.Module):
    def __init__(self, dim, heads, mlp_hidden):
        super().__init__()
        self.heads = heads
        self.head_dim = dim // heads
        self.norm1_img = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
        self.norm1_txt = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
        self.qkv_img = nn.Linear(dim, 3 * dim)
        self.qkv_txt = nn.Linear(dim, 3 * dim)
        self.qn_img = RMSNormHead(self.head_dim)
        self.kn_img = RMSNormHead(self.head_dim)
        self.qn_txt = RMSNormHead(self.head_dim)
        self.kn_txt = RMSNormHead(self.head_dim)
        self.proj_img = nn.Linear(dim, dim)
        self.proj_txt = nn.Linear(dim, dim)
        self.norm2_img = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
        self.norm2_txt = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
        self.mlp_img = SwiGLU(dim, mlp_hidden)
        self.mlp_txt = SwiGLU(dim, mlp_hidden)
        self.ada_img = nn.Sequential(nn.SiLU(), nn.Linear(dim, 6 * dim))
        self.ada_txt = nn.Sequential(nn.SiLU(), nn.Linear(dim, 6 * dim))

    def forward(self, img, txt, c, rope_img, rope_txt, key_valid):
        s1i, sc1i, g1i, s2i, sc2i, g2i = self.ada_img(c).chunk(6, dim=-1)
        s1t, sc1t, g1t, s2t, sc2t, g2t = self.ada_txt(c).chunk(6, dim=-1)

        xi = modulate(self.norm1_img(img), s1i, sc1i)
        xt = modulate(self.norm1_txt(txt), s1t, sc1t)

        B, Ni, C = xi.shape
        Nt = xt.shape[1]
        H, D = self.heads, self.head_dim

        qi, ki, vi = self.qkv_img(xi).reshape(B, Ni, 3, H, D).permute(2, 0, 3, 1, 4)
        qt, kt, vt = self.qkv_txt(xt).reshape(B, Nt, 3, H, D).permute(2, 0, 3, 1, 4)

        qi, ki = self.qn_img(qi), self.kn_img(ki)
        qt, kt = self.qn_txt(qt), self.kn_txt(kt)

        row_cos, row_sin, col_cos, col_sin = rope_img
        qi = apply_rope_2d(qi, row_cos, row_sin, col_cos, col_sin)
        ki = apply_rope_2d(ki, row_cos, row_sin, col_cos, col_sin)

        t_cos, t_sin = rope_txt
        qt = apply_rope(qt, t_cos, t_sin)
        kt = apply_rope(kt, t_cos, t_sin)

        q = torch.cat([qi, qt], dim=2)
        k = torch.cat([ki, kt], dim=2)
        v = torch.cat([vi, vt], dim=2)

        mask = key_valid[:, None, None, :]
        o = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
        o = o.transpose(1, 2).reshape(B, Ni + Nt, C)
        oi, ot = o[:, :Ni], o[:, Ni:]

        img = img + g1i.unsqueeze(1) * self.proj_img(oi)
        txt = txt + g1t.unsqueeze(1) * self.proj_txt(ot)

        img = img + g2i.unsqueeze(1) * self.mlp_img(modulate(self.norm2_img(img), s2i, sc2i))
        txt = txt + g2t.unsqueeze(1) * self.mlp_txt(modulate(self.norm2_txt(txt), s2t, sc2t))
        return img, txt

class MMDiT(nn.Module):
    def __init__(self, latent_ch=4, latent_size=32, patch=2, dim=512, depth=16, heads=8,
                 t5_dim=768, clip_dim=512, t5_len=32, mlp_hidden=1408,
                 repa_dim=384, repa_layer=8):
        super().__init__()
        self.latent_ch = latent_ch
        self.latent_size = latent_size
        self.patch = patch
        self.grid = latent_size // patch
        self.patch_dim = latent_ch * patch * patch
        self.dim = dim
        self.depth = depth
        self.heads = heads
        self.head_dim = dim // heads
        self.t5_len = t5_len
        self.repa_layer = repa_layer

        self.x_embed = nn.Linear(self.patch_dim, dim)
        self.t_mlp = nn.Sequential(nn.Linear(dim, dim), nn.SiLU(), nn.Linear(dim, dim))
        self.clip_proj = nn.Linear(clip_dim, dim)
        self.t5_proj = nn.Linear(t5_dim, dim)

        self.blocks = nn.ModuleList([JointBlock(dim, heads, mlp_hidden) for _ in range(depth)])

        self.norm_out = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
        self.ada_out = nn.Sequential(nn.SiLU(), nn.Linear(dim, 2 * dim))
        self.head = nn.Linear(dim, self.patch_dim)

        self.repa_head = nn.Sequential(nn.Linear(dim, dim), nn.GELU(approximate="tanh"), nn.Linear(dim, repa_dim))

        hd2 = self.head_dim // 2
        rows = torch.arange(self.grid).repeat_interleave(self.grid)
        cols = torch.arange(self.grid).repeat(self.grid)
        row_cos, row_sin = rope_cos_sin(rope_freqs(rows, hd2))
        col_cos, col_sin = rope_cos_sin(rope_freqs(cols, hd2))
        self.register_buffer("row_cos", row_cos, persistent=False)
        self.register_buffer("row_sin", row_sin, persistent=False)
        self.register_buffer("col_cos", col_cos, persistent=False)
        self.register_buffer("col_sin", col_sin, persistent=False)
        t_cos, t_sin = rope_cos_sin(rope_freqs(torch.arange(t5_len), self.head_dim))
        self.register_buffer("t_cos", t_cos, persistent=False)
        self.register_buffer("t_sin", t_sin, persistent=False)

        self._init()

    def _init(self):
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.xavier_uniform_(m.weight)
                if m.bias is not None:
                    nn.init.zeros_(m.bias)
        for b in self.blocks:
            nn.init.zeros_(b.ada_img[-1].weight); nn.init.zeros_(b.ada_img[-1].bias)
            nn.init.zeros_(b.ada_txt[-1].weight); nn.init.zeros_(b.ada_txt[-1].bias)
        nn.init.zeros_(self.ada_out[-1].weight); nn.init.zeros_(self.ada_out[-1].bias)
        nn.init.zeros_(self.head.weight); nn.init.zeros_(self.head.bias)

    def patchify(self, x):
        B, C, H, W = x.shape
        p = self.patch
        x = x.reshape(B, C, H // p, p, W // p, p)
        x = x.permute(0, 2, 4, 1, 3, 5).reshape(B, (H // p) * (W // p), C * p * p)
        return x

    def unpatchify(self, x):
        B, N, _ = x.shape
        p = self.patch
        g = self.grid
        C = self.latent_ch
        x = x.reshape(B, g, g, C, p, p).permute(0, 3, 1, 4, 2, 5)
        return x.reshape(B, C, g * p, g * p)

    def forward(self, x, t, t5_seq, t5_mask, clip_pool, return_repa=False, use_checkpoint=False):
        B = x.shape[0]
        img = self.x_embed(self.patchify(x))
        txt = self.t5_proj(t5_seq)
        c = self.t_mlp(timestep_embedding(t, self.dim)) + self.clip_proj(clip_pool)

        key_valid = torch.cat([
            torch.ones(B, img.shape[1], dtype=torch.bool, device=x.device),
            t5_mask.bool(),
        ], dim=1)

        rope_img = (self.row_cos, self.row_sin, self.col_cos, self.col_sin)
        rope_txt = (self.t_cos, self.t_sin)

        repa_hidden = None
        for i, blk in enumerate(self.blocks):
            if use_checkpoint and self.training:
                img, txt = torch.utils.checkpoint.checkpoint(
                    blk, img, txt, c, rope_img, rope_txt, key_valid, use_reentrant=False)
            else:
                img, txt = blk(img, txt, c, rope_img, rope_txt, key_valid)
            if return_repa and i == self.repa_layer:
                repa_hidden = img

        shift, scale = self.ada_out(c).chunk(2, dim=-1)
        img = modulate(self.norm_out(img), shift, scale)
        out = self.unpatchify(self.head(img))

        if return_repa:
            return out, self.repa_head(repa_hidden)
        return out

    def num_params(self):
        return sum(p.numel() for p in self.parameters())

    def num_backbone_params(self):
        return sum(p.numel() for n, p in self.named_parameters() if not n.startswith("repa_head"))