#!/usr/bin/env python
# coding: utf-8

# Installing the transformers library and additional libraries if looking process
    
# Install the required Transformers library (if running in a notebook)
get_ipython().system('pip install -q transformers')


# Importing stock ml libraries
# Import libraries for data handling, modeling, and evaluation
import numpy as np
import pandas as pd
from sklearn import metrics
import transformers
from torchviz import make_dot
import torch
from torch.utils.data import Dataset, DataLoader, RandomSampler, SequentialSampler
from transformers import BertTokenizer, BertModel, BertConfig
from sklearn.preprocessing import LabelEncoder

# Preparing for TPU usage
# import torch_xla
# import torch_xla.core.xla_model as xm
# device = xm.xla_device()

# # Setting up the device for GPU usage

from torch import cuda
# Set up device: use GPU if available, else fallback to CPU
device = 'cuda' if cuda.is_available() else 'cpu'


# Load the dataset from Excel file
df = pd.read_excel("DatiPerClassificazioneIncrociata.xlsx")

colonna_rif = 'ricercatore2'

# allowed_values (set of valid categories to include in the training subset)

# NOTE: For 'ricercatore2', exclude 'comunicazione normativa' since it was not used
# Define the set of valid categories to keep in the dataset
valori_ammessi = {'Comunicazione dei /nei processi partecipativi', 'Comunicazione di servizio', 'Comunicazione delle attività istituzionali', 'Comunicazione delle issues socialmente rilevanti', 'Comunicazione delle politiche pubbliche', 'Comunicazione interna e valorizzazione delle risorse umane', 'Comunicazione di brand'}

# Filter dataset to include only allowed categories
df_filtrato = df[df[colonna_rif].isin(valori_ammessi)]

valori_distinti = df_filtrato[colonna_rif ].unique()

# Encode string labels into integers
label_encoder = LabelEncoder()
df_filtrato[colonna_rif] = label_encoder.fit_transform(df_filtrato[colonna_rif])

df_filtrato['list'] = df_filtrato[colonna_rif]
new_df = df_filtrato[['Text', 'list']].copy()
new_df.head()


# Defining some key variables that will be used later on in the training
# Define model training hyperparameters
MAX_LEN = 200
TRAIN_BATCH_SIZE = 8
VALID_BATCH_SIZE = 4
EPOCHS = 5
LEARNING_RATE = 1e-05
tokenizer = BertTokenizer.from_pretrained('dbmdz/bert-base-italian-cased')


# Define custom PyTorch Dataset to prepare model inputs
class CustomDataset(Dataset):

    def __init__(self, dataframe, tokenizer, max_len):
        self.tokenizer = tokenizer
        self.data = dataframe
        self.comment_text = dataframe.Text
        self.targets = self.data.list
        self.max_len = max_len

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

    def __getitem__(self, index):
        comment_text = str(self.comment_text[index])
        comment_text = " ".join(comment_text.split())

        inputs = self.tokenizer.encode_plus(
            comment_text,
            None,
            add_special_tokens=True,
            max_length=self.max_len,
            pad_to_max_length=True,
            return_token_type_ids=True
        )
        ids = inputs['input_ids']
        mask = inputs['attention_mask']
        token_type_ids = inputs["token_type_ids"]


        return {
            'ids': torch.tensor(ids, dtype=torch.long),
            'mask': torch.tensor(mask, dtype=torch.long),
            'token_type_ids': torch.tensor(token_type_ids, dtype=torch.long),
            'targets': torch.tensor(self.targets[index], dtype=torch.long)
        }


# Creating the dataset and dataloader for the neural network

train_size = 0.8
# Split dataset into training and testing subsets
train_dataset=new_df.sample(frac=train_size,random_state=200)
test_dataset=new_df.drop(train_dataset.index).reset_index(drop=True)
train_dataset = train_dataset.reset_index(drop=True)


print("FULL Dataset: {}".format(new_df.shape))
print("TRAIN Dataset: {}".format(train_dataset.shape))
print("TEST Dataset: {}".format(test_dataset.shape))

# Define model training hyperparameters
training_set = CustomDataset(train_dataset, tokenizer, MAX_LEN)
# Define model training hyperparameters
testing_set = CustomDataset(test_dataset, tokenizer, MAX_LEN)

train_params = {'batch_size': TRAIN_BATCH_SIZE,
                'shuffle': True,
                'num_workers': 0
                }

test_params = {'batch_size': VALID_BATCH_SIZE,
                'shuffle': True,
                'num_workers': 0
                }

# Create PyTorch DataLoaders for batching
training_loader = DataLoader(training_set, **train_params)
testing_loader = DataLoader(testing_set, **test_params)

testing_set[50]


# Creating the customized model, by adding a drop out and a dense layer on top of distil bert to get the final output for the model.
# Define BERT-based classification model with dropout and linear layer
class BERTClass(torch.nn.Module):
    def __init__(self):
        super(BERTClass, self).__init__()
        self.l1 = transformers.BertModel.from_pretrained('dbmdz/bert-base-italian-cased')
        self.l2 = torch.nn.Dropout(0.3)
        # For 'ricercatore2', use 7 output classes because "comunicazione normativa" was not used and therefore does not appear among the classes
        self.l3 = torch.nn.Linear(768, 7)

    def forward(self, ids, mask, token_type_ids):
        _, output_1= self.l1(ids, attention_mask = mask, token_type_ids = token_type_ids, return_dict=False)
        output_2 = self.l2(output_1)
        output = self.l3(output_2)
        return output

model = BERTClass()
model.to(device)


# Define the loss function (CrossEntropy for multiclass classification)
def loss_fn(outputs, targets):
    # Using CrossEntropyLoss for multiclass classification (only one label per input)
    return torch.nn.CrossEntropyLoss()(outputs, targets)

# Define the optimizer (Adam)
optimizer = torch.optim.Adam(params =  model.parameters(), lr=LEARNING_RATE)


# Training loop for one epoch
def train(epoch):
    model.train()
    for _,data in enumerate(training_loader, 0):
        ids = data['ids'].to(device, dtype = torch.long)
        mask = data['mask'].to(device, dtype = torch.long)
        token_type_ids = data['token_type_ids'].to(device, dtype = torch.long)
        targets = data['targets'].to(device, dtype = torch.long)

        outputs = model(ids, mask, token_type_ids)

        optimizer.zero_grad()
# Define the loss function (CrossEntropy for multiclass classification)
        loss = loss_fn(outputs, targets)
        if _%5000==0:
            print(f'Epoch: {epoch}, Loss:  {loss.item()}')

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

for epoch in range(EPOCHS):
    train(epoch)


# Validation loop to evaluate model on unseen data
def validation(epoch):
    model.eval()
    fin_targets=[]
    fin_outputs=[]
    with torch.no_grad():
        for _, data in enumerate(testing_loader, 0):
            ids = data['ids'].to(device, dtype = torch.long)
            mask = data['mask'].to(device, dtype = torch.long)
            token_type_ids = data['token_type_ids'].to(device, dtype = torch.long)
            targets = data['targets'].to(device, dtype = torch.long)
            outputs = model(ids, mask, token_type_ids)
            fin_targets.extend(targets.cpu().detach().numpy().tolist())
            fin_outputs.extend(torch.softmax(outputs, dim=1).cpu().detach().numpy().tolist())

    return fin_outputs, fin_targets


for epoch in range(EPOCHS):
    outputs, targets = validation(epoch)
    outputs = np.argmax(outputs, axis=1)
    accuracy = metrics.accuracy_score(targets, outputs)
    f1_score_micro = metrics.f1_score(targets, outputs, average='micro')
    f1_score_macro = metrics.f1_score(targets, outputs, average='macro')
    f1_score_weighted = metrics.f1_score(targets, outputs, average='weighted')
    print(f"Accuracy Score = {accuracy}")
    print(f"F1 Score (Micro) = {f1_score_micro}")
    print(f"F1 Score (Macro) = {f1_score_macro}")
    print(f"F1 Score (Weighted) = {f1_score_weighted}")

# Save the trained model
# Save model weights and optimizer state
def save_model(model, epoch, path="model.pth"):
    state = {
        'epoch': epoch,
        'state_dict': model.state_dict(),
        'optimizer': optimizer.state_dict(),
    }
    torch.save(state, path)
    print(f"Model saved to {path}")


# Save model weights and optimizer state
save_model(model, epoch, path=f"model_epoch_{epoch+1}_{colonna_rif}.pth")


# Print the confusion matrix

import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay


# Assuming you have called the validation function and stored the results
fin_outputs, fin_targets = validation(epoch)

# Convert model output probabilities into predicted class labels
fin_outputs = np.argmax(fin_outputs, axis=1)

# Compute the confusion matrix
# Compute and plot normalized confusion matrix
cm = confusion_matrix(fin_targets, fin_outputs)

# Normalize the confusion matrix by row (true labels)
cm_normalized = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]

# Convert to percentage
cm_normalized_percentage = np.round(cm_normalized * 100,2)

# Plot the confusion matrix
plt.figure(figsize=(10, 8))
sns.heatmap(cm_normalized_percentage, annot=True, fmt='.2f', cmap='Blues', xticklabels=range(8), yticklabels=range(8))
plt.xlabel('Predicted Labels')
plt.ylabel('True Labels')
plt.title('Confusion Matrix')

# Save the figure as a PNG file
plt.savefig(f'normalized_confusion_matrix_{colonna_rif}.png')

plt.show()