import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors as colors

def gen_cmap(name=None, rgb=None, consecutive=True):


    def gen_consecutive_cmap_from_name(list_name):

        values = range(len(list_name))
        vmax = np.ceil(np.max(values))
        list_color = [(v/vmax, name) for v, name in zip(values, list_name)]

        return colors.LinearSegmentedColormap.from_list('custom', list_color)


    def gen_consecutive_cmap_from_rgb(list_rgb):

        imax = float(len(list_rgb)) - 1
        list_color = {'red':[], 'green':[], 'blue':[]}
        for i, rgb in enumerate(list_rgb):
            list_color['red' ].append((i/imax, rgb[0], rgb[0]))
            list_color['green'].append((i/imax, rgb[1], rgb[1]))
            list_color['blue'].append((i/imax, rgb[2], rgb[2]))

        return colors.LinearSegmentedColormap('custom', list_color)



    if consecutive:
        if name is not None:
            return gen_consecutive_cmap_from_name(name)

        elif rgb is not None:
            return gen_consecutive_cmap_from_rgb(rgb)

    else:
        if name is not None:
            return colors.ListedColormap(name)

        elif rgb is not None:
            return colors.ListedColormap(rgb)