import math, torch, torch.nn as nn, torch.nn.functional as F
import sentencepiece as spm

sp = spm.SentencePieceProcessor()
sp.load("tok.model")
VOCAB = sp.get_piece_size()

def encode(t):return sp.encode(t, out_type=int)

def rope(T, d, device):
    inv = 1.0 / (10000 ** (torch.arange(0, d, 2, device=device).float() / d))
    pos = torch.arange(T, device=device).float()
    freqs = torch.outer(pos, inv)
    cos = freqs.cos()[None, None, :, :]
    sin = freqs.sin()[None, None, :, :]
    return cos, sin

def apply_rope(x, cos, sin):
    x1 = x[..., ::2]
    x2 = x[..., 1::2]
    out = torch.empty_like(x)
    out[..., ::2] = x1 * cos - x2 * sin
    out[..., 1::2] = x1 * sin + x2 * cos
    return out

def causal_mask(T, device):
    return torch.triu(torch.ones(T, T, device=device, dtype=torch.bool), diagonal=1)

class RMSNorm(nn.Module):
    def __init__(self, d, eps=1e-6):
        super().__init__()
        self.w = 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.w

class SwiGLU(nn.Module):
    def __init__(self, d):
        super().__init__()
        self.fc = nn.Linear(d, d * 4 * 2)
        self.out = nn.Linear(d * 4, d)

    def forward(self, x):
        a, b = self.fc(x).chunk(2, dim=-1)
        return self.out(F.silu(a) * b)

class Attention(nn.Module):
    def __init__(self, d, h):
        super().__init__()
        assert d % h == 0
        self.h = h
        self.dh = d // h
        self.qkv = nn.Linear(d, d * 3)
        self.out = nn.Linear(d, d)

    def forward(self, x):
        B, T, C = x.shape
        q, k, v = self.qkv(x).chunk(3, dim=-1)
        q = q.view(B, T, self.h, self.dh).transpose(1, 2)
        k = k.view(B, T, self.h, self.dh).transpose(1, 2)
        v = v.view(B, T, self.h, self.dh).transpose(1, 2)
        cos, sin = rope(T, self.dh, x.device)
        q = apply_rope(q, cos, sin)
        k = apply_rope(k, cos, sin)
        att = (q @ k.transpose(-2, -1)) / math.sqrt(self.dh)
        mask = causal_mask(T, x.device)
        att = att.masked_fill(mask, float("-inf"))
        att = F.softmax(att, dim=-1)
        y = att @ v
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        return self.out(y)

class Block(nn.Module):
    def __init__(self, d, h):
        super().__init__()
        self.n1 = RMSNorm(d)
        self.attn = Attention(d, h)
        self.n2 = RMSNorm(d)
        self.mlp = SwiGLU(d)

    def forward(self, x):
        x = x + self.attn(self.n1(x))
        x = x + self.mlp(self.n2(x))
        return x

class GPT(nn.Module):
    def __init__(self):
        super().__init__()
        d = 1024
        L = 12
        H = 16
        self.emb = nn.Embedding(VOCAB, d)
        self.blocks = nn.ModuleList([Block(d, H) for _ in range(L)])
        self.norm = RMSNorm(d)
        self.head = nn.Linear(d, VOCAB, bias=False)
        self.head.weight = self.emb.weight

    def forward(self, x):
        x = self.emb(x)
        for b in self.blocks:
            x = b(x)
        return self.head(self.norm(x))

Qdevice = "cuda" if torch.cuda.is_available() else "cpu"

def load_qed(path, device="cpu",dw=False):
    package = torch.load(path,map_location=device)
    model = GPT()
    model.load_state_dict(package["state_dict"])
    None if dw else print(f"Loading {package['model_name']} by {package['author']}")
    model.to(device)
    model.eval()
    return model