# Inferno — self-contained Transformers model code (trust_remote_code=True).
# Pre-RMSNorm decoder; fused QKV grouped-query attention with full-head half-split RoPE;
# packed gate/up SwiGLU (optional causal gate shift: the gate reads token t-1); tied embeddings.
import torch
from torch import nn
from torch.nn import functional as F
from transformers import PretrainedConfig, PreTrainedModel
from transformers.modeling_outputs import CausalLMOutputWithPast
try:
    from transformers.generation import GenerationMixin
except ImportError:  # older transformers
    from transformers.generation.utils import GenerationMixin


class InfernoConfig(PretrainedConfig):
    model_type = "inferno"

    def __init__(self, vocab_size=32064, hidden_size=640, num_hidden_layers=14, num_attention_heads=8,
                 num_key_value_heads=2, head_dim=80, intermediate_size=1536, rope_theta=10000.0,
                 max_position_embeddings=2048, gate_shift_start=-1, tie_word_embeddings=True, **kwargs):
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.num_key_value_heads = num_key_value_heads
        self.head_dim = head_dim
        self.intermediate_size = intermediate_size
        self.rope_theta = rope_theta
        self.max_position_embeddings = max_position_embeddings
        self.gate_shift_start = gate_shift_start   # -1 = plain SwiGLU everywhere
        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)


class RMSNorm(nn.RMSNorm):
    def forward(self, x):
        return F.rms_norm(x.float(), self.normalized_shape, self.weight, self.eps).to(x.dtype)


def rope_tables(length, dim, base):
    inv = 1. / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
    freqs = torch.outer(torch.arange(length, dtype=torch.float32), inv)
    return freqs.cos(), freqs.sin()


def apply_rope(x, cos, sin):
    work = x.float() if x.dtype in (torch.float16, torch.bfloat16) else x
    a, b = work.chunk(2, -1)
    return torch.cat((a * cos - b * sin, a * sin + b * cos), -1).to(x.dtype)


class GQAttention(nn.Module):
    def __init__(self, c):
        super().__init__()
        self.H = c.num_attention_heads; self.K = c.num_key_value_heads; self.hd = c.head_dim
        self.q_width = self.H * self.hd; self.kv_width = self.K * self.hd
        self.qkv_proj = nn.Linear(c.hidden_size, self.q_width + 2 * self.kv_width, bias=False)
        self.o_proj = nn.Linear(self.q_width, c.hidden_size, bias=False)

    def forward(self, x, cos, sin, attn_mask):
        B, T, _ = x.shape
        q, k, v = self.qkv_proj(x).split((self.q_width, self.kv_width, self.kv_width), -1)
        q = q.view(B, T, self.H, self.hd).transpose(1, 2)
        k = k.view(B, T, self.K, self.hd).transpose(1, 2)
        v = v.view(B, T, self.K, self.hd).transpose(1, 2)
        q = apply_rope(q, cos, sin); k = apply_rope(k, cos, sin)
        if attn_mask is None:
            y = F.scaled_dot_product_attention(q, k, v, is_causal=True, enable_gqa=self.H != self.K)
        else:
            y = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, enable_gqa=self.H != self.K)
        return self.o_proj(y.transpose(1, 2).reshape(B, T, self.q_width))


class FeedForward(nn.Module):
    def __init__(self, c, index):
        super().__init__()
        self.shift = c.gate_shift_start >= 0 and index >= c.gate_shift_start
        self.project = nn.Linear(c.hidden_size, 2 * c.intermediate_size, bias=False)
        self.down = nn.Linear(c.intermediate_size, c.hidden_size, bias=False)

    def forward(self, n, own_gate):
        g, u = self.project(n).chunk(2, -1)
        if self.shift:
            previous = torch.cat((g[:, :1], g[:, :-1]), 1)
            g = previous if own_gate is None else torch.where(own_gate[..., None], g, previous)
        return self.down(F.silu(g) * u)


class Block(nn.Module):
    def __init__(self, c, index):
        super().__init__()
        self.attn_norm = RMSNorm(c.hidden_size); self.attn = GQAttention(c)
        self.ffn_norm = RMSNorm(c.hidden_size); self.ffn = FeedForward(c, index)

    def forward(self, x, cos, sin, attn_mask, own_gate):
        x = x + self.attn(self.attn_norm(x), cos, sin, attn_mask)
        return x + self.ffn(self.ffn_norm(x), own_gate)


class InfernoForCausalLM(PreTrainedModel, GenerationMixin):
    config_class = InfernoConfig
    base_model_prefix = "model"
    _tied_weights_keys = {"lm_head.weight": "embed.weight"}   # {target: source}; works on Transformers 4 and 5
    _no_split_modules = ["Block"]

    def __init__(self, config):
        super().__init__(config)
        self.embed = nn.Embedding(config.vocab_size, config.hidden_size)
        self.blocks = nn.ModuleList([Block(config, i) for i in range(config.num_hidden_layers)])
        self.ln_f = RMSNorm(config.hidden_size)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        self.lm_head.weight = self.embed.weight
        self._rope = {}   # per-device RoPE tables, built like the trainer's (CPU fp32, then moved)
        self.post_init()

    def _init_weights(self, module):
        if isinstance(module, (nn.Linear, nn.Embedding)):
            nn.init.normal_(module.weight, std=.02)

    def get_input_embeddings(self): return self.embed
    def set_input_embeddings(self, value): self.embed = value
    def get_output_embeddings(self): return self.lm_head
    def set_output_embeddings(self, value): self.lm_head = value

    def _tables(self, device):
        if device not in self._rope:
            c = self.config
            cos, sin = rope_tables(c.max_position_embeddings + 1, c.head_dim, c.rope_theta)
            self._rope[device] = (cos.to(device), sin.to(device))
        return self._rope[device]

    def prepare_inputs_for_generation(self, input_ids, attention_mask=None, **kwargs):
        # No KV cache: every step re-reads the (context-clipped) sequence. Simple and exact.
        limit = self.config.max_position_embeddings
        inputs = {"input_ids": input_ids[:, -limit:]}
        if attention_mask is not None:
            inputs["attention_mask"] = attention_mask[:, -limit:]
        return inputs

    def forward(self, input_ids=None, attention_mask=None, labels=None, return_dict=None, **kwargs):
        B, T = input_ids.shape
        if T > self.config.max_position_embeddings:
            raise ValueError(f"Sequence length {T} exceeds the trained context {self.config.max_position_embeddings}")
        cos_t, sin_t = self._tables(input_ids.device)
        attn_mask = own_gate = None
        if attention_mask is None or bool(attention_mask.all()):
            cos, sin = cos_t[:T], sin_t[:T]
        else:   # padded batch: real positions restart at 0 after left padding; pads are never attended
            keep = attention_mask.bool()
            positions = (keep.long().cumsum(-1) - 1).clamp(min=0)
            cos, sin = cos_t[positions].unsqueeze(1), sin_t[positions].unsqueeze(1)
            causal = torch.ones(T, T, dtype=torch.bool, device=input_ids.device).tril()
            eye = torch.eye(T, dtype=torch.bool, device=input_ids.device)
            attn_mask = (causal[None, None] & keep[:, None, None, :]) | eye[None, None]
            own_gate = ~torch.cat((torch.zeros_like(keep[:, :1]), keep[:, :-1]), 1)
        x = self.embed(input_ids)
        for block in self.blocks:
            x = block(x, cos, sin, attn_mask, own_gate)
        logits = self.lm_head(self.ln_f(x)).float()
        loss = None
        if labels is not None:
            loss = F.cross_entropy(logits[:, :-1].reshape(-1, logits.shape[-1]), labels[:, 1:].reshape(-1), ignore_index=-100)
        return CausalLMOutputWithPast(loss=loss, logits=logits)