import torch


def train_loop(dataloader, model, loss_fn, optimizer, device):
    """
    training loop
    Arguments:
        dataloader: pytorch dataloader for training
        model: pytorch model
        loss_fn: the loss function to train the model
        optimizer: the optimizer to train the model
        device: pytorch device (cuda or cpu)
    Returns:
        mean_loss(float): the mean loss value
        correct(float): the mean accuracy
    """
    model.train()
    size = len(dataloader.dataset)
    total_loss = 0
    correct = 0
    n_iter = len(dataloader)
    is_cuda = device.type.find('cuda')!=-1
    for batch, (X, y, _) in enumerate(dataloader):
        if is_cuda:
            X = X.to(device)
            y = y.to(device)
        # Compute prediction and loss
        pred = model(X)
        loss = loss_fn(pred, y)

        # get the accuracy
        with torch.no_grad():
            correct += (pred.softmax(1).argmax(1) == y).type(torch.float).sum().item()
            total_loss += loss.item()

        # Backpropagation
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if batch % 2 == 0:
            loss, current = loss.item(), batch * dataloader.batch_size
            print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")
    mean_loss = total_loss/n_iter
    correct /= size
    print(f"Train Stats: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {mean_loss:>8f}")
    return mean_loss,correct