from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments, pipeline, DataCollatorWithPadding
from sklearn.metrics import accuracy_score, f1_score
import torch
import numpy as np
import torch.nn.functional as F
import matplotlib.pyplot as plt
from sklearn.metrics import ConfusionMatrixDisplay, confusion_matrix

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



class TransformersSequenceClassifier:
    def __init__(self,
                 model_output_dir,
                 num_labels,
                 tokenizer : AutoTokenizer,
                 model_checkpoint="distilbert-base-uncased"
                 ):
        self.model_output_dir = model_output_dir
        self.tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
        self.model = AutoModelForSequenceClassification.from_pretrained(model_checkpoint, num_labels=num_labels).to(device)

    def tokenizer_func(self, batch):
        return self.tokenizer(batch["inputs"], truncation=True) #, max_len=386

    def train(self, train_dataset, eval_dataset, epochs=2, batch_size=64):
        
        train_tok_dataset = train_dataset.map(self.tokenizer_func, batched=True, remove_columns=('inputs', '__index_level_0__'))
        val_tok_dataset = eval_dataset.map(self.tokenizer_func, batched=True, remove_columns=('inputs', '__index_level_0__'))
        data_collator = DataCollatorWithPadding(tokenizer=self.tokenizer, padding='longest')
        training_args = TrainingArguments(output_dir=f"{self.model_output_dir}/predicting_effective_arguments_distilbert",
                                          num_train_epochs=epochs,
                                          learning_rate=2e-5,
                                          per_device_train_batch_size=batch_size,
                                          per_device_eval_batch_size=batch_size,
                                          weight_decay=0.01,
                                          evaluation_strategy="epoch",
                                          disable_tqdm=False,
                                          logging_steps=len(train_dataset)// batch_size,
                                          push_to_hub=True,
                                          log_level="error")
        self.trainer = Trainer(
            model=self.model,
            args=training_args,
            compute_metrics=self._compute_metrics,
            train_dataset=train_tok_dataset,
            eval_dataset=val_tok_dataset,
            tokenizer=self.tokenizer,
            data_collator=data_collator
        )
        self.trainer.train()

    @staticmethod
    def _compute_metrics(pred):
        labels = pred.label_ids
        preds = pred.predictions.argmax(-1)
        f1 = f1_score(labels, preds, average="weighted")
        acc = accuracy_score(labels, preds)
        return {"accuracy": acc, "f1": f1}

    def forward_pass_with_label(self, batch):
        # Place all input tensors on the same device as the model
        inputs = {k:v.to(device) for k,v in batch.items()
                if k in self.tokenizer.model_input_names}

        with torch.no_grad():
            output = self.model(**inputs)
            pred_label = torch.argmax(output.logits, axis=-1)
            loss = F.cross_entropy(output.logits, batch["label"].to(device), 
                                reduction="none")

        # Place outputs on CPU for compatibility with other dataset columns
        return {"loss": loss.cpu().numpy(), 
                "predicted_label": pred_label.cpu().numpy()}

    def compute_loss_per_pred(self, valid_dataset):
        # Compute loss values
        return valid_dataset.map(self.forward_pass_with_label, batched=True, batch_size=16)

    @staticmethod
    def plot_confusion_matrix(y_preds, y_true, labels):
        cm = confusion_matrix(y_true, y_preds, normalize="true")
        fig, ax = plt.subplots(figsize=(6, 6))
        disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=labels)
        disp.plot(cmap="Blues", values_format=".2f", ax=ax, colorbar=False)
        plt.title("Normalized confusion matrix")
        plt.show()

    def predict_valid_data(self, valid_dataset):
        #trainer = Trainer(model=self.model)
        preds_output = self.trainer.predict(valid_dataset)
        print(preds_output.metrics)
        y_preds = np.argmax(preds_output.predictions, axis=1)
        return y_preds
    
    @staticmethod
    def predict_test_data(model_checkpoint, test_data):
        pipe_classifier = pipeline("text-classification", model=model_checkpoint)
        preds = pipe_classifier(test_data, return_all_scores=True)
        return preds