import torch
import tiktoken
from Zuba.Model_training import LanguageClassifier

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

# Label ID map
id_label = {
    0: "hausa",
    1: "igbo",
    2: "Broken English",
    3: "Yoruba"
}

# Load your model architecture
model = LanguageClassifier(num_labels=4).to(device)

print("Loading model weights...")
model.load_state_dict(torch.load("../model/language_classifier_model2.pth", map_location=device))
model.eval()
print("Model loaded successfully")

class Classify:
    def __init__(self):
        self.tokenizer = tiktoken.get_encoding("gpt2")
        self.max_len = 277
        self.pad_id = 50256

    def encode_text(self, text):
        """Tokenize → truncate → pad"""
        ids = self.tokenizer.encode(text)

        # Truncate
        ids = ids[:self.max_len]

        # Pad
        if len(ids) < self.max_len:
            ids += [self.pad_id] * (self.max_len - len(ids))

        return ids

    def predict(self, text_or_list):
        """Predict label for one text or many."""

        # Handle single vs batch
        if isinstance(text_or_list, str):
            texts = [text_or_list]
            single = True
        else:
            texts = text_or_list
            single = False

        encoded = [self.encode_text(t) for t in texts]

        # Convert to tensor
        input_ids = torch.tensor(encoded, dtype=torch.long).to(device)

        # Model prediction
        with torch.no_grad():
            logits = model(input_ids)

        pred_ids = torch.argmax(logits, dim=1).tolist()
        labels = [id_label[p] for p in pred_ids]

        return labels[0] if single else labels
