"""
tiny_llama.py — крошечная LLaMA для llama.cpp.

Архитектура как у настоящей LLaMA: RMSNorm + RoPE + SwiGLU + causal attention.
Байт-вокаб (256 токенов) чтоб не возиться со спм/бпе.

Запуск:
    python tiny_llama.py train russian.txt eblangpt1984.gguf
    python tiny_llama.py test eblangpt1984.gguf    # проверить что корректно читается

После экспорта пробуй:
    llama-cli -m eblangpt1984.gguf -p "привет" -n 200 --temp 0.8
"""

import sys
import math
import time
import struct
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from gguf import GGUFWriter, TokenType

# ===================== КОНФИГ =====================
VOCAB       = 256
N_EMBD      = 64
N_LAYERS    = 2
N_HEADS     = 4
HEAD_DIM    = N_EMBD // N_HEADS     # 16
N_FF        = 128
CTX_LEN     = 64
ROPE_THETA  = 10000.0
RMS_EPS     = 1e-5

ARCH = "llama"
MODEL_NAME = "eblangpt1984"

# ===================== МОДЕЛЬ =====================
class RMSNorm(nn.Module):
    def __init__(self, d, eps=RMS_EPS):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(d))
        self.eps = eps
    def forward(self, x):
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight


def precompute_rope(seqlen, head_dim, theta=ROPE_THETA, device="cpu"):
    freqs = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
    t = torch.arange(seqlen, device=device).float()
    f = torch.outer(t, freqs)              # [T, D/2]
    return torch.cos(f), torch.sin(f)      # каждый [T, D/2]


def apply_rope(x, cos, sin):
    # x: [B, H, T, D]. Используется "interleaved" схема — как в llama.cpp.
    T = x.size(-2)
    cos = cos[:T].unsqueeze(0).unsqueeze(0)
    sin = sin[:T].unsqueeze(0).unsqueeze(0)
    x1, x2 = x[..., 0::2], x[..., 1::2]
    y1 = x1 * cos - x2 * sin
    y2 = x1 * sin + x2 * cos
    return torch.stack((y1, y2), dim=-1).flatten(-2)


class Block(nn.Module):
    def __init__(self):
        super().__init__()
        self.attn_norm = RMSNorm(N_EMBD)
        self.wq = nn.Linear(N_EMBD, N_EMBD, bias=False)
        self.wk = nn.Linear(N_EMBD, N_EMBD, bias=False)
        self.wv = nn.Linear(N_EMBD, N_EMBD, bias=False)
        self.wo = nn.Linear(N_EMBD, N_EMBD, bias=False)

        self.ffn_norm = RMSNorm(N_EMBD)
        self.w_gate = nn.Linear(N_EMBD, N_FF, bias=False)
        self.w_up   = nn.Linear(N_EMBD, N_FF, bias=False)
        self.w_down = nn.Linear(N_FF, N_EMBD, bias=False)

    def forward(self, x, cos, sin, mask):
        B, T, D = x.shape
        h = self.attn_norm(x)
        q = self.wq(h).view(B, T, N_HEADS, HEAD_DIM).transpose(1, 2)
        k = self.wk(h).view(B, T, N_HEADS, HEAD_DIM).transpose(1, 2)
        v = self.wv(h).view(B, T, N_HEADS, HEAD_DIM).transpose(1, 2)
        q = apply_rope(q, cos, sin)
        k = apply_rope(k, cos, sin)

        att = (q @ k.transpose(-2, -1)) / math.sqrt(HEAD_DIM)
        att = att.masked_fill(mask[:T, :T], float("-inf"))
        att = F.softmax(att, dim=-1)
        out = (att @ v).transpose(1, 2).contiguous().view(B, T, D)
        x = x + self.wo(out)

        h = self.ffn_norm(x)
        x = x + self.w_down(F.silu(self.w_gate(h)) * self.w_up(h))
        return x


class TinyLlama(nn.Module):
    def __init__(self):
        super().__init__()
        self.embed = nn.Embedding(VOCAB, N_EMBD)
        self.blocks = nn.ModuleList([Block() for _ in range(N_LAYERS)])
        self.norm = RMSNorm(N_EMBD)
        self.lm_head = nn.Linear(N_EMBD, VOCAB, bias=False)

        cos, sin = precompute_rope(CTX_LEN, HEAD_DIM)
        self.register_buffer("cos", cos, persistent=False)
        self.register_buffer("sin", sin, persistent=False)
        mask = torch.triu(torch.ones(CTX_LEN, CTX_LEN, dtype=torch.bool), diagonal=1)
        self.register_buffer("mask", mask, persistent=False)

    def forward(self, x):
        h = self.embed(x)
        for b in self.blocks:
            h = b(h, self.cos, self.sin, self.mask)
        return self.lm_head(self.norm(h))


# ===================== ОБУЧЕНИЕ =====================
def train_model(text_path, out_path, steps=3000, lr=3e-3, bs=16):
    with open(text_path, "rb") as f:
        data = f.read()
    print(f"текст: {len(data)} байт")

    ids = np.frombuffer(data, dtype=np.uint8).astype(np.int64)
    torch.manual_seed(42)
    model = TinyLlama()
    n_params = sum(p.numel() for p in model.parameters())
    print(f"модель: {n_params:,} параметров")

    opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)

    def sample_batch():
        idx = np.random.randint(0, len(ids) - CTX_LEN - 1, size=bs)
        x = np.stack([ids[i:i + CTX_LEN] for i in idx])
        y = np.stack([ids[i + 1:i + CTX_LEN + 1] for i in idx])
        return torch.from_numpy(x), torch.from_numpy(y)

    model.train()
    t0 = time.time()
    run = 0.0
    for step in range(steps):
        x, y = sample_batch()
        logits = model(x)
        loss = F.cross_entropy(logits.view(-1, VOCAB), y.view(-1))
        opt.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        opt.step()
        run = 0.98 * run + 0.02 * loss.item() if step else loss.item()
        if step % 100 == 0 or step == steps - 1:
            dt = time.time() - t0
            print(f"  шаг {step:5d}/{steps}  loss={run:.3f}  [{dt:.1f}s]")

    export_gguf(model, out_path)


# ===================== ЭКСПОРТ В GGUF (архитектура "llama") =====================
def export_gguf(model, out_path):
    print(f"экспорт в {out_path}...")
    w = GGUFWriter(out_path, ARCH)

    # --- метадата LLaMA ---
    w.add_name(MODEL_NAME)
    w.add_context_length(CTX_LEN)
    w.add_embedding_length(N_EMBD)
    w.add_block_count(N_LAYERS)
    w.add_feed_forward_length(N_FF)
    w.add_head_count(N_HEADS)
    w.add_head_count_kv(N_HEADS)           # без GQA
    w.add_layer_norm_rms_eps(RMS_EPS)
    w.add_rope_dimension_count(HEAD_DIM)
    w.add_rope_freq_base(ROPE_THETA)
    w.add_file_type(0)                     # all F32

    # --- байт-токенайзер ---
    tokens = [f"<0x{b:02X}>".encode("utf-8") for b in range(VOCAB)]
    scores = [-1000.0 + float(i) for i in range(VOCAB)]
    types  = [TokenType.BYTE.value] * VOCAB

    w.add_tokenizer_model("llama")
    w.add_tokenizer_pre("default")
    w.add_token_list(tokens)
    w.add_token_scores(scores)
    w.add_token_types(types)
    w.add_bos_token_id(0)
    w.add_eos_token_id(0)
    w.add_unk_token_id(0)
    w.add_add_bos_token(False)
    w.add_add_eos_token(False)

    # --- тензоры ---
    sd = model.state_dict()

    def add(name, tensor):
        arr = tensor.detach().to(torch.float32).cpu().numpy()
        w.add_tensor(name, arr)

    add("token_embd.weight", sd["embed.weight"])       # [V, E]
    add("output_norm.weight", sd["norm.weight"])       # [E]
    add("output.weight",      sd["lm_head.weight"])    # [V, E]

    for i in range(N_LAYERS):
        p = f"blocks.{i}"
        q = f"blk.{i}"
        add(f"{q}.attn_norm.weight", sd[f"{p}.attn_norm.weight"])
        add(f"{q}.attn_q.weight",    sd[f"{p}.wq.weight"])
        add(f"{q}.attn_k.weight",    sd[f"{p}.wk.weight"])
        add(f"{q}.attn_v.weight",    sd[f"{p}.wv.weight"])
        add(f"{q}.attn_output.weight", sd[f"{p}.wo.weight"])
        add(f"{q}.ffn_norm.weight",  sd[f"{p}.ffn_norm.weight"])
        add(f"{q}.ffn_gate.weight",  sd[f"{p}.w_gate.weight"])
        add(f"{q}.ffn_up.weight",    sd[f"{p}.w_up.weight"])
        add(f"{q}.ffn_down.weight",  sd[f"{p}.w_down.weight"])

    w.write_header_to_file()
    w.write_kv_data_to_file()
    w.write_tensors_to_file()
    w.close()
    print(f"готово: {out_path}")


# ===================== ПРОВЕРКА ФАЙЛА =====================
def test_gguf(path):
    with open(path, "rb") as f:
        magic, ver = struct.unpack("<II", f.read(8))
        tc, kv = struct.unpack("<QQ", f.read(16))
    print(f"GGUF v{ver}, магия=0x{magic:08X}")
    print(f"  тензоров: {tc}")
    print(f"  метадата записей: {kv}")
    print(f"  размер файла: {__import__('os').path.getsize(path)} байт")
    assert magic == 0x46554747, "битая магия"
    print("формат валидный ✓")


# ===================== MAIN =====================
if __name__ == "__main__":
    if len(sys.argv) < 2:
        print(__doc__)
        sys.exit(1)

    cmd = sys.argv[1]
    if cmd == "train":
        text = sys.argv[2] if len(sys.argv) > 2 else "russian_mini.txt"
        out  = sys.argv[3] if len(sys.argv) > 3 else "eblangpt1984.gguf"
        steps = int(sys.argv[4]) if len(sys.argv) > 4 else 3000
        train_model(text, out, steps=steps)
        test_gguf(out)
    elif cmd == "test":
        test_gguf(sys.argv[2])
    else:
        print(__doc__)