"""TinyLM — from-scratch transformer used for the VIVEK series.
Character-level language model: embedding + positional encoding +
causal Transformer encoder blocks + output head."""
import torch
import torch.nn as nn


class TinyLM(nn.Module):
    def __init__(self, vocab_size, d_model=128, n_layers=2, n_heads=4,
                 ff=256, max_len=256, dropout=0.1):
        super().__init__()
        self.config = dict(vocab_size=vocab_size, d_model=d_model,
                           n_layers=n_layers, n_heads=n_heads, ff=ff,
                           max_len=max_len, dropout=dropout)
        self.tok = nn.Embedding(vocab_size, d_model)
        self.pos = nn.Embedding(max_len, d_model)
        layer = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=n_heads, dim_feedforward=ff,
            dropout=dropout, activation="gelu", batch_first=True,
            norm_first=True)
        self.enc = nn.TransformerEncoder(layer, num_layers=n_layers)
        self.head = nn.Linear(d_model, vocab_size)

    def forward(self, x):
        B, L = x.shape
        h = self.tok(x) + self.pos(torch.arange(L, device=x.device))
        mask = torch.triu(torch.ones(L, L, dtype=torch.bool), diagonal=1)
        h = self.enc(h, mask=mask)
        return self.head(h)

    @staticmethod
    def from_config(path="config.json"):
        import json
        cfg = json.load(open(path))
        return TinyLM(**{k: cfg[k] for k in
                         ("vocab_size", "d_model", "n_layers", "n_heads",
                          "ff", "max_len", "dropout") if k in cfg})