"""A tiny text-to-video diffusion transformer, trained from random init with rectified flow.

Video [B,3,T,H,W] is cut into (1 x p x p) patches. Each block runs spatial attention
within a frame, temporal attention across frames at the same patch position,
cross-attention to the caption, and an MLP, all modulated (adaLN-Zero) by the
diffusion time and pooled caption. The caption encoder is a 2-layer transformer
over the small word vocabulary in data.py.
"""
from __future__ import annotations

import copy
import math
from dataclasses import asdict, dataclass

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

from .data import NULL, PAD, VOCAB, MAX_TOKENS
from .world import BYTE_VOCAB, HYBRID_VOCAB, MAX_BYTES, MAX_HYBRID, NULL as BYTE_NULL, PAD as BYTE_PAD


def specials(text: str) -> tuple[int, int]:
    """(pad id, null id) for the caption tokenizer a model was trained with."""
    return (BYTE_PAD, BYTE_NULL) if text in ("bytes", "hybrid") else (PAD, NULL)


@dataclass
class ModelConfig:
    frames: int = 16
    size: int = 48
    patch: int = 4
    dim: int = 384
    depth: int = 8
    heads: int = 6
    text_depth: int = 2
    vocab_size: int = len(VOCAB)
    max_tokens: int = MAX_TOKENS
    mlp_ratio: float = 4.0
    attn: str = "factorized"
    rope: bool = False
    text: str = "words"  # "bytes": raw UTF-8 byte tokens (world.ByteTokenizer), no vocabulary at all  # 3D rotary positions (t, y, x) in self-attention: motion becomes a relative pattern  # "joint": the first attention in each block spans all frames, not just one
    grad_checkpoint: bool = False  # recompute block activations in backward: far less VRAM, ~30% slower


# ---------------------------------------------------------------- building blocks

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


class Attention(nn.Module):
    def __init__(self, dim: int, heads: int, cross: bool = False):
        super().__init__()
        self.heads = heads
        self.q = nn.Linear(dim, dim)
        self.kv = nn.Linear(dim, 2 * dim)
        self.out = nn.Linear(dim, dim)
        hd = dim // heads
        self.q_norm = nn.RMSNorm(hd)
        self.k_norm = nn.RMSNorm(hd)

    def forward(self, x, ctx=None, mask=None, rope=None):
        B, N, D = x.shape
        ctx = x if ctx is None else ctx
        q = self.q(x).view(B, N, self.heads, -1).transpose(1, 2)
        k, v = self.kv(ctx).view(B, ctx.shape[1], 2, self.heads, -1).permute(2, 0, 3, 1, 4)
        q, k = self.q_norm(q), self.k_norm(k)
        if rope is not None:
            q, k = apply_rope(q, rope), apply_rope(k, rope)
        o = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
        return self.out(o.transpose(1, 2).reshape(B, N, D))


def rope_table(t, y, x, head_dim: int, base: float = 100.0):
    """cos/sin [L, head_dim/2] for integer positions t, y, x (each [L]). Head dims split 1/4 t, 3/8 y, 3/8 x."""
    pairs = head_dim // 2
    pt = pairs // 4
    py = (pairs - pt) // 2
    angles = []
    for pos, n in ((t, pt), (y, py), (x, pairs - pt - py)):
        freqs = base ** (-torch.arange(n, device=pos.device, dtype=torch.float32) / max(n, 1))
        angles.append(pos.float()[:, None] * freqs[None])
    a = torch.cat(angles, dim=1)
    return a.cos(), a.sin()


def apply_rope(x, rope):
    """Rotate consecutive channel pairs of x [B,H,L,hd] by the angles in rope."""
    cos, sin = rope
    x1, x2 = x.float().unflatten(-1, (-1, 2)).unbind(-1)
    out = torch.stack([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1).flatten(-2)
    return out.to(x.dtype)


class Mlp(nn.Sequential):
    def __init__(self, dim: int, ratio: float):
        hidden = int(dim * ratio)
        super().__init__(nn.Linear(dim, hidden), nn.GELU(approximate="tanh"), nn.Linear(hidden, dim))


class Block(nn.Module):
    def __init__(self, cfg: ModelConfig):
        super().__init__()
        D = cfg.dim
        self.norms = nn.ModuleList([nn.LayerNorm(D, elementwise_affine=False, eps=1e-6) for _ in range(4)])
        self.spatial = Attention(D, cfg.heads)
        self.temporal = Attention(D, cfg.heads)
        self.cross = Attention(D, cfg.heads, cross=True)
        self.mlp = Mlp(D, cfg.mlp_ratio)
        self.ada = nn.Sequential(nn.SiLU(), nn.Linear(D, 11 * D))
        assert cfg.attn in ("factorized", "joint"), cfg.attn
        self.joint = cfg.attn == "joint"

    def forward(self, x, c, ctx, ctx_mask, ropes=None):
        # x [B,T,N,D], c [B,D], ctx [B,L,D], ctx_mask [B,1,1,L] bool (True = attend)
        B, T, N, D = x.shape
        (s1, sc1, g1, s2, sc2, g2, g3, s4, sc4, g4, sc3) = self.ada(c)[:, None, None].chunk(11, dim=-1)

        r_joint, r_space, r_time = ropes or (None, None, None)
        if self.joint:
            h = modulate(self.norms[0](x), s1, sc1).reshape(B, T * N, D)
            x = x + g1 * self.spatial(h, rope=r_joint).view(B, T, N, D)
        else:
            h = modulate(self.norms[0](x), s1, sc1).reshape(B * T, N, D)
            x = x + g1 * self.spatial(h, rope=r_space).view(B, T, N, D)

        h = modulate(self.norms[1](x), s2, sc2).transpose(1, 2).reshape(B * N, T, D)
        x = x + g2 * self.temporal(h, rope=r_time).view(B, N, T, D).transpose(1, 2)

        h = (self.norms[2](x) * (1 + sc3)).reshape(B, T * N, D)
        x = x + g3 * self.cross(h, ctx, ctx_mask).view(B, T, N, D)

        h = modulate(self.norms[3](x), s4, sc4)
        return x + g4 * self.mlp(h)


class TextEncoder(nn.Module):
    def __init__(self, cfg: ModelConfig):
        super().__init__()
        self.pad = specials(cfg.text)[0]
        self.embed = nn.Embedding(cfg.vocab_size, cfg.dim)
        self.pos = nn.Parameter(torch.randn(cfg.max_tokens, cfg.dim) * 0.02)
        layer = nn.TransformerEncoderLayer(cfg.dim, cfg.heads, int(cfg.dim * cfg.mlp_ratio), dropout=0.0,
                                           activation="gelu", batch_first=True, norm_first=True)
        self.layers = nn.TransformerEncoder(layer, cfg.text_depth, enable_nested_tensor=False)
        self.norm = nn.LayerNorm(cfg.dim)

    def forward(self, tokens):
        pad = tokens == self.pad
        h = self.embed(tokens) + self.pos[: tokens.shape[1]]
        h = self.norm(self.layers(h, src_key_padding_mask=pad))
        keep = (~pad).unsqueeze(-1).to(h.dtype)
        pooled = (h * keep).sum(1) / keep.sum(1).clamp(min=1)
        return h, pooled, (~pad)[:, None, None, :]


def timestep_embedding(t, dim: int):
    half = dim // 2
    freqs = torch.exp(-math.log(10000) * torch.arange(half, device=t.device, dtype=torch.float32) / half)
    args = (t.float() * 1000)[:, None] * freqs[None]
    return torch.cat([torch.cos(args), torch.sin(args)], dim=-1)


# ---------------------------------------------------------------- the model

class VideoDiT(nn.Module):
    def __init__(self, cfg: ModelConfig | None = None):
        super().__init__()
        self.cfg = cfg = cfg or ModelConfig()
        assert cfg.size % cfg.patch == 0 and cfg.dim % cfg.heads == 0 and cfg.attn in ("factorized", "joint")
        D, p = cfg.dim, cfg.patch
        self.grid = cfg.size // p
        self.patch_embed = nn.Conv3d(3, D, kernel_size=(1, p, p), stride=(1, p, p))
        self.pos_space = nn.Parameter(torch.randn(self.grid * self.grid, D) * 0.02)
        self.pos_time = nn.Parameter(torch.randn(cfg.frames, D) * 0.02)
        self.t_embed = nn.Sequential(nn.Linear(256, D), nn.SiLU(), nn.Linear(D, D))
        self.text = TextEncoder(cfg)
        self.text_proj = nn.Linear(D, D)
        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_ada = nn.Sequential(nn.SiLU(), nn.Linear(D, 2 * D))
        self.final = nn.Linear(D, p * p * 3)
        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
        nn.init.xavier_uniform_(w.view(w.shape[0], -1))
        nn.init.normal_(self.text.embed.weight, std=0.02)
        for b in self.blocks:  # adaLN-Zero: every residual branch starts switched off
            nn.init.zeros_(b.ada[-1].weight)
            nn.init.zeros_(b.ada[-1].bias)
        for m in (self.final_ada[-1], self.final):
            nn.init.zeros_(m.weight)
            nn.init.zeros_(m.bias)

    def _ropes(self, T: int, device):
        g, hd = self.grid, self.cfg.dim // self.cfg.heads
        tt, yy, xx = torch.meshgrid(torch.arange(T, device=device), torch.arange(g, device=device),
                                    torch.arange(g, device=device), indexing="ij")
        joint = rope_table(tt.flatten(), yy.flatten(), xx.flatten(), hd)
        zeros = torch.zeros(g * g, device=device, dtype=torch.long)
        space = rope_table(zeros, yy[0].flatten(), xx[0].flatten(), hd)
        time = rope_table(torch.arange(T, device=device), zeros[:T], zeros[:T], hd)
        return joint, space, time

    def forward(self, x, t, tokens):
        """x [B,3,T,H,W], t [B] in [0,1], tokens [B,L] -> predicted velocity [B,3,T,H,W]."""
        B, _, T, H, W = x.shape
        h = self.patch_embed(x)                                  # [B,D,T,g,g]
        h = h.flatten(3).permute(0, 2, 3, 1)                     # [B,T,N,D]
        h = h + self.pos_space[None, None] + self.pos_time[None, :T, None]
        ctx, pooled, ctx_mask = self.text(tokens)
        c = self.t_embed(timestep_embedding(t, 256)) + self.text_proj(pooled)
        ropes = self._ropes(T, x.device) if self.cfg.rope else None
        for blk in self.blocks:
            if self.cfg.grad_checkpoint and self.training and torch.is_grad_enabled():
                h = checkpoint(blk, h, c, ctx, ctx_mask, ropes, use_reentrant=False)
            else:
                h = blk(h, c, ctx, ctx_mask, ropes)
        shift, scale = self.final_ada(c)[:, None, None].chunk(2, dim=-1)
        h = self.final(modulate(self.final_norm(h), shift, scale))  # [B,T,N,p*p*3]
        p, g = self.cfg.patch, self.grid
        h = h.view(B, T, g, g, p, p, 3).permute(0, 6, 1, 2, 4, 3, 5)
        return h.reshape(B, 3, T, H, W)


# ---------------------------------------------------------------- training / sampling

def null_tokens(tokens: torch.Tensor, text: str = "words") -> torch.Tensor:
    pad, null = specials(text)
    out = torch.full_like(tokens, pad)
    out[:, 0] = null
    return out


def flow_loss(model: VideoDiT, x0: torch.Tensor, tokens: torch.Tensor, p_uncond: float = 0.1) -> torch.Tensor:
    """Rectified flow: x_t = (1-t) x0 + t eps, regress v = eps - x0. Logit-normal t."""
    B = x0.shape[0]
    t = torch.sigmoid(torch.randn(B, device=x0.device))
    eps = torch.randn_like(x0)
    tt = t.view(B, 1, 1, 1, 1)
    xt = (1 - tt) * x0 + tt * eps
    drop = torch.rand(B, device=x0.device) < p_uncond
    tokens = torch.where(drop[:, None], null_tokens(tokens, model.cfg.text), tokens)
    v = model(xt, t, tokens)
    return F.mse_loss(v.float(), (eps - x0).float())


@torch.no_grad()
def sample(model: VideoDiT, tokens: torch.Tensor, steps: int = 40, cfg_scale: float = 4.0,
           generator: torch.Generator | None = None, shift: float = 2.0) -> torch.Tensor:
    """Euler integration of the flow from noise (t=1) to data (t=0). Returns [B,3,T,H,W] in [-1,1].

    `shift` > 1 spends more steps at high noise, where the layout and motion are decided.
    """
    cfg = model.cfg
    device = next(model.parameters()).device
    tokens = tokens.to(device)
    B = tokens.shape[0]
    x = torch.randn(B, 3, cfg.frames, cfg.size, cfg.size, device=device, generator=generator)
    u = torch.linspace(1, 0, steps + 1, device=device)
    ts = shift * u / (1 + (shift - 1) * u)
    cond = torch.cat([tokens, null_tokens(tokens, cfg.text)]) if cfg_scale != 1 else tokens
    for i in range(steps):
        t = ts[i].expand(cond.shape[0])
        xin = torch.cat([x, x]) if cfg_scale != 1 else x
        with torch.autocast(device.type, dtype=torch.bfloat16, enabled=device.type == "cuda"):
            v = model(xin, t, cond).float()
        if cfg_scale != 1:
            vc, vu = v.chunk(2)
            v = vu + cfg_scale * (vc - vu)
        x = x + (ts[i + 1] - ts[i]) * v
    return x.clamp(-1, 1)


def to_uint8(x: torch.Tensor) -> torch.Tensor:
    """[B,3,T,H,W] in [-1,1] -> uint8 [B,T,H,W,3] on CPU."""
    return ((x.float().clamp(-1, 1) + 1) * 127.5).round().to(torch.uint8).permute(0, 2, 3, 4, 1).cpu()


class EMA:
    def __init__(self, model: nn.Module, decay: float = 0.999):
        self.decay = decay
        self.model = copy.deepcopy(model).eval().requires_grad_(False)

    @torch.no_grad()
    def update(self, model: nn.Module):
        for e, p in zip(self.model.parameters(), model.parameters()):
            e.lerp_(p.detach(), 1 - self.decay)


def save_checkpoint(path, model: VideoDiT, opt=None, step: int = 0, extra: dict | None = None, ema: EMA | None = None):
    torch.save({
        "config": asdict(model.cfg),
        "model": model.state_dict(),
        "ema": ema.model.state_dict() if ema else None,
        "opt": opt.state_dict() if opt else None,
        "step": step,
        "extra": extra or {},
    }, path)


def load_checkpoint(path, device="cpu", use_ema: bool = True):
    ckpt = torch.load(path, map_location=device, weights_only=False)
    model = VideoDiT(ModelConfig(**ckpt["config"]))
    state = ckpt["ema"] if use_ema and ckpt.get("ema") else ckpt["model"]
    model.load_state_dict(state)
    return model.to(device).eval(), ckpt


def tokenizer_for(cfg: ModelConfig):
    """The caption tokenizer matching a model's config (v1 word vocabulary or v2 raw bytes)."""
    if cfg.text == "bytes":
        from .world import ByteTokenizer
        return ByteTokenizer(cfg.max_tokens)
    if cfg.text == "hybrid":
        from .world import HybridTokenizer
        return HybridTokenizer(cfg.max_tokens)
    from .data import Tokenizer
    return Tokenizer(cfg.max_tokens)


def byte_model_config(**kw) -> ModelConfig:
    """Defaults for world-v2 models: byte captions, joint space-time attention, 3D RoPE."""
    base = dict(frames=16, size=32, patch=4, dim=384, depth=8, heads=6, text_depth=4, vocab_size=BYTE_VOCAB,
                max_tokens=MAX_BYTES, attn="joint", rope=True, text="bytes")
    return ModelConfig(**{**base, **kw})


def hybrid_model_config(**kw) -> ModelConfig:
    """byte_model_config with hashed-word tokens added to the caption (text="hybrid")."""
    return byte_model_config(**{"text": "hybrid", "vocab_size": HYBRID_VOCAB, "max_tokens": MAX_HYBRID, **kw})


def load_matching(model: nn.Module, state: dict) -> dict:
    """Copy every tensor whose name and shape match; for embedding tables that grew, copy the overlapping rows.
    Returns {"loaded": [...], "partial": [...], "skipped": [...]} for provenance."""
    own = model.state_dict()
    report = {"loaded": [], "partial": [], "skipped": []}
    with torch.no_grad():
        for k, v in state.items():
            if k not in own:
                report["skipped"].append(k)
            elif own[k].shape == v.shape:
                own[k].copy_(v)
                report["loaded"].append(k)
            elif own[k].ndim == v.ndim and all(a >= b for a, b in zip(own[k].shape, v.shape)):
                own[k][tuple(slice(0, n) for n in v.shape)].copy_(v)
                report["partial"].append(k)
            else:
                report["skipped"].append(k)
    return report