# bora_dit.py -- Bora V1: a flow-matching video DiT over Wan2.1-VAE latents, plus the loss and the block sampler.
#
# Input: a block of F=4 latent frames [B,16,4,32,32] (normalized latents from bora_dataset.BoraLatents).
# Conditioning: T5 token embeddings [B,32,768] + mask (cross-attention), timestep t in [0,1] (adaLN),
#               and optionally the clean first frame of the block (channel-concat + mask channel) so that block k
#               can continue from the last frame of block k-1 at generation time.
# Objective: rectified flow. x_t = (1-t) x + t eps, model predicts v = eps - x, MSE on the frames it has to generate.
# Sampling: Euler, 8 steps by default, classifier-free guidance on the text, blocks chained with 1 frame of overlap.
#
# Sizes (dim/depth/heads): v1 = 1024/20/16 -> 340M params, 1024 tokens per block with patch (1,2,2).
import math
from dataclasses import dataclass, asdict

import torch
import torch.nn as nn
import torch.nn.functional as F


# ----------------------------------------------------------------------------- config

@dataclass
class DiTConfig:
    in_channels: int = 16
    frames: int = 4
    height: int = 32
    width: int = 32
    patch: tuple = (1, 2, 2)          # (t, h, w) -> 4 x 16 x 16 = 1024 tokens
    dim: int = 1024
    depth: int = 20
    heads: int = 16
    mlp_ratio: float = 4.0
    text_dim: int = 768               # T5 v1.1 base
    attn: str = "softmax"             # "softmax" (3D RoPE) or "linear" (ReLU linear attention, learned positions)
    qk_norm: bool = True
    frame_cond: bool = True           # accept a clean first frame (adds 16 + 1 input channels)
    time_shift: float = 3.0           # SD3-style timestep shift toward noise, used in training and sampling

    @property
    def tokens(self):
        return (self.frames // self.patch[0]) * (self.height // self.patch[1]) * (self.width // self.patch[2])

    @property
    def grid(self):
        return (self.frames // self.patch[0], self.height // self.patch[1], self.width // self.patch[2])


PRESETS = {
    "tiny": dict(dim=128, depth=2, heads=4),                 # smoke tests
    "small": dict(dim=512, depth=12, heads=8),               # ~50M, quick experiments
    "v1": dict(dim=1024, depth=20, heads=16),                # ~340M, the plan
}


# ----------------------------------------------------------------------------- pieces

def timestep_embedding(t, dim, max_period=10000):
    """t in [0,1] -> sinusoidal [B, dim]. Scaled by 1000 so the frequencies match the usual 0..1000 diffusion range."""
    half = dim // 2
    freqs = torch.exp(-math.log(max_period) * torch.arange(half, device=t.device, dtype=torch.float32) / half)
    args = (t.float() * 1000.0)[:, None] * freqs[None]
    return torch.cat([args.cos(), args.sin()], dim=-1)


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

    def forward(self, x):
        xf = x.float()
        return (xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + self.eps)).to(x.dtype) * self.weight


def rope_3d(grid, head_dim, theta=10000.0):
    """cos/sin tables [L, head_dim/2] for a (T,H,W) token grid. head_dim/2 angles split 1:1.5:1.5 between t, h, w."""
    T, H, W = grid
    n = head_dim // 2
    nt = (n // 4) if n >= 4 else 1
    nh = (n - nt) // 2
    nw = n - nt - nh
    def axis(num, pos):
        freqs = 1.0 / (theta ** (torch.arange(num, dtype=torch.float32) / num))
        return pos.float()[:, None] * freqs[None]                     # [P, num]
    t_idx, h_idx, w_idx = torch.meshgrid(torch.arange(T), torch.arange(H), torch.arange(W), indexing="ij")
    ang = torch.cat([axis(nt, t_idx.flatten()), axis(nh, h_idx.flatten()), axis(nw, w_idx.flatten())], dim=-1)  # [L, n]
    return ang.cos(), ang.sin()


def apply_rope(x, cos, sin):
    """x [B, heads, L, d]; pairs (i, i + d/2) are rotated by the angle i."""
    x1, x2 = x.float().chunk(2, dim=-1)
    out = torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
    return out.to(x.dtype)


class Attention(nn.Module):
    def __init__(self, dim, heads, qk_norm, kind="softmax", ctx_dim=None):
        super().__init__()
        self.heads, self.hd, self.kind = heads, dim // heads, kind
        ctx_dim = ctx_dim or dim
        self.q = nn.Linear(dim, dim)
        self.kv = nn.Linear(ctx_dim, 2 * dim)
        self.o = nn.Linear(dim, dim)
        self.qn = RMSNorm(self.hd) if qk_norm else nn.Identity()
        self.kn = RMSNorm(self.hd) if qk_norm else nn.Identity()

    def forward(self, x, ctx=None, ctx_mask=None, rope=None):
        B, L, _ = x.shape
        ctx = x if ctx is None else ctx
        q = self.q(x).view(B, L, self.heads, self.hd).transpose(1, 2)                     # [B,H,L,d]
        k, v = self.kv(ctx).view(B, ctx.shape[1], 2, self.heads, self.hd).permute(2, 0, 3, 1, 4)
        q, k = self.qn(q), self.kn(k)
        if rope is not None:
            q, k = apply_rope(q, *rope), apply_rope(k, *rope)
        if self.kind == "linear":
            out = self._linear(q, k, v)
        else:
            mask = None
            if ctx_mask is not None:
                mask = ctx_mask[:, None, None, :]                                           # True = attend
            out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
        return self.o(out.transpose(1, 2).reshape(B, L, -1))

    @staticmethod
    def _linear(q, k, v, eps=1e-6):
        """ReLU linear attention (SANA). fp32 for the sums; O(L d^2) instead of O(L^2 d)."""
        q, k, v = F.relu(q.float()), F.relu(k.float()), v.float()
        kv = torch.einsum("bhld,bhle->bhde", k, v)                                        # [B,H,d,d]
        z = 1.0 / (torch.einsum("bhld,bhd->bhl", q, k.sum(2)) + eps)                        # [B,H,L]
        out = torch.einsum("bhld,bhde->bhle", q, kv) * z[..., None]
        return out.to(v.dtype) if v.dtype != torch.float32 else out


class Block(nn.Module):
    """adaLN-zero DiT block with a shared modulation table (PixArt-alpha): the 6 x dim modulation comes from one global
    MLP on the timestep plus a per-block learned offset, instead of a 6 x dim^2 Linear per block."""
    def __init__(self, cfg):
        super().__init__()
        d = cfg.dim
        self.n1 = nn.LayerNorm(d, elementwise_affine=False, eps=1e-6)
        self.attn = Attention(d, cfg.heads, cfg.qk_norm, cfg.attn)
        self.n2 = nn.LayerNorm(d, elementwise_affine=False, eps=1e-6)
        self.cross = Attention(d, cfg.heads, cfg.qk_norm, "softmax", ctx_dim=cfg.text_dim)
        self.n3 = nn.LayerNorm(d, elementwise_affine=False, eps=1e-6)
        hidden = int(d * cfg.mlp_ratio)
        self.mlp = nn.Sequential(nn.Linear(d, hidden), nn.GELU(approximate="tanh"), nn.Linear(hidden, d))
        self.table = nn.Parameter(torch.randn(6, d) / d ** 0.5)

    def forward(self, x, t6, text, text_mask, rope):
        # t6: [B, 6, dim] from the shared timestep MLP
        shift1, scale1, gate1, shift2, scale2, gate2 = (self.table[None] + t6).unbind(1)
        h = self.n1(x) * (1 + scale1[:, None]) + shift1[:, None]
        x = x + gate1[:, None] * self.attn(h, rope=rope)
        x = x + self.cross(self.n2(x), ctx=text, ctx_mask=text_mask)
        h = self.n3(x) * (1 + scale2[:, None]) + shift2[:, None]
        x = x + gate2[:, None] * self.mlp(h)
        return x


# ----------------------------------------------------------------------------- model

class BoraDiT(nn.Module):
    def __init__(self, cfg: DiTConfig):
        super().__init__()
        self.cfg = cfg
        self.grad_checkpoint = False          # set True to trade ~2x block compute for ~5x less activation memory
        d, p = cfg.dim, cfg.patch
        in_ch = cfg.in_channels + (cfg.in_channels + 1 if cfg.frame_cond else 0)
        self.patch_embed = nn.Conv3d(in_ch, d, kernel_size=p, stride=p)
        self.t_embed = nn.Sequential(nn.Linear(256, d), nn.SiLU(), nn.Linear(d, d))
        self.t_mod = nn.Sequential(nn.SiLU(), nn.Linear(d, 6 * d))                       # shared adaLN
        self.text_proj = nn.Sequential(nn.Linear(cfg.text_dim, cfg.text_dim), nn.GELU(approximate="tanh"), nn.Linear(cfg.text_dim, cfg.text_dim))
        self.blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.depth)])
        self.final_norm = nn.LayerNorm(d, elementwise_affine=False, eps=1e-6)
        self.final_mod = nn.Sequential(nn.SiLU(), nn.Linear(d, 2 * d))
        self.final = nn.Linear(d, p[0] * p[1] * p[2] * cfg.in_channels)
        if cfg.attn == "linear":                                                             # no RoPE for linear attention
            self.pos = nn.Parameter(torch.randn(1, cfg.tokens, d) * 0.02)
        else:
            cos, sin = rope_3d(cfg.grid, d // cfg.heads)
            self.register_buffer("rope_cos", cos, persistent=False)
            self.register_buffer("rope_sin", sin, persistent=False)
        self._init()

    def _init(self):
        def basic(m):
            if isinstance(m, nn.Linear):
                nn.init.xavier_uniform_(m.weight)
                if m.bias is not None:
                    nn.init.zeros_(m.bias)
        self.apply(basic)
        w = self.patch_embed.weight.data
        nn.init.xavier_uniform_(w.view(w.shape[0], -1))
        nn.init.zeros_(self.patch_embed.bias)
        nn.init.normal_(self.t_embed[0].weight, std=0.02)
        nn.init.normal_(self.t_embed[2].weight, std=0.02)
        nn.init.zeros_(self.t_mod[1].weight); nn.init.zeros_(self.t_mod[1].bias)           # adaLN-zero: gates start at 0
        nn.init.zeros_(self.final_mod[1].weight); nn.init.zeros_(self.final_mod[1].bias)
        nn.init.zeros_(self.final.weight); nn.init.zeros_(self.final.bias)                 # output starts at 0
        for b in self.blocks:
            nn.init.zeros_(b.cross.o.weight); nn.init.zeros_(b.cross.o.bias)              # cross-attn starts silent

    def unpatchify(self, x):
        B = x.shape[0]
        T, H, W = self.cfg.grid
        pt, ph, pw = self.cfg.patch
        c = self.cfg.in_channels
        x = x.view(B, T, H, W, pt, ph, pw, c)
        x = x.permute(0, 7, 1, 4, 2, 5, 3, 6)                                              # B c T pt H ph W pw
        return x.reshape(B, c, T * pt, H * ph, W * pw)

    def forward(self, x_t, t, text, text_mask, cond=None, cond_mask=None):
        """x_t [B,16,F,H,W], t [B] in [0,1], text [B,L,768], text_mask [B,L] bool,
        cond [B,16,F,H,W] clean frames (zeros where unknown), cond_mask [B,1,F,1,1] in {0,1}. Returns v [B,16,F,H,W]."""
        B = x_t.shape[0]
        if self.cfg.frame_cond:
            if cond is None:
                cond = torch.zeros_like(x_t)
                cond_mask = torch.zeros(B, 1, x_t.shape[2], 1, 1, device=x_t.device, dtype=x_t.dtype)
            m = cond_mask.expand(B, 1, *x_t.shape[2:]).to(x_t.dtype)
            x_in = torch.cat([x_t, cond * m, m], dim=1)
        else:
            x_in = x_t
        x = self.patch_embed(x_in).flatten(2).transpose(1, 2)                               # [B,L,dim]
        if self.cfg.attn == "linear":
            x = x + self.pos
            rope = None
        else:
            rope = (self.rope_cos, self.rope_sin)
        te = self.t_embed(timestep_embedding(t, 256).to(x.dtype))                            # [B,dim]
        t6 = self.t_mod(te).view(B, 6, -1)
        text = self.text_proj(text.to(x.dtype))
        for blk in self.blocks:
            if self.grad_checkpoint and self.training:
                x = torch.utils.checkpoint.checkpoint(blk, x, t6, text, text_mask, rope, use_reentrant=False)
            else:
                x = blk(x, t6, text, text_mask, rope)
        shift, scale = self.final_mod(te).chunk(2, dim=-1)
        x = self.final_norm(x) * (1 + scale[:, None]) + shift[:, None]
        return self.unpatchify(self.final(x))


# ----------------------------------------------------------------------------- flow matching

def shift_t(u, shift):
    """SD3 timestep shift: pushes samples toward t=1 (more noise). Identity for shift=1."""
    return shift * u / (1 + (shift - 1) * u)


def sample_timesteps(B, shift, device):
    u = torch.sigmoid(torch.randn(B, device=device))                                        # logit-normal(0,1)
    return shift_t(u, shift)


def make_frame_cond(x, p_cond, generator=None):
    """With probability p_cond per sample, give the clean first frame of the block as conditioning.
    Returns cond [B,16,F,H,W], cond_mask [B,1,F,1,1], loss_mask [B,1,F,1,1] (0 on the given frame)."""
    B, C, Fr, H, W = x.shape
    use = (torch.rand(B, device=x.device, generator=generator) < p_cond).to(x.dtype)        # [B]
    cond_mask = torch.zeros(B, 1, Fr, 1, 1, device=x.device, dtype=x.dtype)
    cond_mask[:, :, 0] = use.view(B, 1, 1, 1)
    cond = torch.zeros_like(x)
    cond[:, :, 0] = x[:, :, 0]
    return cond, cond_mask, 1 - cond_mask


def flow_loss(model, x, text, text_mask, p_cond=0.5, shift=3.0, t=None, noise=None):
    """Rectified-flow MSE. x: clean block [B,16,F,H,W]. Returns (loss, dict of parts)."""
    B = x.shape[0]
    t = sample_timesteps(B, shift, x.device) if t is None else t
    noise = torch.randn_like(x) if noise is None else noise
    tt = t.view(B, 1, 1, 1, 1).to(x.dtype)
    x_t = (1 - tt) * x + tt * noise
    cond, cond_mask, loss_mask = make_frame_cond(x, p_cond)
    v = model(x_t, t, text, text_mask, cond, cond_mask)
    target = noise - x
    err = (v.float() - target.float()) ** 2 * loss_mask
    loss = err.sum() / (loss_mask.sum() * x.shape[1] * x.shape[3] * x.shape[4])
    return loss, {"t_mean": float(t.mean()), "cond_frac": float(cond_mask[:, :, 0].mean())}


@torch.no_grad()
def sample_block(model, text, text_mask, null_text, null_mask, steps=8, cfg_scale=4.0, shift=3.0,
                 first_frame=None, generator=None, dtype=torch.bfloat16):
    """Euler sampler for one block. first_frame [B,16,H,W] (clean) continues a previous block; None starts fresh.
    Returns clean latents [B,16,F,H,W] (float32)."""
    c = model.cfg
    B = text.shape[0]
    dev = text.device
    x = torch.randn(B, c.in_channels, c.frames, c.height, c.width, device=dev, generator=generator)
    cond = cond_mask = None
    eps0 = None
    if first_frame is not None:
        cond = torch.zeros_like(x)
        cond[:, :, 0] = first_frame
        cond_mask = torch.zeros(B, 1, c.frames, 1, 1, device=dev)
        cond_mask[:, :, 0] = 1
        eps0 = torch.randn(B, c.in_channels, c.height, c.width, device=dev, generator=generator)
    ts = shift_t(torch.linspace(1, 0, steps + 1, device=dev), shift)
    for i in range(steps):
        t, t_next = ts[i], ts[i + 1]
        if first_frame is not None:                                                            # keep the given frame on its own noising path
            x[:, :, 0] = (1 - t) * first_frame + t * eps0
        tb = t.expand(B)
        with torch.autocast(device_type=dev.type, dtype=dtype, enabled=dev.type == "cuda"):
            if cfg_scale != 1.0:
                xx = torch.cat([x, x]); tt = torch.cat([tb, tb])
                te = torch.cat([text, null_text]); tm = torch.cat([text_mask, null_mask])
                cc = None if cond is None else torch.cat([cond, cond])
                cm = None if cond_mask is None else torch.cat([cond_mask, cond_mask])
                v_c, v_u = model(xx.to(dtype), tt, te, tm, cc if cc is None else cc.to(dtype), cm if cm is None else cm.to(dtype)).float().chunk(2)
                v = v_u + cfg_scale * (v_c - v_u)
            else:
                v = model(x.to(dtype), tb, text, text_mask, None if cond is None else cond.to(dtype), None if cond_mask is None else cond_mask.to(dtype)).float()
        x = x + (t_next - t) * v
    if first_frame is not None:
        x[:, :, 0] = first_frame
    return x


@torch.no_grad()
def sample_video(model, text, text_mask, null_text, null_mask, n_blocks=3, overlap=1, **kw):
    """Chain blocks with `overlap` shared frames: 4 + 3 + 3 = 10 latent frames for n_blocks=3 (one full 37-frame clip)."""
    blocks = []
    first = None
    for _ in range(n_blocks):
        blk = sample_block(model, text, text_mask, null_text, null_mask, first_frame=first, **kw)
        blocks.append(blk if first is None else blk[:, :, overlap:])
        first = blk[:, :, -1]
    return torch.cat(blocks, dim=2)


# ----------------------------------------------------------------------------- helpers

def build(preset="v1", **overrides):
    cfg = DiTConfig(**{**PRESETS[preset], **overrides})
    return BoraDiT(cfg), cfg


def count_params(model):
    return sum(p.numel() for p in model.parameters())


def config_dict(cfg):
    d = asdict(cfg)
    d["patch"] = list(d["patch"])
    return d