"""Mikro model językowy w PyTorch (~108M parametrów, 10x mniejszy od MiniPLLM 1.08B) z RoPE.

Architektura:
- Embedding: d_model = 768, tokenizer SentencePiece BPE 32k (tied z lm_head)
- Pozycjonowanie: RoPE (Rotary Position Embedding) na głowicach atencji (brak klasycznego pos_emb)
- 16 bloków (pre-norm + rezyduły):
    h = MLP_down(norm1(x))          4× (Linear+ReLU): 768 -> ... -> 512
    h = h + Attention(norm2(h))     MHA na szyjce, d=512, 8 głowic x 64, causal + RoPE
    y = x + MLP_up(h)               4× (Linear+ReLU): 512 -> ... -> 768
- LayerNorm finalny -> lm_head (tied z embeddingiem)
"""

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


def mlp_widths(d_from: int, d_to: int, n_layers: int) -> list[int]:
    """n_layers szerokości równomiernie rozłożonych między d_from a d_to."""
    step = (d_to - d_from) / n_layers
    return [int(round(d_from + step * (i + 1))) for i in range(n_layers)]


def rotate_half(x: torch.Tensor) -> torch.Tensor:
    """Rotacja połówek wektora: [-x2, x1] dla RoPE."""
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)


class RotaryEmbedding(nn.Module):
    """Rotary Position Embedding (RoPE) zgodne ze standardem LLaMA / Mistral."""

    def __init__(self, dim: int, max_seq_len: int = 4096, theta: float = 10000.0):
        super().__init__()
        self.dim = dim
        self.max_seq_len = max_seq_len
        self.theta = theta
        inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)
        self._set_cos_sin_cache(max_seq_len)

    def _set_cos_sin_cache(self, seq_len: int):
        t = torch.arange(seq_len, dtype=torch.float32)
        freqs = torch.outer(t, self.inv_freq)
        emb = torch.cat((freqs, freqs), dim=-1)
        self.register_buffer("cos_cached", emb.cos(), persistent=False)
        self.register_buffer("sin_cached", emb.sin(), persistent=False)

    def forward(self, x: torch.Tensor, seq_len: int) -> torch.Tensor:
        # x: (B, n_heads, T, d_head)
        if seq_len > self.cos_cached.shape[0]:
            self._set_cos_sin_cache(seq_len)
        cos = self.cos_cached[:seq_len].to(dtype=x.dtype, device=x.device)
        sin = self.sin_cached[:seq_len].to(dtype=x.dtype, device=x.device)
        # Rozszerzenie wymiarów do (1, 1, T, d_head)
        return (x * cos.unsqueeze(0).unsqueeze(0)) + (rotate_half(x) * sin.unsqueeze(0).unsqueeze(0))


class SimpleMLP(nn.Module):
    """Ciąg warstw Linear+ReLU; pierwszy wymiar wejściowy i ostatni wyjściowy podane."""

    def __init__(self, d_in: int, widths: list[int], d_out: int):
        super().__init__()
        dims = [d_in] + widths + [d_out]
        layers = []
        for i in range(len(dims) - 1):
            layers.append(nn.Linear(dims[i], dims[i + 1]))
            if i < len(dims) - 2:  # ReLU między warstwami, nie po ostatniej
                layers.append(nn.ReLU())
        self.net = nn.Sequential(*layers)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)


class MultiHeadAttention(nn.Module):
    """MHA z causal mask i Rotary Position Embedding (RoPE)."""

    def __init__(self, d_model: int, n_heads: int, max_seq_len: int = 4096):
        super().__init__()
        assert d_model % n_heads == 0, f"d_model ({d_model}) musi dzielić się przez n_heads ({n_heads})"
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_head = d_model // n_heads

        self.wq = nn.Linear(d_model, d_model, bias=False)
        self.wk = nn.Linear(d_model, d_model, bias=False)
        self.wv = nn.Linear(d_model, d_model, bias=False)
        self.wo = nn.Linear(d_model, d_model, bias=False)

        self.rotary_emb = RotaryEmbedding(self.d_head, max_seq_len=max_seq_len)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, _ = x.shape
        q = self.wq(x).view(B, T, self.n_heads, self.d_head).transpose(1, 2)
        k = self.wk(x).view(B, T, self.n_heads, self.d_head).transpose(1, 2)
        v = self.wv(x).view(B, T, self.n_heads, self.d_head).transpose(1, 2)

        # Zastosowanie RoPE do zapytań i kluczy
        q = self.rotary_emb(q, T)
        k = self.rotary_emb(k, T)

        att = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        out = att.transpose(1, 2).contiguous().view(B, T, self.d_model)
        return self.wo(out)


class Block(nn.Module):
    """Pre-norm blok z rezydułami:
    h  = MLP_down(norm1(x));  h = h + Attn_RoPE(norm2(h));  y = x + MLP_up(h).
    """

    def __init__(self, d_model: int, n_heads: int, mlp_layers: int, d_mlp: int, max_seq_len: int = 4096):
        super().__init__()
        down = mlp_widths(d_model, d_mlp, mlp_layers)
        up = mlp_widths(d_mlp, d_model, mlp_layers)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_mlp)
        self.mlp_pre = SimpleMLP(d_model, down, d_mlp)
        self.attn = MultiHeadAttention(d_mlp, n_heads, max_seq_len=max_seq_len)
        self.mlp_post = SimpleMLP(d_mlp, up, d_model)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        h = self.mlp_pre(self.norm1(x))
        h = h + self.attn(self.norm2(h))
        return x + self.mlp_post(h)


class MicroPLLM(nn.Module):
    def __init__(
        self,
        vocab_size: int = 32000,
        d_model: int = 768,
        n_layers: int = 16,
        n_heads: int = 8,
        mlp_layers: int = 4,
        d_mlp: int = 512,
        max_seq_len: int = 4096,
        tie_weights: bool = True,
    ):
        super().__init__()
        self.vocab_size = vocab_size
        self.d_model = d_model
        self.n_layers = n_layers
        self.n_heads = n_heads
        self.mlp_layers = mlp_layers
        self.d_mlp = d_mlp
        self.max_seq_len = max_seq_len

        self.tok_emb = nn.Embedding(vocab_size, d_model)
        self.blocks = nn.ModuleList(
            [Block(d_model, n_heads, mlp_layers, d_mlp, max_seq_len=max_seq_len) for _ in range(n_layers)]
        )
        self.ln_f = nn.LayerNorm(d_model)
        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)
        if tie_weights:
            self.lm_head.weight = self.tok_emb.weight
        self._init_stable(n_layers)

    def _init_stable(self, n_layers: int) -> None:
        """Skalowana inicjalizacja (GPT-2 style): std=0.02, końcowe warstwy tuneli
        i lm_head przeskalowane 1/sqrt(2*n_layers) — stabilny start treningu."""
        def init(module: nn.Module) -> None:
            if isinstance(module, nn.Linear):
                nn.init.normal_(module.weight, mean=0.0, std=0.02)
                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=0.02)

        self.apply(init)
        scale = 1.0 / math.sqrt(2 * n_layers)
        for block in self.blocks:
            block.mlp_post.net[-1].weight.data *= scale
        self.lm_head.weight.data *= scale

    def forward(self, idx: torch.Tensor) -> torch.Tensor:
        x = self.tok_emb(idx)
        for block in self.blocks:
            x = block(x)
        x = self.ln_f(x)
        return self.lm_head(x)


if __name__ == "__main__":
    model = MicroPLLM()
    n_total = sum(p.numel() for p in model.parameters())
    n_block = sum(p.numel() for p in model.blocks[0].parameters())
    n_emb = sum(p.numel() for p in [model.tok_emb.weight])
    print(f"Architektura MicroPLLM z RoPE:")
    print(f"MLP (w dół):  {mlp_widths(768, 512, 4)}")
    print(f"MLP (w górę): {mlp_widths(512, 768, 4)}")
    print(f"1 blok: {n_block:,} param")
    print(f"embeddingi (tok): {n_emb:,} param (RoPE nie wymaga wag pozycji)")
    print(f"CAŁKOWITE (tied): {n_total:,} param ({n_total / 1e6:.2f} M)")

    # Test forward causal
    dummy = torch.randint(0, 32000, (2, 64))
    logits = model(dummy)
    print(f"Forward test OK: logits shape {logits.shape}")