"""BF16 Dory with frozen checkpoint weights and trainable linear LoRA adapters.

The nGPT equations follow the independent dory-checkpoint-reference.py oracle
and vLLM revision 55076950b6c155f384d6766d63563c3bc20a0e3f. This is a training
experiment, not a replacement for Eider's CUDA serving implementation.
"""

from __future__ import annotations

import json
import math
from pathlib import Path

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


def norm(x):
    value = x.float()
    return (value * torch.rsqrt(value.square().sum(-1, keepdim=True) + 1e-6)).to(
        x.dtype
    )


class Linear(nn.Module):
    def __init__(self, weight, rank, alpha):
        super().__init__()
        self.register_buffer("weight", weight)
        self.scale = alpha / rank if rank else 0
        if rank:
            self.lora_a = nn.Parameter(
                torch.empty(rank, weight.shape[1], device=weight.device)
            )
            self.lora_b = nn.Parameter(
                torch.zeros(weight.shape[0], rank, device=weight.device)
            )
            nn.init.kaiming_uniform_(self.lora_a, a=math.sqrt(5))
        else:
            self.register_parameter("lora_a", None)
            self.register_parameter("lora_b", None)

    def forward(self, x):
        base = F.linear(x, self.weight)
        if self.lora_a is None:
            return base
        delta = F.linear(F.linear(x, self.lora_a.to(x.dtype)), self.lora_b.to(x.dtype))
        return (base.float() + self.scale * delta.float()).to(x.dtype)


class Layer(nn.Module):
    def __init__(self, config, index, read, rank, alpha):
        super().__init__()
        self.config = config
        self.kind = config["layer_types"][index]
        prefix = f"backbone.layers.{index}"
        mixer = prefix + ".mixer"
        std = config["initializer_range"]
        if self.kind == "mlp":
            names = ("gate_proj", "up_proj", "down_proj")
            scale = math.sqrt(config["hidden_size"])
            self.register_buffer("suv_gate", read(mixer + ".suv_gate").float() * scale)
            self.register_buffer("suv_up", read(mixer + ".suv_up").float() * scale)
            rate = read(prefix + ".mlp_alpha")
        else:
            names = ("q_proj", "k_proj", "v_proj", "g_proj", "o_proj")
            dim = config["head_dim"]
            self.register_buffer(
                "sqk", (read(mixer + ".sqk").float() * (math.sqrt(dim) / std)).square()
            )
            g0 = config.get("attention_output_gate_logit_scale_init") or math.sqrt(
                config["hidden_size"]
            )
            self.register_buffer(
                "s_gate",
                (read(mixer + ".s_gate").float() * (math.sqrt(g0) / std)).square(),
            )
            rate = read(prefix + ".attn_alpha")
        self.register_buffer(
            "rate", (rate.float() * config["ngpt_alpha_init"] / std).abs()
        )
        self.projections = nn.ModuleDict(
            {
                name: Linear(read(mixer + "." + name + ".weight"), rank, alpha)
                for name in names
            }
        )

    def forward(self, h, rotary, sliding_mask, reference=False):
        p, c, dtype = self.projections, self.config, h.dtype
        if self.kind == "mlp":
            gate = (p["gate_proj"](h).float() * self.suv_gate).to(dtype)
            up = (p["up_proj"](h).float() * self.suv_up).to(dtype)
            out = p["down_proj"](norm((F.silu(gate.float()) * up.float()).to(dtype)))
        else:
            nq, nk, dim = (
                c["num_attention_heads"],
                c["num_key_value_heads"],
                c["head_dim"],
            )
            cos, sin = rotary[self.kind]

            def rope(x):
                a, b = x.float().chunk(2, -1)
                return torch.cat((a * cos - b * sin, b * cos + a * sin), -1).to(dtype)

            q = norm(rope(p["q_proj"](h).reshape(*h.shape[:-1], nq, dim)))
            k = norm(rope(p["k_proj"](h).reshape(*h.shape[:-1], nk, dim)))
            q = (q.float() * self.sqk.reshape(nq, dim)).to(dtype)
            v = p["v_proj"](h).reshape(*h.shape[:-1], nk, dim)
            # Repeat GQA heads explicitly so both SDPA backends and the oracle
            # use the same layout. Batch rows retain independent attention.
            q = q.transpose(1, 2)
            k = k.repeat_interleave(nq // nk, 2).transpose(1, 2)
            v = v.repeat_interleave(nq // nk, 2).transpose(1, 2)
            mask = sliding_mask if self.kind == "sliding_attention" else None
            if reference:
                if mask is None:
                    mask = torch.ones(
                        h.shape[1], h.shape[1], device=h.device, dtype=torch.bool
                    ).tril()
                scores = (q.float() @ k.float().transpose(-1, -2)) / math.sqrt(dim)
                out = (
                    scores.masked_fill(~mask, -torch.inf).softmax(-1) @ v.float()
                ).to(dtype)
            else:
                out = F.scaled_dot_product_attention(
                    q, k, v, attn_mask=mask, is_causal=mask is None
                )
            out = out.transpose(1, 2).reshape(*h.shape[:-1], nq * dim)
            gate = p["g_proj"](h)
            out = norm(
                (out.float() * torch.sigmoid(gate.float() * self.s_gate)).to(dtype)
            )
            out = p["o_proj"](out)
        a, b = norm(h).float(), norm(out).float()
        return norm(a + self.rate * (b - a)).to(dtype)


class Dory(nn.Module):
    def __init__(self, config, read, *, rank=8, alpha=16, loops=2, checkpointing=False):
        super().__init__()
        if config.get("model_type") != "dory" or config.get("quantization_config"):
            raise ValueError("requires an unquantized Dory checkpoint")
        if loops < 1 or rank < 0 or alpha <= 0:
            raise ValueError("invalid loop count or adapter dimensions")
        self.config, self.loops, self.checkpointing = config, loops, checkpointing
        self.register_buffer("embedding", read("backbone.embeddings.weight"))
        self.register_buffer("head", read("lm_head.weight"))
        self.register_buffer(
            "head_scale",
            read("lm_head.logit_scale").float() / config["initializer_range"],
        )
        self.layers = nn.ModuleList(
            Layer(config, i, read, rank, alpha)
            for i in range(config["num_hidden_layers"])
        )

    @classmethod
    def load(cls, path: Path, *, device="cuda", **kwargs):
        config = json.loads((path / "config.json").read_text())
        with safe_open(
            path / "model.safetensors", framework="pt", device="cpu"
        ) as weights:
            consumed = set()

            def read(name):
                value = weights.get_tensor(name)
                if value.dtype != torch.bfloat16:
                    raise ValueError(f"expected BF16 weight: {name}")
                consumed.add(name)
                return value.to(device)

            model = cls(config, read, **kwargs)
            if consumed != set(weights.keys()):
                raise ValueError(
                    f"unconsumed checkpoint tensors: {set(weights.keys()) - consumed}"
                )
        return model

    def forward(self, ids, labels=None, *, lengths=None, reference=False):
        c = self.config
        single = ids.ndim == 1
        if single:
            ids = ids.unsqueeze(0)
        positions = torch.arange(ids.shape[1], device=ids.device)
        dim = c["head_dim"]
        rotary = {}
        for kind, params in c["rope_parameters"].items():
            angles = positions.float()[:, None] * params["rope_theta"] ** (
                -torch.arange(dim // 2, device=ids.device).float() / (dim // 2)
            )
            rotary[kind] = (angles.cos()[:, None, :], angles.sin()[:, None, :])
        # For short training rows the sliding mask equals the causal mask.
        sliding_mask = None
        if ids.shape[1] > c["sliding_window"]:
            distance = positions[:, None] - positions[None, :]
            sliding_mask = (distance >= 0) & (distance < c["sliding_window"])

        def apply(h, start, end):
            for layer in self.layers[start:end]:
                if self.training and self.checkpointing and not reference:
                    h = checkpoint(layer, h, rotary, sliding_mask, use_reentrant=False)
                else:
                    # Keep the serial evaluator and oracle eager. Their batch-one,
                    # inference-only graphs otherwise compete with training graphs.
                    forward = (
                        layer if self.training and not reference else layer.forward
                    )
                    h = forward(h, rotary, sliding_mask, reference)
            return h

        ni, nr = c["n_input_layers"], c["n_recurrent_layers"]
        h = apply(F.embedding(ids, self.embedding), 0, ni)
        input_h = h
        for _ in range(self.loops):
            h = norm(h + input_h)
            h = norm(apply(h, ni, ni + nr))
        h = apply(h, ni + nr, len(self.layers))
        # Right-padding is always after each real query, so causal attention
        # excludes it without a padding mask. Read only each row's final token.
        last = (
            h[:, -1]
            if lengths is None
            else h[torch.arange(ids.shape[0], device=ids.device), lengths - 1]
        )
        weight = self.head if labels is None else self.head[labels]
        scale = self.head_scale if labels is None else self.head_scale[labels]
        logits = (
            torch.bmm(weight, last.unsqueeze(-1)).squeeze(-1)
            if labels is not None and labels.ndim == 2
            else F.linear(last, weight)
        ).float() * scale
        return logits.squeeze(0) if single else logits

    def compile_layers(self):
        # Compile individual layers so recurrent traversals reuse their graphs.
        # Keep the BF16 rounding points used by the eager nGPT equations.
        # All layer kinds share Layer.forward's code cache. Attention masks,
        # input gradients and shape guards produce more than eight variants when
        # training batches cross the sliding-window boundary. Keep a bounded
        # cache large enough for those variants; fullgraph raises on overflow.
        torch._dynamo.config.recompile_limit = max(
            torch._dynamo.config.recompile_limit, 32
        )
        for layer in self.layers:
            layer.compile(
                fullgraph=True,
                dynamic=True,
                options={"emulate_precision_casts": True},
            )

    def save_adapter(self, path):
        save_file(
            {
                name: p.detach().cpu().contiguous()
                for name, p in self.named_parameters()
            },
            str(path),
        )

    def load_adapter(self, path):
        values = load_file(str(path))
        params = dict(self.named_parameters())
        if values.keys() != params.keys():
            raise ValueError("adapter does not match the model's projection modules")
        with torch.no_grad():
            for name, p in params.items():
                if p.shape != values[name].shape:
                    raise ValueError(f"adapter shape mismatch: {name}")
                p.copy_(values[name])