import json, numpy as np, torch
import torch.nn as nn
from safetensors.torch import load_file

class TransformerBlock(nn.Module):
    def __init__(self, dim, heads, ff_dim, dropout):
        super().__init__()
        self.mha = nn.MultiheadAttention(dim, heads, dropout=dropout, batch_first=True)
        self.ln1 = nn.LayerNorm(dim); self.ln2 = nn.LayerNorm(dim)
        self.drop = nn.Dropout(dropout)
        self.ffn = nn.Sequential(nn.Linear(dim, ff_dim), nn.ReLU(), nn.Dropout(dropout), nn.Linear(ff_dim, dim))
    def forward(self, x, key_padding_mask=None):
        attn_out, _ = self.mha(x, x, x, key_padding_mask=key_padding_mask, need_weights=False)
        x = self.ln1(x + self.drop(attn_out))
        ff_out = self.ffn(x)
        return self.ln2(x + self.drop(ff_out))

class BiGRUResidualBlock(nn.Module):
    def __init__(self, dim, dropout):
        super().__init__()
        self.gru = nn.GRU(dim, dim//2, num_layers=1, batch_first=True, bidirectional=True)
        self.ln = nn.LayerNorm(dim); self.drop = nn.Dropout(dropout)
    def forward(self, x):
        out, _ = self.gru(x)
        return self.ln(x + self.drop(out))

class MeterModel(nn.Module):
    def __init__(self, vocab_size, num_classes, T, emb_dim=64, num_heads=4, ff_dim=256, gru_blocks=3, dropout=0.1):
        super().__init__()
        self.emb = nn.Embedding(vocab_size, emb_dim, padding_idx=0)
        self.pos = nn.Embedding(T, emb_dim)
        self.drop = nn.Dropout(dropout)
        self.tr = TransformerBlock(emb_dim, num_heads, ff_dim, dropout)
        self.gru_blocks = nn.ModuleList([BiGRUResidualBlock(emb_dim, dropout) for _ in range(gru_blocks)])
        self.head = nn.Sequential(nn.Linear(emb_dim, 128), nn.ReLU(), nn.Dropout(dropout), nn.Linear(128, num_classes))
    def forward(self, x):
        B, T = x.shape
        pos = torch.arange(T, device=x.device).unsqueeze(0).expand(B, T)
        h = self.drop(self.emb(x) + self.pos(pos))
        pad_mask = (x == 0)
        h = self.tr(h, key_padding_mask=pad_mask)
        for blk in self.gru_blocks:
            h = blk(h)
        mask = (~pad_mask).float().unsqueeze(-1)
        pooled = (h * mask).sum(1) / mask.sum(1).clamp_min(1.0)
        return self.head(pooled)

def load_local(repo_dir="."):
    cfg = json.load(open(f"{repo_dir}/config.json", "r", encoding="utf-8"))
    stoi = json.load(open(f"{repo_dir}/stoi.json", "r", encoding="utf-8"))
    classes = cfg["classes"]
    model = MeterModel(
        vocab_size=cfg["vocab_size"], num_classes=cfg["num_classes"], T=cfg["max_length"],
        emb_dim=cfg["emb_dim"], num_heads=cfg["num_heads"], ff_dim=cfg["ff_dim"],
        gru_blocks=cfg["gru_blocks"], dropout=cfg["dropout"]
    )
    sd = load_file(f"{repo_dir}/model.safetensors")
    model.load_state_dict(sd)
    model.eval()
    return model, stoi, classes, cfg["max_length"]

def encode(text, stoi, max_len):
    ids = np.zeros((max_len,), dtype=np.int64)
    for i, ch in enumerate(text[:max_len]):
        ids[i] = stoi.get(ch, 0)
    return torch.from_numpy(ids).unsqueeze(0)

@torch.no_grad()
def predict(text, model, stoi, classes, max_len, topk=5):
    x = encode(text, stoi, max_len)
    logits = model(x)[0]
    probs = torch.softmax(logits, dim=-1)
    topv, topi = torch.topk(probs, k=min(topk, probs.numel()))
    pred = classes[int(topi[0])]
    top = [(classes[int(i)], float(v)) for v, i in zip(topv, topi)]
    return pred, top