import pandas as pd
from sklearn.model_selection import train_test_split
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    Trainer,
    TrainingArguments,
)
import torch
import numpy as np
from sklearn.metrics import accuracy_score, f1_score

# --- 1. Định nghĩa các hằng số và đường dẫn ---
MODEL_NAME = "vinai/phobert-base-v2"
DATA_PATH = "data/labeled_listings.csv"
OUTPUT_DIR = "models/phobert_scam_classifier"
NUM_LABELS = 2

# --- 2. Tạo class Dataset cho PyTorch ---
class ListingDataset(torch.utils.data.Dataset):
    def __init__(self, encodings, labels):
        self.encodings = encodings
        self.labels = labels

    def __getitem__(self, idx):
        item = {key: torch.tensor(val[idx]) for key, val in self.encodings.items()}
        item['labels'] = torch.tensor(self.labels[idx])
        return item

    def __len__(self):
        return len(self.labels)

# --- 3. Hàm tính toán độ đo ---
def compute_metrics(pred):
    labels = pred.label_ids
    preds = np.argmax(pred.predictions, axis=1)
    f1 = f1_score(labels, preds, average="weighted")
    acc = accuracy_score(labels, preds)
    return {"accuracy": acc, "f1": f1}

def main():
    # --- 4. Tải và chuẩn bị dữ liệu ---
    print("Loading and preparing data...")
    df = pd.read_csv(DATA_PATH)
    texts = df['text'].tolist()
    labels = df['label'].tolist()

    # Chia dữ liệu thành tập train và validation
    train_texts, val_texts, train_labels, val_labels = train_test_split(
        texts, labels, test_size=0.2, random_state=42
    )

    # Tải tokenizer
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

    # Tokenize văn bản
    train_encodings = tokenizer(train_texts, truncation=True, padding=True, max_length=256)
    val_encodings = tokenizer(val_texts, truncation=True, padding=True, max_length=256)

    # Tạo đối tượng Dataset
    train_dataset = ListingDataset(train_encodings, train_labels)
    val_dataset = ListingDataset(val_encodings, val_labels)

    # --- 5. Huấn luyện mô hình ---
    print("Starting model training...")
    # Tải mô hình PhoBERT đã được huấn luyện trước
    model = AutoModelForSequenceClassification.from_pretrained(
        MODEL_NAME, num_labels=NUM_LABELS
    )

    # Định nghĩa các tham số huấn luyện
    training_args = TrainingArguments(
        output_dir=OUTPUT_DIR,
        num_train_epochs=3,
        per_device_train_batch_size=4,
        per_device_eval_batch_size=4,
        warmup_steps=500,
        weight_decay=0.01,
        logging_dir='./logs',
        logging_steps=10,
        evaluation_strategy="epoch", # Đánh giá sau mỗi epoch
        save_strategy="epoch",     # Lưu model sau mỗi epoch
        load_best_model_at_end=True, # Tải model tốt nhất khi kết thúc
    )

    # Khởi tạo Trainer
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=val_dataset,
        compute_metrics=compute_metrics,
    )

    # Bắt đầu huấn luyện
    trainer.train()

    # --- 6. Lưu lại model và tokenizer tốt nhất ---
    print(f"Training complete. Saving model to {OUTPUT_DIR}")
    trainer.save_model(OUTPUT_DIR)
    tokenizer.save_pretrained(OUTPUT_DIR)
    print("Model saved successfully!")


if __name__ == "__main__":
    main()