import torch


def train_model(train_loader, model, optimizer, criterion, device):
    """
    Note: train_loss and train_acc is accurate only if set drop_last=False in loader

    :param train_loader: y: one_hot float tensor
    :param model:
    :param optimizer:
    :param criterion: set reduction='sum'
    :param device:
    :return:
    """
    model.train(mode=True)
    train_loss = 0
    correct = 0
    for batch_idx, (x, y) in enumerate(train_loader):
        x, y = x.to(device), y.to(device)
        global_prob = model(x)[0]
        if isinstance(criterion, torch.nn.CrossEntropyLoss):
            _, yi = y.max(dim=1)
            loss = criterion(global_prob, yi)
        else:
            loss = criterion(global_prob, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        train_loss += loss.item()
        with torch.no_grad():
            pred = global_prob.max(1, keepdim=True)[1]  # get the index of the max log-probability
            _, y_idx = y.max(dim=1)
            correct += pred.eq(y_idx.view_as(pred)).sum().item()

    train_loss /= len(train_loader.dataset)
    train_acc = correct / len(train_loader.dataset)
    return {'loss': train_loss, 'acc': train_acc}