import torch


def train(model, device, epochs, criterion, optimizer, trainloader, validationloader):
    """
        This function Trains the model with the given parameters.
        Arguments:
            model - The Pre-trained model for transfer learning
            device - The device for training can be cuda or cpu
            epochs - Number of cycles for training
            criterion - Measure of training Loss
            optimizer - Training optimizer for updating weights
            trainloader - processed dataloader images for training
            validationloader - processed dataloader images for validation
        Returns:
            model - Trained model
    """
    model.to(device)
    for i in range(epochs):
        running_loss = 0
        for images, labels in trainloader:
            images, labels = images.to(device), labels.to(device)

            optimizer.zero_grad()

            output = model.forward(images)

            loss = criterion(output, labels)

            loss.backward()

            optimizer.step()

            running_loss += loss.item()

        else:
            test_losses = 0
            accuracy = 0
            model.eval()
            with torch.no_grad():
                for images, labels in validationloader:
                    images, labels = images.to(device), labels.to(device)
                    output = model.forward(images)

                    test_losses += criterion(output, labels).item()

                    ps = torch.exp(output)
                    top_p, top_class = ps.topk(1, dim=1)
                    equals = top_class == labels.view(*top_class.shape)

                    accuracy += torch.mean(equals.type(torch.FloatTensor)).item()
            model.train()

            print(f'Training Loss: {running_loss / len(trainloader):.3f} '
                  f'Validation Loss: {test_losses / len(validationloader):.3f} '
                  f'Validation Accuracy: {(accuracy / len(validationloader)) * 100:.3f}%')
    return model