#!/usr/bin/env python3
# SPDX-License-Identifier: MIT
# Copyright (c) 2026 Hamid Wakili <hamid@ideployed.com>
"""Llama-style decoder-only Transformer for German Commons pretraining.

This module is the larger German-language base-model counterpart to the
original ``MiniLeipzigGPT``. It intentionally uses only PyTorch and standard
Python. There are no imports from Hugging Face ``transformers`` or any other
pre-built model library, and no pre-trained weights are consumed.

Architecture
------------
The implementation follows the decoder-only Transformer family introduced by
Vaswani et al. (2017), using modern small-model choices that are common in
Llama-style language models:

* causal self-attention with PyTorch scaled-dot-product attention;
* RoPE rotary position embeddings instead of learned absolute positions;
* RMSNorm pre-normalization;
* SwiGLU feed-forward layers;
* bias-free linear layers;
* tied token embedding / output projection by default;
* optional grouped-query attention via ``n_kv_heads``.

Training provenance
-------------------
The model is designed to train from random weights on locally prepared German
Commons token shards produced with the repository's own byte-level BPE
tokenizer. The implementation keeps the architecture and training inputs
inspectable within the project.
"""

from __future__ import annotations

from dataclasses import asdict, dataclass
from typing import Dict, Optional, Tuple

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


@dataclass
class GermanGPTConfig:
    """Configuration for :class:`GermanGPT`.

    The field names mirror ``configs.ModelConfig`` so checked training plans can
    be passed directly into this model.
    """

    name: str = "german-commons"
    vocab_size: int = 32768
    d_model: int = 1024
    n_layers: int = 24
    n_heads: int = 16
    n_kv_heads: int = 16
    d_ff: int = 2816
    context_len: int = 2048
    rope_theta: float = 10000.0
    norm: str = "rmsnorm"
    mlp: str = "swiglu"
    tie_embeddings: bool = True
    dropout: float = 0.0

    @property
    def head_dim(self) -> int:
        """Return the per-head hidden dimension."""

        return self.d_model // self.n_heads

    @classmethod
    def from_dict(cls, data: Dict[str, object]) -> "GermanGPTConfig":
        """Build a config from a plain checkpoint dictionary."""

        return cls(
            name=str(data.get("name", "german-commons")),
            vocab_size=int(data["vocab_size"]),
            d_model=int(data.get("d_model", 1024)),
            n_layers=int(data.get("n_layers", 24)),
            n_heads=int(data.get("n_heads", 16)),
            n_kv_heads=int(data.get("n_kv_heads", data.get("n_heads", 16))),
            d_ff=int(data.get("d_ff", 2816)),
            context_len=int(data.get("context_len", 2048)),
            rope_theta=float(data.get("rope_theta", 10000.0)),
            norm=str(data.get("norm", "rmsnorm")),
            mlp=str(data.get("mlp", "swiglu")),
            tie_embeddings=bool(data.get("tie_embeddings", True)),
            dropout=float(data.get("dropout", 0.0)),
        )

    def to_dict(self) -> Dict[str, object]:
        """Return a JSON/checkpoint-friendly dictionary."""

        return asdict(self)


def config_from_model_config(config: object, vocab_size: Optional[int] = None) -> GermanGPTConfig:
    """Convert ``configs.ModelConfig`` or a similar object to ``GermanGPTConfig``."""

    data = config.__dict__.copy()
    if vocab_size is not None:
        data["vocab_size"] = vocab_size
    return GermanGPTConfig.from_dict(data)


class RMSNorm(nn.Module):
    """Root-mean-square normalization used in Llama-style blocks."""

    def __init__(self, dim: int, eps: float = 1e-6) -> None:
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Normalize the final dimension of ``x`` without mean subtraction."""

        dtype = x.dtype
        x_float = x.float()
        variance = x_float.pow(2).mean(dim=-1, keepdim=True)
        normalized = x_float * torch.rsqrt(variance + self.eps)
        return (normalized.to(dtype) * self.weight)


class RotaryEmbedding(nn.Module):
    """Precomputed RoPE cosine/sine tables for causal attention."""

    def __init__(self, dim: int, max_seq_len: int, theta: float = 10000.0) -> None:
        super().__init__()
        if dim % 2 != 0:
            raise ValueError("RoPE head dimension must be even")
        inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
        positions = torch.arange(max_seq_len, dtype=torch.float32)
        freqs = torch.outer(positions, inv_freq)
        self.register_buffer("cos", freqs.cos(), persistent=False)
        self.register_buffer("sin", freqs.sin(), persistent=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Apply rotary embeddings to ``x`` shaped ``(B, T, H, D)``."""

        seq_len = x.size(1)
        cos = self.cos[:seq_len].to(dtype=x.dtype, device=x.device).view(1, seq_len, 1, -1)
        sin = self.sin[:seq_len].to(dtype=x.dtype, device=x.device).view(1, seq_len, 1, -1)
        even = x[..., 0::2]
        odd = x[..., 1::2]
        rotated_even = even * cos - odd * sin
        rotated_odd = even * sin + odd * cos
        return torch.stack((rotated_even, rotated_odd), dim=-1).flatten(-2)


class CausalSelfAttention(nn.Module):
    """Causal multi-head attention with optional grouped-query attention."""

    def __init__(self, config: GermanGPTConfig) -> None:
        super().__init__()
        if config.d_model % config.n_heads != 0:
            raise ValueError("d_model must be divisible by n_heads")
        if config.n_heads % config.n_kv_heads != 0:
            raise ValueError("n_heads must be divisible by n_kv_heads")

        self.n_heads = config.n_heads
        self.n_kv_heads = config.n_kv_heads
        self.head_dim = config.head_dim
        self.kv_repeat = config.n_heads // config.n_kv_heads

        self.q_proj = nn.Linear(config.d_model, config.n_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(config.d_model, config.n_kv_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(config.d_model, config.n_kv_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(config.n_heads * self.head_dim, config.d_model, bias=False)
        self.dropout = float(config.dropout)
        self.rope = RotaryEmbedding(config.head_dim, config.context_len, config.rope_theta)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Apply causal self-attention to hidden states ``x``."""

        batch_size, seq_len, _ = x.shape
        q = self.q_proj(x).view(batch_size, seq_len, self.n_heads, self.head_dim)
        k = self.k_proj(x).view(batch_size, seq_len, self.n_kv_heads, self.head_dim)
        v = self.v_proj(x).view(batch_size, seq_len, self.n_kv_heads, self.head_dim)

        q = self.rope(q)
        k = self.rope(k)

        if self.kv_repeat > 1:
            k = k.repeat_interleave(self.kv_repeat, dim=2)
            v = v.repeat_interleave(self.kv_repeat, dim=2)

        q = q.transpose(1, 2)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)

        y = F.scaled_dot_product_attention(
            q,
            k,
            v,
            attn_mask=None,
            dropout_p=self.dropout if self.training else 0.0,
            is_causal=True,
        )
        y = y.transpose(1, 2).contiguous().view(batch_size, seq_len, -1)
        return self.o_proj(y)


class SwiGLU(nn.Module):
    """Bias-free SwiGLU feed-forward layer."""

    def __init__(self, config: GermanGPTConfig) -> None:
        super().__init__()
        self.gate_proj = nn.Linear(config.d_model, config.d_ff, bias=False)
        self.up_proj = nn.Linear(config.d_model, config.d_ff, bias=False)
        self.down_proj = nn.Linear(config.d_ff, config.d_model, bias=False)
        self.dropout = nn.Dropout(config.dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Apply ``SiLU(gate) * up`` followed by down projection."""

        return self.dropout(self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)))


class TransformerBlock(nn.Module):
    """Pre-normalized decoder block."""

    def __init__(self, config: GermanGPTConfig) -> None:
        super().__init__()
        self.attn_norm = RMSNorm(config.d_model)
        self.attn = CausalSelfAttention(config)
        self.ffn_norm = RMSNorm(config.d_model)
        self.ffn = SwiGLU(config)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Apply attention and feed-forward residual paths."""

        x = x + self.attn(self.attn_norm(x))
        x = x + self.ffn(self.ffn_norm(x))
        return x


class GermanGPT(nn.Module):
    """Decoder-only German base language model."""

    def __init__(self, config: GermanGPTConfig) -> None:
        super().__init__()
        validate_config(config)
        self.config = config
        self.token_embedding = nn.Embedding(config.vocab_size, config.d_model)
        self.dropout = nn.Dropout(config.dropout)
        self.blocks = nn.ModuleList([TransformerBlock(config) for _ in range(config.n_layers)])
        self.norm = RMSNorm(config.d_model)
        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
        if config.tie_embeddings:
            self.lm_head.weight = self.token_embedding.weight

        self.apply(self._init_weights)

    def forward(
        self,
        input_ids: torch.Tensor,
        labels: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
        """Run the model and optionally compute next-token loss."""

        if input_ids.ndim != 2:
            raise ValueError("input_ids must have shape (batch, time)")
        _, seq_len = input_ids.shape
        if seq_len > self.config.context_len:
            raise ValueError(
                f"Sequence length {seq_len} exceeds context_len {self.config.context_len}"
            )

        x = self.dropout(self.token_embedding(input_ids))
        for block in self.blocks:
            x = block(x)
        x = self.norm(x)
        logits = self.lm_head(x)

        loss = None
        if labels is not None:
            loss = F.cross_entropy(
                logits.reshape(-1, logits.size(-1)),
                labels.reshape(-1),
                ignore_index=-100,
            )
        return logits, loss

    @staticmethod
    def _init_weights(module: nn.Module) -> None:
        """Initialize trainable weights with a GPT/Llama-compatible normal init."""

        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)


def validate_config(config: GermanGPTConfig) -> None:
    """Validate architectural invariants before allocating the model."""

    if config.vocab_size <= 0:
        raise ValueError("vocab_size must be positive")
    if config.d_model <= 0 or config.n_layers <= 0:
        raise ValueError("d_model and n_layers must be positive")
    if config.context_len <= 0:
        raise ValueError("context_len must be positive")
    if config.d_model % config.n_heads != 0:
        raise ValueError("d_model must be divisible by n_heads")
    if config.n_heads % config.n_kv_heads != 0:
        raise ValueError("n_heads must be divisible by n_kv_heads")
    if config.head_dim % 2 != 0:
        raise ValueError("head_dim must be even for RoPE")
    if config.norm != "rmsnorm":
        raise ValueError("Only rmsnorm is implemented")
    if config.mlp != "swiglu":
        raise ValueError("Only swiglu is implemented")
    if not 0.0 <= config.dropout < 1.0:
        raise ValueError("dropout must be in [0, 1)")


def count_parameters(model: nn.Module) -> int:
    """Return the number of unique trainable parameters."""

    seen = set()
    total = 0
    for parameter in model.parameters():
        if id(parameter) in seen:
            continue
        seen.add(id(parameter))
        if parameter.requires_grad:
            total += parameter.numel()
    return total