from datasets import load_dataset, Dataset
import random
import numpy as np
from transformers import (
    AutoTokenizer,
    DataCollatorWithPadding,
    AutoModelForSequenceClassification,
    TrainingArguments,
    Trainer,
    PreTrainedTokenizer,
    ElectraForSequenceClassification,
    EarlyStoppingCallback
)
from dataclasses import dataclass
from sklearn.metrics import accuracy_score, precision_recall_fscore_support

def label_mapper(label: str) -> int:
    if label == "SP":
        return 0
    elif label == "WP":
        return 1
    elif label == "WN":
        return 2
    elif label == "SN":
        return 3
    elif label == "NU":
        return 4
    else:
        raise ValueError(f"Invalid label: {label}")


def process(batch: dict, tokenizer: PreTrainedTokenizer) -> dict:
    new_labels = [label_mapper(label) for label in batch["Polarity"]]
    inputs = tokenizer(batch["Text"], truncation=True)
    batch["input_ids"] = inputs["input_ids"]
    batch["attention_mask"] = inputs["attention_mask"]
    batch["labels"] = new_labels
    return batch


def compute_metrics(eval_pred):
    logits, labels = eval_pred
    predictions = logits.argmax(-1)
    accuracy = accuracy_score(labels, predictions)
    precision, recall, f1, _ = precision_recall_fscore_support(
        labels, predictions, average="macro"
    )
    return {
        "accuracy": accuracy,
        "precision": precision,
        "recall": recall,
        "f1": f1,
    }


def pipeline(args):
    model = AutoModelForSequenceClassification.from_pretrained(args.model_name, num_labels=5)
    tokenizer = AutoTokenizer.from_pretrained(args.model_name)
    dataset = load_dataset(args.dataset_name)
    dataset = dataset.map(process, batched=True, fn_kwargs={'tokenizer': tokenizer})
    dataset =  dataset["train"].train_test_split(args.split_ratio)
    train_dataset = dataset["train"]
    test_dataset = dataset["test"]
    data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

    trainer = Trainer(
        model=model,
        args=TrainingArguments(
            output_dir="./results",
            learning_rate=args.learning_rate,
            per_device_train_batch_size=args.batch_size,
            per_device_eval_batch_size=args.batch_size,
            num_train_epochs=args.epochs,
            weight_decay=0.01,
            eval_strategy="steps",
            save_strategy="steps",
            load_best_model_at_end=True,
            report_to="none",
            save_steps=500,
            eval_steps=500,
            save_total_limit=1,
            logging_steps=500,
            fp16=args.fp16,
            greater_is_better=True,
            metric_for_best_model="f1",
        ),
        train_dataset=train_dataset,
        eval_dataset=test_dataset,
        processing_class=tokenizer,
        data_collator=data_collator,
        compute_metrics=compute_metrics,
        callbacks=[EarlyStoppingCallback(early_stopping_patience=5)]
    )

    trainer.train()
    trainer.evaluate()
    trainer.predict(test_dataset)

    # Push to Hub
    tokenizer.push_to_hub(args.hub_location)
    model.push_to_hub(args.hub_location)

@dataclass
class Arguments:
    model_name: str = "csebuetnlp/banglabert"
    dataset_name: str = "SayedShaun/sentigold"
    split_ratio: float = 0.1
    batch_size: int = 128
    epochs: int = 40
    learning_rate: float = 1e-5
    fp16: bool = True
    hub_location: str = "SayedShaun/bangla-classifier-multiclass"

if __name__=="__main__":
    args = Arguments()
    pipeline(args)