import math


def sched_dsd(start, end, pos, middle=None):
    """
    Pruning schedule inspired on Dense-Sparse-Dense: prunes to middle % of the
    features, then unprunes to end %, following a cosine-shaped pattern.
    DSD original paper (Han et al, 2016): https://arxiv.org/pdf/1607.04381.pdf
    Implementation based on the Fasterai tutorial at:
    https://nathanhubens.github.io/fasterai/schedules.html#Dense-Sparse-Dense

    Args:
        start (int or float) - percentage corresponding to the network's
            initial sparsity.
        end (int or float) - percentage corresponding to the network's
            final sparsity.
        pos (float) - relative position within the pruning schedule.
        middle (int or float or None) - percentage corresponding to the
            network's sparsity at the middle of the pruning (should be lower
            than both start and end), if None (default) corresponds to halfway
            between end and 100%.
    """
    if middle is None:
        middle = end + (100-end)/2
    if pos < 0.5:
        return start + (1 + math.cos(math.pi*(1-pos*2))) * (middle-start) / 2
    else:
        return middle + (1 - math.cos(math.pi*(1-pos*2))) * (end-middle) / 2