"""
Standalone model definition for swahili-gpt-71m (decoder-only Transformer).

Self-contained — no training repo needed. Loads model_config.json from the same
directory as the weights. See inference.py for an end-to-end example.
"""
import os
import json
import torch
import torch.nn as nn
import torch.nn.functional as F


def _load_config(config_path):
    with open(config_path, "r", encoding="utf-8") as f:
        return json.load(f)


def create_causal_mask(seq_len, device):
    mask = torch.full((seq_len, seq_len), float("-inf"), device=device)
    return torch.triu(mask, diagonal=1)


class TransformerBlock(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        h, eps = cfg["hidden_size"], cfg["layer_norm_epsilon"]
        self.attn = nn.MultiheadAttention(
            embed_dim=h, num_heads=cfg["num_attention_heads"],
            dropout=cfg["dropout"], batch_first=True,
        )
        self.norm1 = nn.LayerNorm(h, eps=eps)
        self.norm2 = nn.LayerNorm(h, eps=eps)
        self.ffn = nn.Sequential(
            nn.Linear(h, cfg["intermediate_size"]), nn.GELU(),
            nn.Linear(cfg["intermediate_size"], h), nn.Dropout(cfg["dropout"]),
        )
        self.dropout = nn.Dropout(cfg["dropout"])

    def forward(self, x, mask=None):
        n = self.norm1(x)
        a, _ = self.attn(n, n, n, attn_mask=mask, need_weights=False)
        x = x + self.dropout(a)
        x = x + self.dropout(self.ffn(self.norm2(x)))
        return x


class KiswahiliLLM(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        h = cfg["hidden_size"]
        self.embedding = nn.Embedding(cfg["vocab_size"], h)
        self.pos_embedding = nn.Embedding(cfg["max_seq_len"], h)
        self.layers = nn.ModuleList(
            [TransformerBlock(cfg) for _ in range(cfg["num_layers"])]
        )
        self.norm = nn.LayerNorm(h, eps=cfg["layer_norm_epsilon"])
        self.output = nn.Linear(h, cfg["vocab_size"])

    def forward(self, input_ids):
        T = input_ids.shape[1]
        pos = torch.arange(T, device=input_ids.device)[None, :]
        x = self.embedding(input_ids) + self.pos_embedding(pos)
        mask = create_causal_mask(T, input_ids.device)
        for layer in self.layers:
            x = layer(x, mask=mask)
        return self.output(self.norm(x))


def load_model(weights_path, config_path=None, device="cpu"):
    if config_path is None:
        config_path = os.path.join(os.path.dirname(weights_path), "model_config.json")
    cfg = _load_config(config_path)
    model = KiswahiliLLM(cfg).to(device)
    state = torch.load(weights_path, map_location=device, weights_only=True)
    model.load_state_dict(state)
    model.eval()
    return model, cfg


@torch.no_grad()
def generate(model, sp, prompt, max_new_tokens=80, temperature=0.8,
             top_k=40, top_p=0.9, device="cpu"):
    ids = torch.tensor(sp.encode(prompt), dtype=torch.long, device=device)[None, :]
    for _ in range(max_new_tokens):
        logits = model(ids)[:, -1, :] / temperature
        if top_k:
            kth = torch.topk(logits, min(top_k, logits.shape[-1])).values[:, -1, None]
            logits = torch.where(logits < kth, torch.full_like(logits, float("-inf")), logits)
        sl, si = torch.sort(logits, descending=True)
        cum = torch.cumsum(F.softmax(sl, dim=-1), dim=-1)
        rm = cum > top_p
        rm[:, 1:] = rm[:, :-1].clone(); rm[:, 0] = False
        sl[rm] = float("-inf")
        logits = sl.gather(-1, si.argsort(-1))
        nxt = torch.multinomial(F.softmax(logits, dim=-1), 1)
        ids = torch.cat([ids, nxt], dim=1)
        if nxt.item() == sp.eos_id():
            break
    return sp.decode(ids[0].tolist()).replace("⁇", "")