#===============Install required packages =====================#

!pip install transformers[torch] datasets evaluate wandb

#================Import Libaries===============================#
# Data processing
import numpy as np

# Modeling

from transformers import RobertaTokenizer, RobertaForSequenceClassification, TrainingArguments, Trainer, EarlyStoppingCallback, TextClassificationPipeline

# Hugging Face Dataset
from datasets import Dataset

# Model performance evaluation
from sklearn import metrics

#for train, validate spliting
from sklearn.model_selection import train_test_split

from sklearn.preprocessing import MultiLabelBinarizer

import seaborn as sns
import matplotlib.pyplot as plt
import wandb
import tqdm

#sets the model to evaluation mode, disables certain layers like dropout and batch normalization that behave differently during training and evaluation
model.eval()

#sets the model to evaluation mode, disables certain layers like dropout and batch normalization that behave differently during training and evaluation
test_preds = []     #store the predicted probabilities
test_labels = []    #store the true labels for each sample

#disables gradient calculation during the forward pass, as we don't need gradients for inference, this save memory and speeds up the computation
with torch.no_grad():

    #iterates over the batches of data from the test data loader
    for batch in tqdm(test_loader, desc="Testing"):

	#move the input IDs tensor, attention mask tensor, labels tensor to the specified device, GPU in this case
	input_ids = batch['input_ids'].to(device)
        attention_mask = batch['attention_mask'].to(device)
        labels = batch['labels'].to(device)

	#forward pass: pass the input IDs and attention mask through the model to obtain the logits (raw outputs before applying the activation function)
        logits = model(input_ids, attention_mask)

	# apply sigmoid activation to obtain predicted probabilities, move the tensor to the CPU and convert it to a numpy array
	probs = torch.sigmoid(logits).cpu().numpy()

	#extend the test predictions and labels lists
	test_preds.extend(probs)
        test_labels.extend(labels.cpu().numpy())

#convert the test predictions and labels to numpy arrays
test_preds = np.array(test_preds)
test_labels = np.array(test_labels)

#Calculate the F1 macro score by comparing the true labels (test_labels) with the predicted labels (test_preds > 0.5). The average='macro' argument computes the F1 score for each label independently and then takes the unweighted mean
f1_macro = metrics.f1_score(test_labels, test_preds > 0.5, average='macro')
print(f"Test F1 Macro Score: {f1_macro:.4f}")

#Login with WeightandBias(WandB) to track the trainloss and validate loss during training the model
from google.colab import userdata
wandb.login(key=userdata.get('secret_wandb'))

import wandb
from torch.optim.lr_scheduler import CosineAnnealingLR

#initial wandb for experiment tracking and update the configuration with the specified hyperparameters
wandb.init(project='scopusclassification', name='run3')
wandb.config.update({
    'learning_rate': 1e-5,
    'batch_size': 16,
    'num_epochs': 20,
    'model_architecture': 'RoBERTa',
    'scheduler': 'CosineAnnealingLR'
})

#set the device to GPU if available, otherwise use CPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

#move the model to the specified device, GPU in this case
model.to(device)

#define the optimizer, loss function, and learning rate scheduler:
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-5)
criterion = nn.BCEWithLogitsLoss()
steps_per_epoch = len(train_loader)
num_epochs = 20
T_max = num_epochs * steps_per_epoch
scheduler = CosineAnnealingLR(optimizer=optimizer, T_max=T_max)

       ################## Training Loop ########################

#apply for tracking validate loss
best_valid_loss = float('inf')
best_model_state = None


#iterate over the specified number of epoch
for epoch in range(num_epochs):

    #set the model to training mode
    model.train()

    #initialize variables to track training loss and steps
    train_loss = 0.0
    train_steps = 0

    #iterate over the training data loader
    for batch in tqdm(train_loader, desc=f"Training Epoch {epoch}"):

	#move the input data to the device
        input_ids = batch['input_ids'].to(device)
        attention_mask = batch['attention_mask'].to(device)
        labels = batch['labels'].to(device)

	#zero the gradients for resetting loss gradients from previous epoch
        optimizer.zero_grad()

	#forward pass
        logits = model(input_ids, attention_mask)

	#calculate the loss on current batch
        loss = criterion(logits, labels.float())

	#backward pass to calculate loss gradient
	loss.backward()

	#update the model parameters(weights) using loss gradient
	optimizer.step()

	#accumulate the training loss and increment the step counter
        train_loss += loss.item()
        train_steps += 1

    #calculate the average training loss
    train_loss /= train_steps

    #calculate and log the average training loss for the epoch, then log the training loss to wandb
    print(f"Training Loss: {train_loss:.4f}")
    wandb.log({'train_loss': train_loss}, step=epoch)


    ################## Validation Loop ########################

    #set the model to evaluation mode
    model.eval()

    #initialize variables to track validation loss, predictions, and labels
    valid_loss = 0.0
    valid_steps = 0
    valid_preds = []
    valid_labels = []

    #disable gradient calculation (likes when inferencing)
    with torch.no_grad():

	#futher processes same as inference process describe above, but for the validation loop, we iterate over the validation data loader
        for batch in tqdm(valid_loader, desc=f"Validation Epoch {epoch}"):
            input_ids = batch['input_ids'].to(device)
            attention_mask = batch['attention_mask'].to(device)
            labels = batch['labels'].to(device)

            logits = model(input_ids, attention_mask)
            loss = criterion(logits, labels.float())

            probs = torch.sigmoid(logits).cpu().numpy()
            valid_preds.extend(probs)
            valid_labels.extend(labels.cpu().numpy())

            valid_loss += loss.item()
            valid_steps += 1

    valid_loss /= valid_steps
    print(f"Validation Loss: {valid_loss:.4f}")


    valid_preds = np.array(valid_preds)
    valid_labels = np.array(valid_labels)
    f1_macro = metrics.f1_score(valid_labels, valid_preds > 0.5, average='macro')
    print(f"Validation F1 Macro Score: {f1_macro:.4f}")
    wandb.log({'val_loss': valid_loss, 'val_f1_macro': f1_macro}, step=epoch)

    # Save the model if the validation loss improves
    if valid_loss < best_valid_loss:
        best_valid_loss = valid_loss
        best_model_state = model.state_dict()
        torch.save(best_model_state, 'best_model.pth')


    #update the learning rate based on the cosine annealing schedule
    scheduler.step()

#finish the wandb run and sync the logged data
wandb.finish()

# Load the best model state for inference
model.load_state_dict(torch.load('best_model.pth'))

import wandb
from torch.optim.lr_scheduler import CosineAnnealingLR

wandb.init(project='scopusclassification', name='run4_continue_finetuning', resume='allow')

wandb.config.update({
    'learning_rate': 1e-5,
    'batch_size': 16,
    'num_epochs': 40,  # Increase the number of epochs
    'model_architecture': 'RoBERTa',
    'scheduler': 'CosineAnnealingLR'
})

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

# Load the best model state
model.load_state_dict(torch.load('best_model.pth'))
model.to(device)

optimizer = torch.optim.AdamW(model.parameters(), lr=1e-5)
criterion = nn.BCEWithLogitsLoss()

steps_per_epoch = len(train_loader)
num_epochs = 40  # Increase the number of epochs
T_max = num_epochs * steps_per_epoch

scheduler = CosineAnnealingLR(optimizer=optimizer, T_max=T_max)

best_valid_loss = float('inf')
best_model_state = None

# Start from the last epoch
start_epoch = 20  # Adjust this if needed

for epoch in range(start_epoch, num_epochs):
    model.train()
    train_loss = 0.0
    train_steps = 0

    for batch in tqdm(train_loader, desc=f"Training Epoch {epoch}"):
        input_ids = batch['input_ids'].to(device)
        attention_mask = batch['attention_mask'].to(device)
        labels = batch['labels'].to(device)

        optimizer.zero_grad()
        logits = model(input_ids, attention_mask)
        loss = criterion(logits, labels.float())
        loss.backward()
        optimizer.step()

        train_loss += loss.item()
        train_steps += 1

    train_loss /= train_steps
    print(f"Training Loss: {train_loss:.4f}")
    wandb.log({'train_loss': train_loss}, step=epoch)

    model.eval()
    valid_loss = 0.0
    valid_steps = 0
    valid_preds = []
    valid_labels = []

    with torch.no_grad():
        for batch in tqdm(valid_loader, desc=f"Validation Epoch {epoch}"):
            input_ids = batch['input_ids'].to(device)
            attention_mask = batch['attention_mask'].to(device)
            labels = batch['labels'].to(device)

            logits = model(input_ids, attention_mask)
            loss = criterion(logits, labels.float())

            probs = torch.sigmoid(logits).cpu().numpy()
            valid_preds.extend(probs)
            valid_labels.extend(labels.cpu().numpy())

            valid_loss += loss.item()
            valid_steps += 1

    valid_loss /= valid_steps
    print(f"Validation Loss: {valid_loss:.4f}")


    valid_preds = np.array(valid_preds)
    valid_labels = np.array(valid_labels)
    f1_macro = metrics.f1_score(valid_labels, valid_preds > 0.5, average='macro')
    f1_micro = metrics.f1_score(valid_labels, valid_preds > 0.5, average='micro')
    print(f"Validation F1 Macro Score: {f1_macro:.4f}")
    print(f"Validation F1 Micro Score: {f1_micro:.4f}")
    wandb.log({'val_loss': valid_loss, 'val_f1_macro': f1_macro, 'val_f1_micro':f1_micro }, step=epoch)


    # Save the model if the validation loss improves
    if valid_loss < best_valid_loss:
        best_valid_loss = valid_loss
        best_model_state = model.state_dict()
        torch.save(best_model_state, 'best_model.pth')



    scheduler.step()

wandb.finish()

# Load the best model state for inference
model.load_state_dict(torch.load('best_model.pth'))

model.eval()
test_preds = []
test_labels = []

with torch.no_grad():
    for batch in tqdm(test_loader, desc="Testing"):
        input_ids = batch['input_ids'].to(device)
        attention_mask = batch['attention_mask'].to(device)
        labels = batch['labels'].to(device)

        logits = model(input_ids, attention_mask)
        probs = torch.sigmoid(logits).cpu().numpy()
        test_preds.extend(probs)
        test_labels.extend(labels.cpu().numpy())

test_preds = np.array(test_preds)
test_labels = np.array(test_labels)

f1_macro = metrics.f1_score(test_labels, test_preds > 0.5, average='macro')
print(f"Test F1 Macro Score: {f1_macro:.4f}")
f1_micro = metrics.f1_score(test_labels, test_preds > 0.5, average='micro')
print(f"Test F1 Micro Score: {f1_micro:.4f}")