from utils.dataloader import Dataloader
from utils.BertArchitecture import BertNER
from utils.BertArchitecture import BioBertNER
from utils.metric_tracking import MetricsTracking

import torch
from torch.optim import SGD
from torch.utils.data import DataLoader

import numpy as np
import pandas as pd

from tqdm import tqdm

def train_loop(model, train_dataset, eval_dataset, optimizer, batch_size, epochs, type, train_sampler=None, eval_sampler=None, verbose=True):
    """
    Usual training loop, including training and evaluation.

    Parameters:
    model (BertNER | BioBertNER): Model to be trained.
    train_dataset (Custom_Dataset): Dataset used for training.
    eval_dataset (Custom_Dataset): Dataset used for testing.
    optimizer (torch.optim): Optimizer used, usually SGD or Adam.
    batch_size (int): Batch size used during training.
    epochs (int): Number of epochs used for training.
    train_sampler (SubsetRandomSampler): Sampler used during hyperparameter-tuning.
    val_subsampler (SubsetRandomSampler): Sampler used during hyperparameter-tuning.
    verbose (bool): Whether the model should be evaluated after each epoch or not.

    Returns:
    tuple:
        - train_res (dict): A dictionary containing the results obtained during training.
        - test_res (dict): A dictionary containing the results obtained during testing.
    """

    if train_sampler == None or eval_sampler == None:
        train_dataloader = DataLoader(train_dataset, batch_size = batch_size, shuffle = False, sampler=train_sampler)
        eval_dataloader = DataLoader(eval_dataset, batch_size = batch_size, shuffle = False, sampler=eval_sampler)
    else:
        train_dataloader = DataLoader(train_dataset, batch_size = batch_size, shuffle = False)
        eval_dataloader = DataLoader(eval_dataset, batch_size = batch_size, shuffle = False)

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

    #training
    for epoch in range(epochs):

        train_metrics = MetricsTracking(type)

        model.train() #train mode

        for train_data in tqdm(train_dataloader):

            train_label = train_data['entity'].to(device)
            mask = train_data['attention_mask'].squeeze(1).to(device)
            input_id = train_data['input_ids'].squeeze(1).to(device)

            optimizer.zero_grad()

            output = model(input_id, mask, train_label)
            loss, logits = output.loss, output.logits
            predictions = logits.argmax(dim=-1)

            #compute metrics
            train_metrics.update(predictions, train_label, loss.item())

            loss.backward()
            optimizer.step()

        if verbose:
            model.eval() #evaluation mode

            eval_metrics = MetricsTracking(type)

            with torch.no_grad():

                for eval_data in eval_dataloader:

                    eval_label = eval_data['entity'].to(device)
                    mask = eval_data['attention_mask'].squeeze(1).to(device)
                    input_id = eval_data['input_ids'].squeeze(1).to(device)

                    output = model(input_id, mask, eval_label)
                    loss, logits = output.loss, output.logits

                    predictions = logits.argmax(dim=-1)

                    eval_metrics.update(predictions, eval_label, loss.item())

            train_results = train_metrics.return_avg_metrics(len(train_dataloader))
            eval_results = eval_metrics.return_avg_metrics(len(eval_dataloader))

            print(f"Epoch {epoch+1} of {epochs} finished!")
            print(f"TRAIN\nMetrics {train_results}\n")
            print(f"VALIDATION\nMetrics {eval_results}\n")

    if not verbose:
        model.eval() #evaluation mode

        eval_metrics = MetricsTracking(type)

        with torch.no_grad():

            for eval_data in eval_dataloader:

                eval_label = eval_data['entity'].to(device)
                mask = eval_data['attention_mask'].squeeze(1).to(device)
                input_id = eval_data['input_ids'].squeeze(1).to(device)

                output = model(input_id, mask, eval_label)
                loss, logits = output.loss, output.logits

                predictions = logits.argmax(dim=-1)

                eval_metrics.update(predictions, eval_label, loss.item())

        train_results = train_metrics.return_avg_metrics(len(train_dataloader))
        eval_results = eval_metrics.return_avg_metrics(len(eval_dataloader))

        print(f"Epoch {epoch+1} of {epochs} finished!")
        print(f"TRAIN\nMetrics {train_results}\n")
        print(f"VALIDATION\nMetrics {eval_results}\n")

    return train_results, eval_results

def testing(model, test_dataset, batch_size, type):
    """
    Function for testing a trained model.

    Parameters:
    model (BertNER | BioBertNER): Model to be tested
    train_dataset (Custom_Dataset): Dataset used for testing
    batch_size (int): Batch size used during training.

    Returns:
    tuple:
        - test_res (dict): A dictionary containing the results obtained during testing.
    """

    test_dataloader = DataLoader(test_dataset, batch_size = batch_size, shuffle = False)

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

    model.eval() #evaluation mode

    test_metrics = MetricsTracking(type)

    with torch.no_grad():

        for test_data in test_dataloader:

            test_label = test_data['entity'].to(device)
            mask = test_data['attention_mask'].squeeze(1).to(device)
            input_id = test_data['input_ids'].squeeze(1).to(device)

            output = model(input_id, mask, test_label)
            loss, logits = output.loss, output.logits

            predictions = logits.argmax(dim=-1)

            test_metrics.update(predictions, test_label, loss.item())

        test_results = test_metrics.return_avg_metrics(len(test_dataloader))

        print(f"TEST\nMetrics {test_results}\n")

    return test_results