from __future__ import annotations

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

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 sincos_2d(dim, grid_size):
    g = np.arange(grid_size, dtype=np.float32)
    gx, gy = np.meshgrid(g, g, indexing="xy")
    assert dim % 4 == 0
    d4 = dim // 4
    omega = 1.0 / (10000 ** (np.arange(d4, dtype=np.float32) / d4))
    def emb1(p):
        out = p.reshape(-1)[:, None] * omega[None]
        return np.concatenate([np.sin(out), np.cos(out)], axis=1)
    pe = np.concatenate([emb1(gx), emb1(gy)], axis=1)
    return torch.from_numpy(pe).float()

class Attention(nn.Module):
    def __init__(self, dim, heads):
        super().__init__()
        self.heads = heads
        self.q = nn.Linear(dim, dim)
        self.kv = nn.Linear(dim, dim * 2)
        self.proj = nn.Linear(dim, dim)

    def forward(self, x, ctx=None):
        ctx = x if ctx is None else ctx
        B, N, C = x.shape
        M = ctx.shape[1]
        h = self.heads
        q = self.q(x).reshape(B, N, h, C // h).transpose(1, 2)
        kv = self.kv(ctx).reshape(B, M, 2, h, C // h).permute(2, 0, 3, 1, 4)
        k, v = kv[0], kv[1]
        o = F.scaled_dot_product_attention(q, k, v)
        o = o.transpose(1, 2).reshape(B, N, C)
        return self.proj(o)

class Block(nn.Module):
    def __init__(self, dim, heads, mlp_ratio=4.0):
        super().__init__()
        self.norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
        self.attn = Attention(dim, heads)
        self.norm_ca = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
        self.cross = Attention(dim, heads)
        self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
        hidden = int(dim * mlp_ratio)
        self.mlp = nn.Sequential(nn.Linear(dim, hidden), nn.GELU(approximate="tanh"),
                                 nn.Linear(hidden, dim))
        self.ada = nn.Sequential(nn.SiLU(), nn.Linear(dim, 6 * dim))
        self.cross_gate = nn.Parameter(torch.zeros(1))

    def forward(self, x, c, text):
        shift1, scale1, gate1, shift2, scale2, gate2 = self.ada(c).chunk(6, dim=1)
        x = x + gate1.unsqueeze(1) * self.attn(modulate(self.norm1(x), shift1, scale1))
        x = x + self.cross_gate * self.cross(self.norm_ca(x), text)
        x = x + gate2.unsqueeze(1) * self.mlp(modulate(self.norm2(x), shift2, scale2))
        return x

class DiT(nn.Module):
    def __init__(self, latent_ch=4, latent_size=32, patch=2, dim=384, depth=12,
                 heads=6, text_dim=512, mlp_ratio=4.0):
        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.x_embed = nn.Linear(self.patch_dim, dim)
        self.register_buffer("pos", sincos_2d(dim, self.grid).unsqueeze(0))
        self.t_mlp = nn.Sequential(nn.Linear(dim, dim), nn.SiLU(), nn.Linear(dim, dim))
        self.text_proj = nn.Linear(text_dim, dim)
        self.text_pool = nn.Linear(text_dim, dim)
        self.blocks = nn.ModuleList([Block(dim, heads, mlp_ratio) 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.dim = dim
        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[-1].weight); nn.init.zeros_(b.ada[-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, text_seq, text_pool):
        x = self.x_embed(self.patchify(x)) + self.pos
        c = self.t_mlp(timestep_embedding(t, self.dim)) + self.text_pool(text_pool)
        text = self.text_proj(text_seq)
        for blk in self.blocks:
            x = blk(x, c, text)
        shift, scale = self.ada_out(c).chunk(2, dim=1)
        x = modulate(self.norm_out(x), shift, scale)
        x = self.head(x)
        return self.unpatchify(x)

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