import time
import copy
import torch


def train_model(nb_epoch, dataloaders, model, opt, criterion, device,
                scheduler=None):
    """
    Function to train a model.
    Args:
        nb_epoch: number of epoch to train
        train_dsld: dataset loader for training
        model: model to optimize
        opt: optimizer
        criterion: metric to calculate the loss
        scheduler: learning rate scheduler
        device: device generated by torch, referring to CPU or GPU

    Returns:
        model: a trained model

    """
    st_start = time.time()
    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0

    nb_data = {}
    for phase in dataloaders:
        nb_data[phase] = dataloaders[phase].batch_size * len(dataloaders[phase])
    for epoch in range(nb_epoch):
        print("Epcoh {}/{}".format(epoch, nb_epoch))
        print("_" * 10)
        for phase in ["train", "val"]:
            if phase == "train":
                model.train()
            else:
                model.eval()

            running_loss = 0.0
            running_corrects = 0

            for inputs, labels in dataloaders[phase]:
                inputs = inputs.to(device)
                labels = labels.to(device)

                opt.zero_grad()
                with torch.set_grad_enabled(phase == "train"):
                    pred_y = model(inputs)
                    _, preds = torch.max(pred_y, 1)
                    loss = criterion(pred_y, labels)
                    if phase == "train":
                        loss.backward()
                        opt.step()
                running_loss += loss.item() * inputs.size(0)
                running_corrects += torch.sum(preds == labels.data)
            if phase == "train" and scheduler is not None:
                scheduler.step()

            epoch_loss = running_loss / nb_data[phase]
            epoch_acc = float(running_corrects) / nb_data[phase]

            print("{} Loss: {:.4f} Acc: {:.4f}".format(
                phase, epoch_loss, epoch_acc))

            if phase == "val" and epoch_acc > best_acc:
                best_acc = epoch_acc
                best_model_wts = copy.deepcopy(model.state_dict())

        print()
    time_elapsed = time.time() - st_start
    print("Training complete in {:.0f}m {:.0f}s".format(
        time_elapsed // 60, time_elapsed % 60))
    model.load_state_dict(best_model_wts)

    return model