from pathlib import Path

import torch
import torch.nn.functional as F
from transformers import BertTokenizer, BertForSequenceClassification

# Folder tempat file ini berada (bareng config.json, model.safetensors, dll.)
BASE_DIR = Path(__file__).resolve().parent


def load_model():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    tokenizer = BertTokenizer.from_pretrained(str(BASE_DIR))
    model = BertForSequenceClassification.from_pretrained(str(BASE_DIR))

    model.to(device)
    model.eval()
    return tokenizer, model, device


def predict_sentiment(text: str, tokenizer, model, device, max_length: int = 128):
    """
    Kembalikan dict:
      - "logits": numpy array (1, 2)
      - "probs": numpy array [p_neg, p_pos]
      - "label_name": "Negatif" atau "Positif"
    """

    encoded = tokenizer(
        text,
        max_length=max_length,
        padding="max_length",
        truncation=True,
        return_tensors="pt",
    )

    encoded = {k: v.to(device) for k, v in encoded.items()}

    with torch.no_grad():
        outputs = model(**encoded)
        logits = outputs.logits  # shape: (1, 2)

    probs = F.softmax(logits, dim=-1).squeeze(0).cpu().numpy()
    label_id = int(probs.argmax())
    label_name = "Negatif" if label_id == 0 else "Positif"

    return {
        "logits": logits.cpu().numpy(),
        "probs": probs,
        "label_name": label_name,
    }