# By Jake VanderPlas
# License: BSD-style

import matplotlib.pyplot as plt
import numpy as np
from epistasis.utils import Bunch

def discrete_cmap(N, base_cmap=None):
    """Create an N-bin discrete colormap from the specified input map.

    Borrowed from Jake VanderPlas's gist found here:
    https://gist.github.com/jakevdp/91077b0cae40f8f8244a
    """
    # Note that if base_cmap is a string or None, you can simply do
    #    return plt.cm.get_cmap(base_cmap, N)
    # The following works for string, None, or a colormap instance:
    base = plt.cm.get_cmap(base_cmap)
    color_list = base(np.linspace(0, 1, N))
    cmap_name = base.name + str(N)
    return base.from_list(cmap_name, color_list, N)
