"""
T.I.T.S.2 — a Diffusion Transformer (DiT), i.e. this one actually is a transformer.

Works on 4x32x32 SD-VAE latents (256x256 images). Patchify 2x2 -> 256 tokens, then a
stack of transformer blocks. Each block does:

    self-attention  : every patch talks to every other patch
    cross-attention : every patch reads the prompt, token by token
    MLP

Timestep + pooled text conditioning goes in via adaLN-Zero (the DiT paper's trick:
each block starts as an identity function and learns its way out, which makes deep
stacks train stably from scratch).

~93M params at defaults (dim=576, depth=12, heads=9).
"""
import math

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


def timestep_embedding(t: torch.Tensor, dim: int) -> torch.Tensor:
    """Sinusoidal embedding. t is continuous in [0,1] here (flow matching), so scale up."""
    half = dim // 2
    freqs = torch.exp(-math.log(10000) * torch.arange(half, device=t.device, dtype=torch.float32) / half)
    args = t.float()[:, None] * 1000.0 * freqs[None]
    return torch.cat([torch.cos(args), torch.sin(args)], dim=-1)


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


class Attention(nn.Module):
    """Self-attention, or cross-attention when `context_dim` is given."""

    def __init__(self, dim: int, heads: int, context_dim: int = None):
        super().__init__()
        self.heads = heads
        self.head_dim = dim // heads
        self.q = nn.Linear(dim, dim, bias=False)
        self.kv = nn.Linear(context_dim or dim, dim * 2, bias=False)
        self.out = nn.Linear(dim, dim)
        self.cross = context_dim is not None

    def forward(self, x, context=None, context_mask=None):
        b, n, _ = x.shape
        ctx = context if self.cross else x
        q = self.q(x).view(b, n, self.heads, self.head_dim).transpose(1, 2)
        k, v = self.kv(ctx).view(b, ctx.shape[1], 2, self.heads, self.head_dim).permute(2, 0, 3, 1, 4)
        mask = None
        if self.cross and context_mask is not None:
            mask = context_mask[:, None, None, :].bool()  # (b, 1, 1, ctx_len)
        x = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
        return self.out(x.transpose(1, 2).reshape(b, n, -1))


class DiTBlock(nn.Module):
    def __init__(self, dim: int, heads: int, text_dim: int, mlp_ratio: float = 4.0):
        super().__init__()
        self.norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
        self.attn = Attention(dim, heads)
        self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
        self.cross_attn = Attention(dim, heads, context_dim=text_dim)
        self.norm3 = 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))
        # adaLN-Zero: predicts shift/scale/gate for self-attn and MLP from the conditioning vector.
        self.ada = nn.Sequential(nn.SiLU(), nn.Linear(dim, 6 * dim))
        nn.init.zeros_(self.ada[1].weight)
        nn.init.zeros_(self.ada[1].bias)

    def forward(self, x, cond, context, context_mask):
        shift1, scale1, gate1, shift2, scale2, gate2 = self.ada(cond).chunk(6, dim=-1)
        x = x + gate1.unsqueeze(1) * self.attn(modulate(self.norm1(x), shift1, scale1))
        x = x + self.cross_attn(self.norm2(x), context, context_mask)
        x = x + gate2.unsqueeze(1) * self.mlp(modulate(self.norm3(x), shift2, scale2))
        return x


class TITS2(nn.Module):
    def __init__(
        self,
        latent_size: int = 32,
        in_ch: int = 4,
        patch_size: int = 2,
        dim: int = 576,
        depth: int = 12,
        heads: int = 9,
        text_dim: int = 768,
        use_checkpoint: bool = False,
    ):
        super().__init__()
        self.patch_size = patch_size
        self.in_ch = in_ch
        self.latent_size = latent_size
        self.use_checkpoint = use_checkpoint
        n_patches = (latent_size // patch_size) ** 2

        self.x_embed = nn.Conv2d(in_ch, dim, patch_size, stride=patch_size)
        self.pos = nn.Parameter(torch.zeros(1, n_patches, dim))
        nn.init.normal_(self.pos, std=0.02)

        self.t_mlp = nn.Sequential(nn.Linear(dim, dim), nn.SiLU(), nn.Linear(dim, dim))
        self.pooled_text = nn.Sequential(nn.Linear(text_dim, dim), nn.SiLU(), nn.Linear(dim, dim))
        self.text_proj = nn.Linear(text_dim, text_dim)

        self.blocks = nn.ModuleList([DiTBlock(dim, heads, text_dim) 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.proj_out = nn.Linear(dim, patch_size * patch_size * in_ch)
        nn.init.zeros_(self.ada_out[1].weight)
        nn.init.zeros_(self.ada_out[1].bias)
        nn.init.zeros_(self.proj_out.weight)
        nn.init.zeros_(self.proj_out.bias)

    def unpatchify(self, x):
        p, c = self.patch_size, self.in_ch
        h = w = self.latent_size // p
        x = x.reshape(x.shape[0], h, w, p, p, c).permute(0, 5, 1, 3, 2, 4)
        return x.reshape(x.shape[0], c, h * p, w * p)

    def forward(self, x, t, text_tokens, text_pooled, text_mask=None):
        """x: (b,4,32,32) noisy latent. t: (b,) in [0,1]. text_tokens: (b,77,768)."""
        h = self.x_embed(x).flatten(2).transpose(1, 2) + self.pos
        cond = self.t_mlp(timestep_embedding(t, h.shape[-1])) + self.pooled_text(text_pooled)
        context = self.text_proj(text_tokens)

        for blk in self.blocks:
            if self.use_checkpoint and self.training:
                h = checkpoint(blk, h, cond, context, text_mask, use_reentrant=False)
            else:
                h = blk(h, cond, context, text_mask)

        shift, scale = self.ada_out(cond).chunk(2, dim=-1)
        h = modulate(self.norm_out(h), shift, scale)
        return self.unpatchify(self.proj_out(h))


if __name__ == "__main__":
    net = TITS2()
    print(f"TITS2 params: {sum(p.numel() for p in net.parameters()) / 1e6:.2f}M")
    x = torch.randn(2, 4, 32, 32)
    out = net(x, torch.rand(2), torch.randn(2, 77, 768), torch.randn(2, 768))
    print("output:", out.shape)