from datasets import load_from_disk
from transformers import AutoModelForSequenceClassification, AutoTokenizer
import numpy as np
import torch
from sklearn.metrics import accuracy_score, precision_recall_fscore_support, confusion_matrix
from torch.utils.data import DataLoader
from tqdm import tqdm

def evaluate_model():
    print("Ewaluacja wytrenowanego modelu klasyfikacji emocji...")

    test_dataset = load_from_disk('./emotion_classification/data/emotion_tokenized')['test']
    model = AutoModelForSequenceClassification.from_pretrained('./emotion_classification/models/emotion_classifier')
    tokenizer = AutoTokenizer.from_pretrained('distilbert-base-uncased')

    test_dataset.set_format(type='torch', columns=['input_ids', 'attention_mask', 'label'])
    test_loader = DataLoader(test_dataset, batch_size=64)
    model.eval()

    all_preds = []
    all_labels = []

    with torch.no_grad():
        for batch in tqdm(test_loader, desc="Ewaluacja"):
            input_ids = batch['input_ids']
            attention_mask = batch['attention_mask']
            labels = batch['label']
            outputs = model(input_ids=input_ids, attention_mask=attention_mask)
            preds = torch.argmax(outputs.logits, dim=1)
            all_preds.append(preds)
            all_labels.append(labels)

    y_pred = torch.cat(all_preds).numpy()
    y_true = torch.cat(all_labels).numpy()

    accuracy = accuracy_score(y_true, y_pred)
    precision, recall, f1, _ = precision_recall_fscore_support(y_true, y_pred, average='weighted')
    cm = confusion_matrix(y_true, y_pred)

    print(f'Dokladnosc: {accuracy:.4f}')
    print(f'Precyzja: {precision:.4f}')
    print(f'Czulosc: {recall:.4f}')
    print(f'Wynik F1: {f1:.4f}')
    print('Macierz bledu:')
    print(cm)