# coding=utf-8
# Modeling file for Oleander (hybrid GDN2 / gated-GQA causal LM).
# Auto-generated by the Oleander training notebook. The arithmetic mirrors the trainer exactly (same parameter names,
# same reference path), so the checkpoints in this repository load with
#
#     AutoModelForCausalLM.from_pretrained(repo, trust_remote_code=True)
#
# Kernel policy - both paths implement the same math:
#   * GDN2 -> flash-linear-attention `chunk_gdn2` when installed AND running on CUDA in bf16/fp16;
#             otherwise a token-by-token recurrent reference (exact, but slow for long prompts).
#   * GQA  -> torch.nn.functional.scaled_dot_product_attention (no KV cache: this repo ships `use_cache=False`).
#   * norm -> F.rms_norm when the installed torch has it.
# Set OLEANDER_REFERENCE_KERNELS=1 to force the reference path (handy when diffing against the trainer).
from __future__ import annotations

import os
from typing import Optional

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel
from transformers.modeling_outputs import CausalLMOutputWithPast

try:                                         # loaded as remote code through config.json -> auto_map
    from .configuration_oleander import OleanderConfig
except ImportError:                          # loaded as a plain module (a script, or the trainer's selftest)
    from configuration_oleander import OleanderConfig

_REF_ONLY = os.environ.get("OLEANDER_REFERENCE_KERNELS", "0") == "1"
try:
    from fla.ops.gdn2 import chunk_gdn2 as _chunk_gdn2
    _HAS_FLA = not _REF_ONLY
except Exception:
    _chunk_gdn2 = None
    _HAS_FLA = False
_HAS_FUSED_RMS = hasattr(F, "rms_norm")


def _rms(x, w, eps):
    if _HAS_FUSED_RMS:
        return F.rms_norm(x, (x.shape[-1],), w, eps)
    xf = x.float()
    y = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + eps)
    if w is not None:
        y = y * w.float()
    return y.to(x.dtype)


class RMSNorm(nn.Module):
    def __init__(self, d, eps):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(d))
        self.eps = eps

    def forward(self, x):
        return _rms(x, self.weight, self.eps)


class OffsetRMSNorm(nn.Module):
    # QK-norm in offset form: scale = 1 + w, w initialised at zero (same as the trainer).
    def __init__(self, d, eps):
        super().__init__()
        self.weight = nn.Parameter(torch.zeros(d))
        self.eps = eps

    def forward(self, x):
        return _rms(x, None, self.eps) * (1.0 + self.weight).to(x.dtype)


class GatedRMSNorm(nn.Module):
    # FLA FusedRMSNormGated(activation="sigmoid") semantics: rmsnorm(x) * w * sigmoid(g).
    def __init__(self, d, eps):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(d))
        self.eps = eps

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


class SwiGLU(nn.Module):
    def __init__(self, d, ff):
        super().__init__()
        self.gate_proj = nn.Linear(d, ff, bias=False)
        self.up_proj = nn.Linear(d, ff, bias=False)
        self.down_proj = nn.Linear(ff, d, bias=False)

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


class ResConsumer(nn.Module):
    # Learned query + RMS-norm scale for one Block-Attention-Residuals read point.
    def __init__(self, d):
        super().__init__()
        self.query = nn.Parameter(torch.zeros(d))
        self.norm_w = nn.Parameter(torch.ones(d))


def attnres_mix(query, norm_w, eps, *srcs):
    # softmax over sources of <rmsnorm(src) * norm_w, query>, weighted sum of the RAW sources (fp32).
    V = torch.stack(srcs, 0).float()
    qeff = (query * norm_w).float()
    r = torch.rsqrt(V.pow(2).mean(-1, keepdim=True) + eps)
    a = torch.softmax(((V * r) * qeff).sum(-1), dim=0)
    return (a.unsqueeze(-1) * V).sum(0)


def build_rope(max_len, rot_dim, theta):
    inv = 1.0 / (theta ** (torch.arange(0, rot_dim, 2, dtype=torch.float32) / rot_dim))
    f = torch.outer(torch.arange(max_len, dtype=torch.float32), inv)
    return f.cos(), f.sin()


def apply_rope(x, cos, sin):
    # Partial RoPE on the first 2*cos.shape[-1] dims of x [B,T,H,hd].
    T = x.shape[1]
    r = cos.shape[-1] * 2
    c = cos[:T].to(x.dtype)[None, :, None, :]
    s = sin[:T].to(x.dtype)[None, :, None, :]
    x1, x2, xp = x[..., : r // 2], x[..., r // 2: r], x[..., r:]
    return torch.cat([x1 * c - x2 * s, x2 * c + x1 * s, xp], dim=-1)


def l2norm(x, eps=1e-6):
    xf = x.float()
    return (xf * torch.rsqrt(xf.pow(2).sum(-1, keepdim=True) + eps)).to(x.dtype)


def causal_conv(x, weight):
    # Depthwise causal conv (+SiLU). x [B,T,C]; weight [C,1,K] (nn.Conv1d) or [C,K].
    T = x.shape[1]
    w = (weight[:, 0, :] if weight.dim() == 3 else weight).to(x.dtype)
    K = w.shape[-1]
    out = torch.zeros_like(x)
    for j in range(K):
        s = K - 1 - j
        xs = F.pad(x, (0, 0, s, 0))[:, :T] if s > 0 else x
        out = out + xs * w[:, j]
    return F.silu(out)


def gdn2_recurrent(q, k, v, g, b, w):
    # Token-by-token GDN2 reference (same arithmetic as fla.ops.gdn2.naive and the trainer).
    # q,k [B,T,H,K] (l2-normalised) | v,w [B,T,H,V] | g,b [B,T,H,K].
    B, T, H, K = k.shape
    V = v.shape[-1]
    qf, kf, vf, gf, bf, wf = (x.float() for x in (q, k, v, g, b, w))
    qf = qf * (K ** -0.5)
    S = torch.zeros(B, H, K, V, dtype=torch.float32, device=k.device)
    outs = []
    for t in range(T):
        S = S * gf[:, t].exp().unsqueeze(-1)
        erase = ((bf[:, t] * kf[:, t]).unsqueeze(-1) * S).sum(-2)
        S = S + kf[:, t].unsqueeze(-1) * (wf[:, t] * vf[:, t] - erase).unsqueeze(-2)
        outs.append((qf[:, t].unsqueeze(-1) * S).sum(-2))
    return torch.stack(outs, 1).to(v.dtype)


class GDN2Mixer(nn.Module):
    def __init__(self, d, cfg):
        super().__init__()
        self.H = int(cfg["heads"])
        self.hd = int(cfg["head_dim"])
        self.K = self.Vd = self.H * self.hd
        self.norm_eps = float(cfg.get("norm_eps", 1e-5))
        self.allow_neg = bool(cfg.get("allow_neg_eigval", False))
        self.gate_in_kernel = bool(cfg.get("gate_in_kernel", False))
        conv = int(cfg.get("conv_size", 4))
        K = self.K
        self.q_proj = nn.Linear(d, K, bias=False)
        self.k_proj = nn.Linear(d, K, bias=False)
        self.v_proj = nn.Linear(d, self.Vd, bias=False)
        self.q_conv1d = nn.Conv1d(K, K, conv, groups=K, padding=conv - 1, bias=False)
        self.k_conv1d = nn.Conv1d(K, K, conv, groups=K, padding=conv - 1, bias=False)
        self.v_conv1d = nn.Conv1d(self.Vd, self.Vd, conv, groups=self.Vd, padding=conv - 1, bias=False)
        self.f_proj = nn.Sequential(nn.Linear(d, self.hd, bias=False), nn.Linear(self.hd, K, bias=False))
        self.b_proj = nn.Linear(d, K, bias=False)
        self.w_proj = nn.Linear(d, self.Vd, bias=False)
        self.A_log = nn.Parameter(torch.zeros(self.H))
        self.dt_bias = nn.Parameter(torch.zeros(K))
        self.g_proj = nn.Sequential(nn.Linear(d, self.hd, bias=False), nn.Linear(self.hd, self.Vd, bias=True))
        self.o_norm = GatedRMSNorm(self.hd, self.norm_eps)
        self.o_proj = nn.Linear(self.Vd, d, bias=False)

    def forward(self, x):
        B, T, _ = x.shape
        H, hd = self.H, self.hd
        q, k, v = self.q_proj(x), self.k_proj(x), self.v_proj(x)
        f = self.f_proj(x)
        gate = self.g_proj(x)
        b = torch.sigmoid(self.b_proj(x))
        w = torch.sigmoid(self.w_proj(x))
        fast = _HAS_FLA and x.is_cuda and x.dtype in (torch.bfloat16, torch.float16)
        if fast:
            q = causal_conv(q, self.q_conv1d.weight).view(B, T, H, hd)
            k = causal_conv(k, self.k_conv1d.weight).view(B, T, H, hd)
            v = causal_conv(v, self.v_conv1d.weight).view(B, T, H, hd)
            bb = b.view(B, T, H, hd)
            if self.allow_neg:
                bb = bb * 2.0
            ww = w.view(B, T, H, hd)
            if self.gate_in_kernel:
                # the Triton kernel applies softplus + dt_bias itself
                g = f.float().view(B, T, H, hd)
                kw = dict(use_gate_in_kernel=True, A_log=self.A_log, dt_bias=self.dt_bias)
            else:
                g = -self.A_log.float().exp().view(1, 1, H, 1) * \
                    F.softplus(f.float() + self.dt_bias).view(B, T, H, hd)
                kw = {}
            o, _ = _chunk_gdn2(q=q, k=k, v=v, g=g, b=bb, w=ww, use_qk_l2norm_in_kernel=True, **kw)
        else:
            q = causal_conv(q, self.q_conv1d.weight).view(B, T, H, hd)
            k = causal_conv(k, self.k_conv1d.weight).view(B, T, H, hd)
            v = causal_conv(v, self.v_conv1d.weight).view(B, T, H, hd)
            g = -self.A_log.float().exp().view(1, 1, H, 1) * \
                F.softplus(f.float() + self.dt_bias).view(B, T, H, hd)
            bb = b.view(B, T, H, hd)
            if self.allow_neg:
                bb = bb * 2.0
            o = gdn2_recurrent(l2norm(q), l2norm(k), v, g, bb, w.view(B, T, H, hd))
        o = self.o_norm(o, gate.view(B, T, H, hd)).reshape(B, T, self.Vd)
        return self.o_proj(o)


class GatedGQA(nn.Module):
    def __init__(self, d, cfg, max_len, eps):
        super().__init__()
        self.nq = int(cfg["q_heads"])
        self.nkv = int(cfg["kv_heads"])
        self.hd = int(cfg["head_dim"])
        self.rot = int(cfg.get("rotary_dim", self.hd))
        self.theta = float(cfg.get("rope_theta", 10000.0))
        self.value_residual = bool(cfg.get("value_residual", True))
        self.qk_norm = bool(cfg.get("qk_norm", True))
        self.max_len = int(max_len)
        self.qg_proj = nn.Linear(d, 2 * self.nq * self.hd, bias=False)
        self.k_proj = nn.Linear(d, self.nkv * self.hd, bias=False)
        self.v_proj = nn.Linear(d, self.nkv * self.hd, bias=False)
        self.o_proj = nn.Linear(self.nq * self.hd, d, bias=False)
        self.q_norm = OffsetRMSNorm(self.hd, eps) if self.qk_norm else None
        self.k_norm = OffsetRMSNorm(self.hd, eps) if self.qk_norm else None
        self.v_mix = nn.Parameter(torch.tensor([0.5, 0.5]))
        # NOTE: the RoPE tables are deliberately NOT buffers. transformers builds models on the `meta` device and
        # materialises them afterwards, which leaves computed non-persistent buffers uninitialised (NaNs).
        self._rope = None
        self._rope_key = None

    def _tables(self, T, device):
        key = (int(max(T, 1)), str(device))
        if self._rope_key != key:
            cos, sin = build_rope(max(key[0], self.max_len), self.rot, self.theta)
            self._rope = (cos.to(device), sin.to(device))
            self._rope_key = key
        return self._rope

    def forward(self, x, v_first, attn_mask):
        B, T, _ = x.shape
        nq, nkv, hd = self.nq, self.nkv, self.hd
        q, gate = self.qg_proj(x).chunk(2, dim=-1)
        q = q.reshape(B, T, nq, hd)
        k = self.k_proj(x).reshape(B, T, nkv, hd)
        v = self.v_proj(x).reshape(B, T, nkv, hd)
        if self.value_residual:
            if v_first is None:
                v_first = v
            v = self.v_mix[0].to(v.dtype) * v + self.v_mix[1].to(v.dtype) * v_first
        if self.q_norm is not None:
            q, k = self.q_norm(q), self.k_norm(k)
        cos, sin = self._tables(T, x.device)
        q = apply_rope(q, cos, sin)
        k = apply_rope(k, cos, sin)
        rep = nq // nkv
        k = k.repeat_interleave(rep, dim=2)
        v = v.repeat_interleave(rep, dim=2)
        o = F.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2),
                                           attn_mask=attn_mask, is_causal=attn_mask is None)
        o = o.transpose(1, 2).reshape(B, T, nq * hd) * torch.sigmoid(gate)
        return self.o_proj(o), v_first


def _causal_mask(attention_mask, B, T, device):
    # None means "plain causal" (the fast path: is_causal=True inside SDPA).
    if attention_mask is None:
        return None
    am = attention_mask.to(device=device)
    if am.dim() != 2 or bool(am.bool().all()):
        return None
    causal = torch.ones(T, T, dtype=torch.bool, device=device).tril()
    return causal[None, None] & am[:, None, None, :].bool()


class OleanderBlock(nn.Module):
    def __init__(self, config, kind):
        super().__init__()
        d, eps = config.hidden_size, config.norm_eps
        self.kind = kind
        self.norm1, self.norm2 = RMSNorm(d, eps), RMSNorm(d, eps)
        if kind == "gdn2":
            self.mixer = GDN2Mixer(d, config.gdn2)
        elif kind == "gqa":
            self.mixer = GatedGQA(d, config.gqa, config.max_seq_len, eps)
        else:
            raise ValueError("unknown layer kind: %r" % (kind,))
        self.mlp = SwiGLU(d, config.mlp_hidden)
        self.attn_res_mixer = ResConsumer(d)
        self.attn_res_mlp = ResConsumer(d)


class OleanderPreTrainedModel(PreTrainedModel):
    config_class = OleanderConfig
    base_model_prefix = ""
    main_input_name = "input_ids"
    _no_split_modules = ["OleanderBlock"]
    supports_gradient_checkpointing = False
    _supports_sdpa = True
    _supports_cache_class = False
    _supports_static_cache = False

    def _init_weights(self, module):
        std = float(getattr(self.config, "init_std", 0.02))
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=std)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=std)

    def tie_weights(self, *args, **kwargs):
        # The LM head *is* the embedding matrix - nothing to tie.
        return None


# transformers >= 4.50 wants an explicit GenerationMixin; older versions already inherit it, in which case the tuple
# stays a single base (listing it twice would be an MRO error).
_OLEANDER_BASES = (OleanderPreTrainedModel,)
try:
    from transformers import GenerationMixin as _GenerationMixin

    if _GenerationMixin not in OleanderPreTrainedModel.__mro__:
        _OLEANDER_BASES = (OleanderPreTrainedModel, _GenerationMixin)
except Exception:  # noqa: BLE001
    pass


class OleanderForCausalLM(*_OLEANDER_BASES):
    """Causal LM. The trunk sits directly on this module (no `model.` submodule), so the parameter names are exactly
    the trainer's - `embed.weight`, `layers.N.mixer.*`, `final_norm.weight` - and the training checkpoints load with
    `strict=True`. The LM head is the tied embedding matrix.
    """

    base_model_prefix = ""

    def __init__(self, config):
        super().__init__(config)
        d = config.hidden_size
        self.group_size = len(config.group_pattern)
        self.eps = config.norm_eps
        self.use_attn_res = bool(config.attn_res)
        self.embed = nn.Embedding(config.vocab_size, d)
        kinds = getattr(config, "oleander_layer_types", None) or config.layer_types
        self.layers = nn.ModuleList([OleanderBlock(config, kinds[i]) for i in range(config.num_hidden_layers)])
        self.res_out = ResConsumer(d)
        self.final_norm = RMSNorm(d, self.eps)
        self.post_init()

    # ---- embeddings / head -----------------------------------------------------------------------------------
    def get_input_embeddings(self):
        return self.embed

    def set_input_embeddings(self, value):
        self.embed = value

    def get_output_embeddings(self):
        return None

    @property
    def head_weight(self):
        return self.embed.weight

    # ---- trunk -----------------------------------------------------------------------------------------------
    def _mix(self, cons, srcs):
        if len(srcs) == 1:
            return srcs[0]
        if not self.use_attn_res:
            out = None
            for s in srcs:
                out = s.float() if out is None else out + s.float()
            return out
        return attnres_mix(cons.query, cons.norm_w, self.eps, *srcs)

    def _trunk(self, input_ids, attention_mask=None):
        """Returns the final-normed hidden states in fp32."""
        x = self.embed(input_ids)
        B, T, _ = x.shape
        attn_mask = _causal_mask(attention_mask, B, T, x.device)
        sources, partial, v_first = [x], None, None
        for i, blk in enumerate(self.layers):
            h = self._mix(blk.attn_res_mixer, sources + ([partial] if partial is not None else []))
            if i > 0 and i % self.group_size == 0:
                sources.append(partial)
                partial = None
            hn = blk.norm1(h)
            if blk.kind == "gdn2":
                out = blk.mixer(hn)
            else:
                out, v_first = blk.mixer(hn, v_first, attn_mask)
            out = out.float()
            partial = out if partial is None else partial + out
            h = self._mix(blk.attn_res_mlp, sources + [partial])
            partial = partial + blk.mlp(blk.norm2(h)).float()
        h = self._mix(self.res_out, sources + [partial])
        return self.final_norm(h)

    # ---- head / loss -----------------------------------------------------------------------------------------
    def _logits(self, h, keep=0, chunk=8192):
        if keep and keep > 0:
            h = h[:, -int(keep):]
        B, T, D = h.shape
        hf = h.reshape(-1, D)
        parts = [F.linear(hf[s:s + chunk], self.head_weight) for s in range(0, hf.shape[0], chunk)]
        logits = parts[0] if len(parts) == 1 else torch.cat(parts, 0)
        return logits.view(B, T, -1).float()

    def _loss(self, h, labels, chunk=8192):
        if labels.dim() == 3:
            labels = labels[..., 0]
        B, T, D = h.shape
        hf = h.reshape(-1, D)
        lf = labels.reshape(-1)
        total = hf.new_zeros((), dtype=torch.float32)
        denom = (lf >= 0).sum().clamp(min=1).float()
        for s in range(0, hf.shape[0], chunk):
            lg = F.linear(hf[s:s + chunk], self.head_weight).float()
            total = total + F.cross_entropy(lg, lf[s:s + chunk], reduction="sum", ignore_index=-100)
        return total / denom

    def forward(self, input_ids=None, attention_mask=None, position_ids=None, past_key_values=None,
                inputs_embeds=None, labels=None, use_cache=None, output_attentions=None,
                output_hidden_states=None, return_dict=None, logits_to_keep=0, num_logits_to_keep=0, **kwargs):
        # No KV cache: the GDN2 mixers are linear-attention and this repo ships `use_cache=False`.
        return_dict = True if return_dict is None else return_dict
        h = self._trunk(input_ids, attention_mask=attention_mask)     # fp32 [B, T, D]
        loss = self._loss(h, labels) if labels is not None else None
        keep = int(logits_to_keep or num_logits_to_keep or 0)
        logits = self._logits(h, keep=keep)
        if not return_dict:
            return (logits,) + ((loss,) if loss is not None else ())
        return CausalLMOutputWithPast(loss=loss, logits=logits)

    def prepare_inputs_for_generation(self, input_ids, attention_mask=None, **kwargs):
        return {"input_ids": input_ids, "attention_mask": attention_mask}