"""
MetaDiffusionLM: masked-diffusion LM converted from Qwen3-0.6B-Instruct.

Architecture: 28L x 1024W, GQA (16Q / 8KV, head_dim 128), QK-norm, RoPE,
timestep conditioning (sinusoidal MLP + zero-init per-block residual).
Bidirectional attention (no causal mask) is the key difference from the AR source.

Parameter init:
- Copied from AR checkpoint: token embeddings, all transformer blocks, norms,
  QK-norm, RoPE buffers.
- New, zero-init (identity at step 0): timestep embedding MLP, per-block
  timestep residual. The model starts as exactly the AR model; diffusion
  behavior is learned on top.
- New, mean-init: [MASK] token row and the 7 rainbow padding token rows
  (appended to both embed_tokens and the untied lm_head).

Training loss (train.py):
- CE on masked positions (the diffusion objective).
"""

import math
from dataclasses import asdict, dataclass, field
from typing import Optional

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


@dataclass
class MetaDiffusionConfig:
    hidden_size: int = 1024
    intermediate_size: int = 3072
    num_hidden_layers: int = 28
    num_attention_heads: int = 16
    num_key_value_heads: int = 8
    head_dim: int = 128
    vocab_size: int = 151669          # Qwen3 real tokenizer vocab (config 151936 is TP-padded)
    mask_vocab_size: int = 151677     # + [MASK] + 7 rainbow tokens
    mask_token_id: int = 151669
    pad_token_id: int = 151643        # <|endoftext|> in Qwen3
    max_position_embeddings: int = 32768
    rope_theta: float = 1000000.0
    rms_norm_eps: float = 1e-6
    hidden_act: str = "silu"
    qk_norm: bool = True
    timestep_emb_hidden: int = 1024
    tie_word_embeddings: bool = False
    mask_ratio_min: float = 0.0
    mask_ratio_max: float = 1.0
    dtype: str = "float32"


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

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


class RotaryEmbedding(nn.Module):
    def __init__(self, dim, max_position_embeddings=32768, base=1000000.0):
        super().__init__()
        self.dim = dim
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)
        self.max_position_embeddings = max_position_embeddings

    def forward(self, x, position_ids):
        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(
            position_ids.shape[0], -1, 1
        )
        position_ids_expanded = position_ids[:, None, :].float()
        freqs = (inv_freq_expanded @ position_ids_expanded).transpose(1, 2)
        emb = torch.cat((freqs, freqs), dim=-1)
        cos = emb.cos().to(dtype=x.dtype)
        sin = emb.sin().to(dtype=x.dtype)
        return cos, sin


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


def apply_rotary_pos_emb(q, k, cos, sin):
    cos = cos.unsqueeze(1)
    sin = sin.unsqueeze(1)
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed, k_embed


class TimestepEmbedding(nn.Module):
    """Sinusoidal timestep embedding with learned MLP projection."""

    def __init__(self, hidden_size):
        super().__init__()
        self.hidden_size = hidden_size
        self.mlp = nn.Sequential(
            nn.Linear(hidden_size, hidden_size * 4),
            nn.SiLU(),
            nn.Linear(hidden_size * 4, hidden_size),
        )

    def forward(self, t):
        half_dim = self.hidden_size // 2
        emb = math.log(10000.0) / (half_dim - 1)
        emb = torch.exp(
            torch.arange(half_dim, device=t.device, dtype=torch.float32) * -emb
        )
        emb = t[:, None].float() * emb[None, :]
        emb = torch.cat([emb.sin(), emb.cos()], dim=-1)
        # cast to the MLP weight dtype: the model may be bf16 while t is fp32
        return self.mlp(emb.to(self.mlp[0].weight.dtype))


class TimestepModulation(nn.Module):
    """adaLN-style timestep conditioning: scale + shift the hidden state.

    Zero-init scale/shift so the model is a pure copy of the AR model at
    step 0. Unlike the old zero-init ADDITIVE residual (TimestepResidual),
    the gradient here is dL/dscale = dL/dx * x with x nonzero, so the
    t-path trains: the additive version deadlocked (zero output through a
    zero weight = zero outer-product gradient forever), leaving the model
    noise-schedule-agnostic."""

    def __init__(self, hidden_size):
        super().__init__()
        self.proj = nn.Linear(hidden_size, hidden_size * 2)
        nn.init.zeros_(self.proj.weight)
        nn.init.zeros_(self.proj.bias)

    def forward(self, x, emb):
        scale, shift = self.proj(emb).chunk(2, dim=-1)
        scale, shift = scale[:, None, :], shift[:, None, :]
        return x * (1.0 + scale) + shift


class SelfAttention(nn.Module):
    """GQA attention, bidirectional (no causal mask), optional QK-norm."""

    def __init__(self, config):
        super().__init__()
        self.config = config
        self.hidden_size = config.hidden_size
        self.num_heads = config.num_attention_heads
        self.num_kv_heads = config.num_key_value_heads
        self.head_dim = config.head_dim
        self.num_kv_groups = self.num_heads // self.num_kv_heads

        self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(self.num_heads * self.head_dim, config.hidden_size, bias=False)

        self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) if config.qk_norm else nn.Identity()
        self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) if config.qk_norm else nn.Identity()

        self.rotary_emb = RotaryEmbedding(
            config.head_dim,
            max_position_embeddings=config.max_position_embeddings,
            base=config.rope_theta,
        )

    def forward(self, x, attention_mask=None, position_ids=None):
        batch, seq, _ = x.shape

        q = self.q_proj(x).view(batch, seq, self.num_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(batch, seq, self.num_kv_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(batch, seq, self.num_kv_heads, self.head_dim).transpose(1, 2)

        q = self.q_norm(q)
        k = self.k_norm(k)

        cos, sin = self.rotary_emb(x, position_ids)
        q, k = apply_rotary_pos_emb(q, k, cos, sin)

        if self.num_kv_groups > 1:
            k = k.repeat_interleave(self.num_kv_groups, dim=1)
            v = v.repeat_interleave(self.num_kv_groups, dim=1)

        out = F.scaled_dot_product_attention(q, k, v, attn_mask=attention_mask)
        out = out.transpose(1, 2).contiguous().view(batch, seq, -1)
        return self.o_proj(out)


class MLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)

    def forward(self, x):
        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))


class TransformerBlock(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.self_attn = SelfAttention(config)
        self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.mlp = MLP(config)
        self.timestep_modulation = TimestepModulation(config.hidden_size)

    def forward(self, x, timestep_emb, attention_mask=None, position_ids=None):
        residual = x
        x = self.input_layernorm(x)
        x = self.self_attn(x, attention_mask, position_ids)
        x = residual + x
        x = self.timestep_modulation(x, timestep_emb)

        residual = x
        x = self.post_attention_layernorm(x)
        x = self.mlp(x)
        x = residual + x
        x = self.timestep_modulation(x, timestep_emb)
        return x


class MetaDiffusionLM(nn.Module):
    def __init__(self, config: MetaDiffusionConfig):
        super().__init__()
        self.config = config
        self.embed_tokens = nn.Embedding(config.mask_vocab_size, config.hidden_size)
        self.timestep_emb = TimestepEmbedding(config.timestep_emb_hidden)
        self.layers = nn.ModuleList([TransformerBlock(config) for _ in range(config.num_hidden_layers)])
        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        if config.tie_word_embeddings:
            self.lm_head = None
        else:
            self.lm_head = nn.Linear(config.hidden_size, config.mask_vocab_size, bias=False)

    def forward(self, input_ids, timesteps, attention_mask=None):
        batch, seq = input_ids.shape
        position_ids = torch.arange(seq, device=input_ids.device).unsqueeze(0).expand(batch, -1)

        x = self.embed_tokens(input_ids)
        t_emb = self.timestep_emb(timesteps)

        attn_mask = None
        if attention_mask is not None:
            attn_mask = ((1.0 - attention_mask[:, None, None, :].float()) * -1e9).to(x.dtype)

        for layer in self.layers:
            x = layer(x, t_emb, attn_mask, position_ids)

        x = self.norm(x)
        if self.lm_head is not None:
            logits = self.lm_head(x)
        else:
            logits = F.linear(x, self.embed_tokens.weight)
        return logits

    def compute_loss(self, logits, labels, mask_positions, pad_token_id=None,
                     eos_token_id=None, eos_weight=1.0):
        """CE on masked positions only (the masked-diffusion objective).

        eos_token_id/eos_weight: boost the loss on the terminator token when
        it is a masked target, so the model learns to emit it (EOS-weighting,
        arXiv 2506.05017)."""
        logits_masked = logits[mask_positions]
        labels_masked = labels[mask_positions]
        if pad_token_id is not None:
            valid = labels_masked != pad_token_id
            logits_masked = logits_masked[valid]
            labels_masked = labels_masked[valid]
        if labels_masked.numel() == 0:
            return torch.tensor(0.0, device=logits.device), 0
        ce = F.cross_entropy(logits_masked, labels_masked, reduction="none")
        if eos_token_id is not None and eos_weight != 1.0:
            w = torch.where(labels_masked == eos_token_id, eos_weight, 1.0)
            ce = ce * w
        return ce.mean(), labels_masked.numel()

    @classmethod
    def from_checkpoint(cls, checkpoint_path, device="cpu"):
        ckpt = torch.load(checkpoint_path, map_location=device, weights_only=False)
        config_dict = ckpt.get("config", ckpt)
        config = MetaDiffusionConfig(
            **{k: v for k, v in config_dict.items() if k in MetaDiffusionConfig.__dataclass_fields__}
        )
        model = cls(config)
        sd = {k.replace("_orig_mod.", "", 1) if k.startswith("_orig_mod.") else k: v
              for k, v in ckpt.get("model_state_dict", ckpt).items()}
        model.load_state_dict(sd, strict=True)
        return model, ckpt