import math


def find_lr(model, optimizer, criterion, trn_loader, device, init_value=1e-8, final_value=10., beta=0.98):
    """
    Basic learning rate finder implemented in fastai
    quoted from https://sgugger.github.io/how-do-you-find-a-good-learning-rate.html
    :param model: model to be trained
    :param optimizer: optimizer to be used
    :param criterion: loss function
    :param trn_loader: training loader
    :param device: 'cpu' or 'cuda'
    :param init_value: the initial value for the learning rate
    :param final_value: the final value for the learning rate
    :param beta: a weighted parameter
    :return: log10 of the lrs and the corresponding loss,
             good to pick the lr an order of magnitude smaller than the best
    """
    num = len(trn_loader) - 1
    mult = (final_value / init_value) ** (1 / num)
    lr = init_value
    optimizer.param_groups[0]['lr'] = lr
    avg_loss = 0.
    best_loss = 0.
    batch_num = 0
    losses = []
    log_lrs = []
    for data in trn_loader:
        batch_num += 1
        # As before, get the loss for this mini-batch of inputs/outputs
        inputs, labels = data
        inputs = inputs.to(device)
        labels = labels.to(device)
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        # Compute the smoothed loss
        avg_loss = beta * avg_loss + (1 - beta) * loss.data
        smoothed_loss = avg_loss / (1 - beta ** batch_num)
        # Stop if the loss is exploding
        if batch_num > 1 and smoothed_loss > 4 * best_loss:
            return log_lrs, losses
        # Record the best loss
        if smoothed_loss < best_loss or batch_num == 1:
            best_loss = smoothed_loss
        # Store the values
        losses.append(smoothed_loss)
        log_lrs.append(math.log10(lr))
        # Do the SGD step
        loss.backward()
        optimizer.step()
        # Update the lr for the next step
        lr *= mult
        optimizer.param_groups[0]['lr'] = lr
    return log_lrs, losses