import math

def batch_sizes(n, batch_size, drop_last=False):
    if n <= 0 or batch_size <= 0:
        raise ValueError('positive sample count and batch size required')
    sizes = [min(batch_size, n-start) for start in range(0, n, batch_size)]
    return [s for s in sizes if not drop_last or s == batch_size]

def count(n, b, epochs, accumulation=1, drop_last=False):
    if epochs <= 0 or accumulation <= 0:
        raise ValueError('positive epochs and accumulation required')
    sizes = batch_sizes(n, b, drop_last)
    # Flush each incomplete accumulation group at every epoch end.
    updates_per_epoch = math.ceil(len(sizes)/accumulation)
    return {'sizes': sizes, 'microsteps': len(sizes)*epochs,
            'optimizer_updates': updates_per_epoch*epochs,
            'sample_visits': sum(sizes)*epochs}

normal = count(103, 32, 3)
dropped = count(103, 32, 3, drop_last=True)
accumulated = count(103, 32, 3, accumulation=2)
print('normal', normal)
print('drop_last', dropped)
print('accumulate 2', accumulated)
assert normal == {'sizes':[32, 32, 32, 7], 'microsteps':12,
                  'optimizer_updates':12, 'sample_visits':309}
assert dropped['optimizer_updates'] == 9 and dropped['sample_visits'] == 288
assert accumulated['optimizer_updates'] == 6 and accumulated['microsteps'] == 12
assert batch_sizes(3, 32, True) == []