import time

import torch

from torchmetrics import Accuracy

def train_model(model: torch.nn.Module,

               data_loader: torch.utils.data.DataLoader, 

               loss_fn: torch.nn.Module, #criterion

               optimizer: torch.optim.Optimizer,

               device: torch.device,

               num_epochs,

               output_shape):

    start_time = time.time()

    accuracy_metric = Accuracy(num_classes= output_shape, task='multiclass').to(device)

    for epoch in range(num_epochs):

        print(f'Epoch {epoch + 1}/{num_epochs}')

        model.train() 

        train_loss = 0

        for signal, class_label in data_loader: 

            signal, class_label = signal.to(device), class_label.to(device) #

            train_pred = model(signal)

            loss = loss_fn(train_pred, class_label)

            train_loss += loss.item() 

            accuracy_metric(train_pred, class_label)

            optimizer.zero_grad()

            loss.backward()

            optimizer.step()

        train_acc = accuracy_metric.compute() * 100  

        print(f"Train loss: {train_loss / len(data_loader):.5f} | Train accuracy: {train_acc:.2f}%")

        accuracy_metric.reset()

    total_time = (time.time() - start_time)

    print(f"\nTotal training time: {total_time} seconds")

    return total_time 

def evaluate_model(model: torch.nn.Module,

                   test_loader: torch.utils.data.DataLoader, 

                   loss_fn: torch.nn.Module, #criterion

                   device: torch.device,

                   output_shape):

    start_time = time.time()

    test_loss = 0

    accuracy_metric = Accuracy(num_classes=output_shape, task='multiclass').to(device)

    model.eval()

    with torch.inference_mode(): 

        for signal, class_label in test_loader:

            signal, class_label = signal.to(device), class_label.to(device)

            test_pred = model(signal)

            loss = loss_fn(test_pred, class_label)

            test_loss +=loss.item()

            accuracy_metric(test_pred, class_label)

    test_acc = accuracy_metric.compute() * 100  

    print(f"\nTest loss: {test_loss/len(test_loader):.5f} | Test accuracy: {test_acc:.2f}%")

    accuracy_metric.reset() 

    total_time = (time.time() - start_time)

    print(f"Total evaluation time: {total_time} seconds\n")

    return test_acc.item(), total_time