import re, json, os, torch, torch.nn as nn

HERE = os.path.dirname(os.path.abspath(__file__))
cfg = json.load(open(os.path.join(HERE, "vibe_config.json")))
vocab = json.load(open(os.path.join(HERE, "vibe_vocab.json")))
MAX_LEN, EMB, LAYERS, HEADS, FF = cfg["max_len"], cfg["emb"], cfg["layers"], cfg["heads"], cfg["ff"]
LABELS = cfg["labels"]

class VibeNet(nn.Module):
    def __init__(self, V):
        super().__init__()
        self.emb = nn.Embedding(V, EMB, padding_idx=0)
        self.pos = nn.Embedding(MAX_LEN, EMB)
        self.drop = nn.Dropout(0.0)
        layer = nn.TransformerEncoderLayer(EMB, HEADS, FF, 0.0, batch_first=True, activation="gelu")
        self.enc = nn.TransformerEncoder(layer, LAYERS, enable_nested_tensor=False)
        self.head = nn.Linear(EMB, 2)
    def forward(self, x):
        mask = (x == 0)
        pos = torch.arange(x.size(1)).unsqueeze(0)
        h = self.emb(x) + self.pos(pos)
        h = self.enc(h, src_key_padding_mask=mask)
        keep = (~mask).unsqueeze(-1).float()
        pooled = (h * keep).sum(1) / keep.sum(1).clamp(min=1)
        return self.head(pooled)

_model = VibeNet(len(vocab))
_model.load_state_dict(torch.load(os.path.join(HERE, "vibe_model.pt"), map_location="cpu"))
_model.eval()

_tok = lambda s: re.findall(r"[a-z']+", s.lower())
def _encode(s):
    ids = [vocab.get(w, 1) for w in _tok(s)][:MAX_LEN]
    if not ids: ids = [1]
    return ids + [0]*(MAX_LEN-len(ids))

@torch.no_grad()
def predict_vibe(text):
    text = (text or "").strip()
    if not text: return {"label": "Neutral", "confidence": 0.0}
    x = torch.tensor([_encode(text)])
    probs = torch.softmax(_model(x), dim=1)[0]
    i = int(probs.argmax())
    return {"label": LABELS[i], "confidence": round(float(probs[i])*100, 1)}