"""Trainable PyTorch port of NVIDIA Dory1.0 (looped normalized-GPT).

Mirrors vLLM's inference implementation (vllm-project/vllm#60149,
vllm/model_executor/models/dory.py) op for op, replacing its paged attention and
tensor-parallel layers with plain modules so the model can be fine-tuned:

  * full-sequence causal attention with per-layer sliding windows; each recurrent
    loop attends over that loop's own keys/values, which is exactly vLLM's
    "per_loop" KV-cache mode
  * loop count selectable per forward pass (n_loops), for elastic-loop training
  * per-layer gradient checkpointing and a chunked cross-entropy over the 131k vocab
  * packed sequences (doc_ids + per-document position_ids); attention via dense-mask
    SDPA (any device, used for checks) or FlexAttention block masks (CUDA training)

Parameter names match the released checkpoint, so state dicts load and save as-is.
"""

import json
import math
from dataclasses import dataclass, fields
from pathlib import Path

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

EPS = 1e-6


@dataclass
class DoryConfig:
    vocab_size: int = 131072
    hidden_size: int = 2560
    num_hidden_layers: int = 72
    intermediate_size: int = 7168
    num_attention_heads: int = 16
    num_key_value_heads: int = 4
    head_dim: int = 256
    initializer_range: float = 0.0177
    sliding_window: int = 1024
    rope_parameters: dict = None
    layer_types: list = None
    ngpt_alpha_init: float = 0.05
    attention_output_gate_logit_scale_init: float = None
    n_input_layers: int = 16
    n_recurrent_layers: int = 40
    n_output_layers: int = 16
    n_recurrent_loops: int = 2

    @classmethod
    def from_json(cls, path):
        raw = json.loads(Path(path).read_text())
        known = {f.name for f in fields(cls)}
        cfg = cls(**{k: v for k, v in raw.items() if k in known})
        assert cfg.n_input_layers + cfg.n_recurrent_layers + cfg.n_output_layers == cfg.num_hidden_layers
        assert not raw.get("tie_word_embeddings") and raw.get("hidden_act", "silu") == "silu"
        assert not raw.get("attention_bias") and not raw.get("mlp_bias")
        return cfg


def justnorm(x, dim=-1):
    """nGPT normalization: x * rsqrt(sum(x^2) + eps), fp32 accumulation."""
    xf = x.float()
    return (xf * torch.rsqrt(xf.pow(2).sum(dim=dim, keepdim=True) + EPS)).to(x.dtype)


def slerp_alpha(a, b, alpha, alpha_scale):
    ya, yb = justnorm(a).float(), justnorm(b).float()
    lr = (alpha.float() * alpha_scale).abs()
    return justnorm(ya + lr * (yb - ya)).to(a.dtype)


def rope_cos_sin(positions, head_dim, theta, device):
    inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device, dtype=torch.float32) / head_dim))
    freqs = positions.to(device=device, dtype=torch.float32)[..., None] * inv_freq
    return freqs.cos(), freqs.sin()  # [B, T, D/2]


def apply_rope_neox(x, cos, sin):
    """x: [B, T, H, D]; cos/sin: [B, T, D/2]. NeoX style (rotate halves)."""
    x1, x2 = x.float().chunk(2, dim=-1)
    c, s = cos[:, :, None, :], sin[:, :, None, :]
    return torch.cat([x1 * c - x2 * s, x2 * c + x1 * s], dim=-1).to(x.dtype)


def causal_mask(t, window, device, doc_ids=None):
    """Bool mask, True = attend: [T, T], or [B, 1, T, T] with doc_ids (packed sequences).
    Window w keeps keys with 0 <= i - j < w."""
    i = torch.arange(t, device=device)[:, None]
    j = torch.arange(t, device=device)[None, :]
    allowed = j <= i
    if window is not None:
        allowed &= (i - j) < window
    if doc_ids is not None:
        allowed = allowed[None] & (doc_ids[:, :, None] == doc_ids[:, None, :])
        allowed = allowed[:, None]
    return allowed


_flex = None


def flex_block_mask(t, window, device, doc_ids):
    from torch.nn.attention.flex_attention import create_block_mask

    def mask_mod(b, h, q, kv):
        ok = (kv <= q) & (doc_ids[b, q] == doc_ids[b, kv])
        if window is not None:
            ok = ok & ((q - kv) < window)
        return ok

    return create_block_mask(mask_mod, doc_ids.shape[0], None, t, t, device=device)


def flex_attention(q, k, v, block_mask, scale):
    global _flex
    if _flex is None:
        from torch.nn.attention.flex_attention import flex_attention as fa

        _flex = torch.compile(fa, dynamic=False)
    return _flex(q, k, v, block_mask=block_mask, scale=scale, enable_gqa=True)


class DoryAttention(nn.Module):
    """GQA attention with nGPT qk-norm (after RoPE), sigmoid output gate and justnorm."""

    def __init__(self, cfg: DoryConfig, layer_type: str):
        super().__init__()
        h, kv, d = cfg.num_attention_heads, cfg.num_key_value_heads, cfg.head_dim
        self.num_heads, self.num_kv_heads, self.head_dim = h, kv, d
        self.q_proj = nn.Linear(cfg.hidden_size, h * d, bias=False)
        self.k_proj = nn.Linear(cfg.hidden_size, kv * d, bias=False)
        self.v_proj = nn.Linear(cfg.hidden_size, kv * d, bias=False)
        self.g_proj = nn.Linear(cfg.hidden_size, h * d, bias=False)
        self.o_proj = nn.Linear(h * d, cfg.hidden_size, bias=False)
        self.sqk = nn.Parameter(torch.empty(h * d))
        self.s_gate = nn.Parameter(torch.empty(h * d))
        std = cfg.initializer_range
        self.sqk_factor = d**0.5 / std
        g0 = cfg.attention_output_gate_logit_scale_init or cfg.hidden_size**0.5
        self.s_gate_factor = g0**0.5 / std
        self.layer_type = layer_type

    def forward(self, x, cos, sin, mask):
        b, t, _ = x.shape
        h, kv, d = self.num_heads, self.num_kv_heads, self.head_dim
        q = self.q_proj(x).view(b, t, h, d)
        k = self.k_proj(x).view(b, t, kv, d)
        v = self.v_proj(x).view(b, t, kv, d)
        gate = self.g_proj(x)

        q, k = apply_rope_neox(q, cos, sin), apply_rope_neox(k, cos, sin)
        sqk = (self.sqk.float() * self.sqk_factor).view(h, d)
        q = (sqk.square() * justnorm(q).float()).to(q.dtype)
        k = justnorm(k)

        q, k, v = (z.transpose(1, 2) for z in (q, k, v))  # [B, H, T, D]
        if isinstance(mask, torch.Tensor):
            if kv != h:
                k = k.repeat_interleave(h // kv, dim=1)
                v = v.repeat_interleave(h // kv, dim=1)
            out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, scale=1.0 / math.sqrt(d))
        else:  # FlexAttention BlockMask
            out = flex_attention(q, k, v, mask, 1.0 / math.sqrt(d))
        out = out.transpose(1, 2).reshape(b, t, h * d)

        sg = (self.s_gate.float() * self.s_gate_factor).square()
        out = (out.float() * torch.sigmoid(gate.float() * sg)).to(out.dtype)
        return self.o_proj(justnorm(out))


class DoryMLP(nn.Module):
    """SwiGLU with nGPT suv scaling and justnorm before down_proj."""

    def __init__(self, cfg: DoryConfig):
        super().__init__()
        self.gate_proj = nn.Linear(cfg.hidden_size, cfg.intermediate_size, bias=False)
        self.up_proj = nn.Linear(cfg.hidden_size, cfg.intermediate_size, bias=False)
        self.down_proj = nn.Linear(cfg.intermediate_size, cfg.hidden_size, bias=False)
        self.suv_gate = nn.Parameter(torch.empty(cfg.intermediate_size))
        self.suv_up = nn.Parameter(torch.empty(cfg.intermediate_size))
        self.suv_factor = cfg.hidden_size**0.5

    def forward(self, x):
        g = (self.gate_proj(x).float() * (self.suv_gate.float() * self.suv_factor)).to(x.dtype)
        u = (self.up_proj(x).float() * (self.suv_up.float() * self.suv_factor)).to(x.dtype)
        m = (F.silu(g.float()) * u.float()).to(x.dtype)
        return self.down_proj(justnorm(m))


class DoryLayer(nn.Module):
    """One physical layer: a single sublayer (attention OR mlp) plus slerp residual."""

    def __init__(self, cfg: DoryConfig, idx: int):
        super().__init__()
        self.layer_type = cfg.layer_types[idx]
        self.is_attention = self.layer_type != "mlp"
        self.alpha_scale = cfg.ngpt_alpha_init / cfg.initializer_range
        if self.is_attention:
            self.mixer = DoryAttention(cfg, self.layer_type)
            self.attn_alpha = nn.Parameter(torch.empty(cfg.hidden_size))
        else:
            self.mixer = DoryMLP(cfg)
            self.mlp_alpha = nn.Parameter(torch.empty(cfg.hidden_size))

    def forward(self, h, ctx):
        if self.is_attention:
            cos, sin = ctx["rope"][self.layer_type]
            out = self.mixer(h, cos, sin, ctx["mask"][self.layer_type])
            return slerp_alpha(h, out, self.attn_alpha, self.alpha_scale)
        return slerp_alpha(h, self.mixer(h), self.mlp_alpha, self.alpha_scale)


class DoryBackbone(nn.Module):
    def __init__(self, cfg: DoryConfig):
        super().__init__()
        self.cfg = cfg
        self.embeddings = nn.Embedding(cfg.vocab_size, cfg.hidden_size)
        self.layers = nn.ModuleList(DoryLayer(cfg, i) for i in range(cfg.num_hidden_layers))
        self.gradient_checkpointing = False

    def _run(self, layer, h, ctx):
        if self.gradient_checkpointing and self.training and torch.is_grad_enabled():
            return checkpoint(layer, h, ctx, use_reentrant=False)
        return layer(h, ctx)

    def forward(self, input_ids, position_ids=None, n_loops=None, doc_ids=None, attn_impl="sdpa"):
        cfg = self.cfg
        b, t = input_ids.shape
        dev = input_ids.device
        if position_ids is None:
            position_ids = torch.arange(t, device=dev).expand(b, t)
        if attn_impl == "flex" and doc_ids is None:
            doc_ids = torch.zeros(b, t, dtype=torch.int32, device=dev)
        ctx = {"rope": {}, "mask": {}}
        for lt in set(cfg.layer_types) - {"mlp"}:
            ctx["rope"][lt] = rope_cos_sin(position_ids, cfg.head_dim, cfg.rope_parameters[lt]["rope_theta"], dev)
            window = cfg.sliding_window if lt == "sliding_attention" else None
            if attn_impl == "flex":
                ctx["mask"][lt] = flex_block_mask(t, window, dev, doc_ids)
            else:
                ctx["mask"][lt] = causal_mask(t, window, dev, doc_ids)

        n_in, n_rec = cfg.n_input_layers, cfg.n_recurrent_layers
        h = self.embeddings(input_ids)
        for layer in self.layers[:n_in]:
            h = self._run(layer, h, ctx)
        input_h = h
        state = input_h
        for _ in range(n_loops or cfg.n_recurrent_loops):
            h = justnorm(state + input_h)
            for layer in self.layers[n_in : n_in + n_rec]:
                h = self._run(layer, h, ctx)
            state = justnorm(h)
        h = state
        for layer in self.layers[n_in + n_rec :]:
            h = self._run(layer, h, ctx)
        return h


class DoryForCausalLM(nn.Module):
    def __init__(self, cfg: DoryConfig):
        super().__init__()
        self.config = cfg
        self.backbone = DoryBackbone(cfg)
        self.lm_head = nn.Linear(cfg.hidden_size, cfg.vocab_size, bias=False)
        self.lm_head.logit_scale = nn.Parameter(torch.empty(cfg.vocab_size))

    @classmethod
    def from_pretrained(cls, path, dtype=torch.bfloat16, device="cpu"):
        path = Path(path)
        cfg = DoryConfig.from_json(path / "config.json")
        with torch.device("meta"):
            model = cls(cfg)
        state = {}
        for shard in sorted(path.glob("*.safetensors")):
            state.update(load_file(shard))
        state = {k: v.to(dtype) for k, v in state.items()}
        missing, unexpected = model.load_state_dict(state, strict=True, assign=True)
        assert not missing and not unexpected, (missing, unexpected)
        return model.to(device)

    def gradient_checkpointing_enable(self):
        self.backbone.gradient_checkpointing = True

    @torch.no_grad()
    def normalize_matrices_(self):
        """Project weights back onto the unit hypersphere after an optimizer step (nGPT).

        The released checkpoint has every matrix unit-norm along the hidden dimension:
        rows of embeddings/lm_head and of maps out of the residual stream (q/k/v/g,
        gate/up), columns of maps into it (o_proj, down_proj).
        """
        hidden = self.config.hidden_size
        for w in (self.backbone.embeddings.weight, self.lm_head.weight):
            w.copy_(justnorm(w, dim=1))
        for mod in self.backbone.modules():
            if isinstance(mod, nn.Linear):
                dim = 1 if mod.in_features == hidden else 0
                mod.weight.copy_(justnorm(mod.weight, dim=dim))

    def logits(self, hidden):
        scale = self.lm_head.logit_scale.float() / self.config.initializer_range
        return self.lm_head(hidden).float() * scale

    def forward(self, input_ids, labels=None, position_ids=None, n_loops=None, doc_ids=None,
                attn_impl="sdpa", loss_chunk=2048):
        """Without labels: returns float32 logits [B, T, V].
        With labels (same shape as input_ids, -100 = ignore; shifted internally, so a
        packed batch must set -100 on each document's first token):
        returns (mean_loss, n_target_tokens) via a chunked, checkpointed CE."""
        hidden = self.backbone(input_ids, position_ids=position_ids, n_loops=n_loops,
                               doc_ids=doc_ids, attn_impl=attn_impl)
        if labels is None:
            return self.logits(hidden)
        h = hidden[:, :-1].reshape(-1, hidden.shape[-1])
        y = labels[:, 1:].reshape(-1)
        keep = y != -100
        h, y = h[keep], y[keep]
        total = h.new_zeros((), dtype=torch.float32)

        def chunk_loss(hc, yc):
            return F.cross_entropy(self.logits(hc), yc, reduction="sum")

        for s in range(0, h.shape[0], loss_chunk):
            hc, yc = h[s : s + loss_chunk], y[s : s + loss_chunk]
            if self.training and torch.is_grad_enabled():
                total = total + checkpoint(chunk_loss, hc, yc, use_reentrant=False)
            else:
                total = total + chunk_loss(hc, yc)
        n = int(keep.sum())
        return total / max(n, 1), n

    @torch.no_grad()
    def generate(self, input_ids, max_new_tokens=128, n_loops=None, temperature=0.0, eos_ids=()):
        """Simple cache-free sampling, for correctness checks (not speed)."""
        ids = input_ids
        for _ in range(max_new_tokens):
            logits = self.forward(ids, n_loops=n_loops)[:, -1]
            if temperature > 0:
                nxt = torch.multinomial(torch.softmax(logits / temperature, -1), 1)
            else:
                nxt = logits.argmax(-1, keepdim=True)
            ids = torch.cat([ids, nxt], dim=1)
            if int(nxt[0, 0]) in eos_ids:
                break
        return ids[:, input_ids.shape[1] :]