# Model definition for Panhapich/pre-train-llama.
#
# Llama-style decoder-only transformer (RoPE, grouped-query attention, SwiGLU,
# RMSNorm), trained from scratch. This is not a `transformers`-library model
# class -- load model.safetensors into TextGenerationModel directly:
#
#     import json
#     from safetensors.torch import load_file
#     from modeling_llama_custom import TextGenerationModel
#
#     config = json.load(open("config.json"))
#     model = TextGenerationModel(**config["model_config"])
#     model.load_state_dict(load_file("model.safetensors"))
#     model.eval()

import math

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

if not hasattr(nn, "RMSNorm"):
    class _RMSNormFallback(nn.Module):
        def __init__(self, dim, eps=1e-6):
            super().__init__()
            self.eps = eps
            self.weight = nn.Parameter(torch.ones(dim))

        def forward(self, x):
            rms = x.pow(2).mean(dim=-1, keepdim=True).add(self.eps).rsqrt()
            return x * rms * self.weight

    nn.RMSNorm = _RMSNormFallback


class RotaryPositionalEncoding(nn.Module):
    def __init__(self, head_dim, max_seq_len, theta=10000.0):
        super().__init__()
        inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2).float() / head_dim))
        t = torch.arange(max_seq_len).float()
        freqs = torch.outer(t, inv_freq)
        self.register_buffer("cos", torch.cos(freqs), persistent=False)
        self.register_buffer("sin", torch.sin(freqs), persistent=False)

    def rotate(self, x):
        T = x.shape[-2]
        cos = self.cos[:T].unsqueeze(0).unsqueeze(0)
        sin = self.sin[:T].unsqueeze(0).unsqueeze(0)
        x1, x2 = x[..., 0::2], x[..., 1::2]
        rotated = torch.stack([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
        return rotated.flatten(-2)


class GQA(nn.Module):
    def __init__(self, hidden_dim, num_heads, num_kv_heads, dropout=0.1):
        super().__init__()
        self.num_heads = num_heads
        self.num_kv_heads = num_kv_heads
        self.n_rep = num_heads // num_kv_heads
        self.head_dim = hidden_dim // num_heads

        self.q_proj = nn.Linear(hidden_dim, num_heads * self.head_dim)
        self.k_proj = nn.Linear(hidden_dim, num_kv_heads * self.head_dim)
        self.v_proj = nn.Linear(hidden_dim, num_kv_heads * self.head_dim)
        self.out_proj = nn.Linear(num_heads * self.head_dim, hidden_dim)
        self.dropout = nn.Dropout(dropout)

    def forward(self, q, k, v, mask=None, rope=None):
        B, T, _ = q.shape
        q = self.q_proj(q).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(k).view(B, T, self.num_kv_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(v).view(B, T, self.num_kv_heads, self.head_dim).transpose(1, 2)

        if rope is not None:
            q = rope.rotate(q)
            k = rope.rotate(k)
        if self.n_rep > 1:
            k = k.repeat_interleave(self.n_rep, dim=1)
            v = v.repeat_interleave(self.n_rep, dim=1)

        scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
        if mask is not None:
            scores = scores.masked_fill(~mask.unsqueeze(1).bool(), float('-inf'))
        attn = F.softmax(scores, dim=-1)
        attn = self.dropout(attn)
        out = attn @ v
        out = out.transpose(1, 2).reshape(B, T, -1)
        return self.out_proj(out)


class SwiGLU(nn.Module):
    def __init__(self, hidden_dim, ff_dim):
        super().__init__()
        self.gate_proj = nn.Linear(hidden_dim, ff_dim)
        self.up_proj = nn.Linear(hidden_dim, ff_dim)
        self.down_proj = nn.Linear(ff_dim, hidden_dim)

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


class DecoderLayer(nn.Module):
    def __init__(self, hidden_dim, num_heads, num_kv_heads, dropout=0.1):
        super().__init__()
        self.self_attn = GQA(hidden_dim, num_heads, num_kv_heads, dropout)
        self.mlp = SwiGLU(hidden_dim, 4 * hidden_dim)
        self.norm1 = nn.RMSNorm(hidden_dim)
        self.norm2 = nn.RMSNorm(hidden_dim)

    def forward(self, x, mask=None, rope=None):
        out = self.norm1(x)
        out = self.self_attn(out, out, out, mask, rope)
        x = out + x
        out = self.norm2(x)
        out = self.mlp(out)
        return out + x

class TextGenerationModel(nn.Module):
    def __init__(self, num_layers, num_heads, num_kv_heads, hidden_dim,
                 max_seq_len, vocab_size, dropout=0.1):
        super().__init__()
        self.rope = RotaryPositionalEncoding(hidden_dim // num_heads, max_seq_len)
        self.embedding = nn.Embedding(vocab_size, hidden_dim)
        self.decoders = nn.ModuleList([
            DecoderLayer(hidden_dim, num_heads, num_kv_heads, dropout)
            for _ in range(num_layers)
        ])
        self.norm = nn.RMSNorm(hidden_dim)
        self.out = nn.Linear(hidden_dim, vocab_size)

    def forward(self, ids, mask=None):
        x = self.embedding(ids)
        for decoder in self.decoders:
            x = decoder(x, mask, self.rope)
        x = self.norm(x)
        return self.out(x)

def create_causal_mask(seq_len, device):
    return torch.tril(torch.ones(seq_len, seq_len, dtype=torch.bool, device=device))