# -*- coding: utf-8 -*-
import numpy as np
import pylab as plt

def sine_wave_triplets(n=8, phi=0, saturation=1.00, start=0.0, stop=1.0):
    x = np.linspace(start, stop, n)
    angle = 2 * np.pi * (x-start)/(stop - start) + phi
    y = 0.5 * saturation * np.sin(angle) + 0.5
    return np.array([x,y,y]).transpose()

def saturated_tripplets(n_sin=5, phi=0, n_sat=1, rel_space=0.0):
    sat_list = np.linspace(0, 1, n_sat)
    sin_list = [sine_wave_triplets(n=n_sin,
                                    phi=phi,
                                    saturation=sat,
                                    start=float(i)/len(sat_list), 
                                    stop=float(i + 1-rel_space)/len(sat_list))
                for i, sat in enumerate(sat_list)]
    if rel_space > 0:
        sin_list.append(np.array([1, 0.5, 0.5]))
    return np.vstack(sin_list)
    
def cmap_color_circle(npoints=10):
    npoints = 5
    cdict = {'red': sine_wave_triplets(n=npoints, phi=0),
             'green': sine_wave_triplets(n=npoints, phi=2*np.pi/3.),
             'blue': sine_wave_triplets(n=npoints, phi=4*np.pi/3.)}
    my_cmap = plt.matplotlib.colors.LinearSegmentedColormap('my_colormap',cdict,256)
    my_cmap.set_under('black')
    my_cmap.set_over('white')
    return my_cmap

def cmap_color_circle_saturation(n_sat=16, n_sin=5, rel_space=0.0):
    cdict = {'red': saturated_tripplets(n_sin, phi=0, n_sat=n_sat, rel_space=rel_space),
             'green': saturated_tripplets(n_sin, phi=2*np.pi/3., n_sat=n_sat, rel_space=rel_space),
             'blue': saturated_tripplets(n_sin, phi=4*np.pi/3., n_sat=n_sat, rel_space=rel_space),}
    n_colors = n_sat*n_sin
    if n_colors < 256:
        n_colors = n_sat*n_sin*(np.floor(256/n_colors)+1)
    my_cmap = plt.matplotlib.colors.LinearSegmentedColormap('circlesat', cdict, n_colors)
    my_cmap.set_under('black')
    my_cmap.set_over('white')
    return my_cmap

def calc_angle(x, y, center):
    dx = x - center[0]
    dy = y - center[1]
    return np.arctan2(dy, dx)    

def point_point_distance(px, py, point):
    dx = px - point[0]
    dy = py - point[1]
    dist = np.sqrt(dx*dx + dy*dy)
    return dist
    
def show_color_circle(cmap=plt.cm.gray, nx=100, ny=100):
    y, x, = np.mgrid[:ny, :nx]
    center = [0.5*nx, 0.5*ny]
    data = calc_angle(x, y, center) + np.pi
    dist = point_point_distance(x, y, center)
    rad1 = 0.3*nx
    rad2 = 0.4*nx
    data[dist < rad1] = -1
    data[dist > rad2] = -1
    data[0.4*nx:0.6*nx, 0.4*ny:0.6*ny] = 10 # square of points above range --> set_over color 
    
    plt.subplot(1, 2, 1)
    norm = plt.matplotlib.colors.Normalize(vmin=0*np.pi, vmax=2*np.pi)
    im = plt.imshow(data, interpolation='nearest', cmap=cmap)   
    im.set_norm(norm)
    plt.colorbar(shrink=0.5)
    
    n_sat = 50
    n_sin = 200
    rel_space=0.1
    sgl_sat_range = 10
    angle_range = (1-rel_space)*sgl_sat_range

    angles = calc_angle(x, y, center) 
    distances = point_point_distance(x, y, center)
    distances[distances > 0.4*nx] = 0
    distances -= 0.2*nx
    middle_indices = distances < 0
    distances[middle_indices] = 0
    
    distances *= (n_sat-1) / np.max(distances.flat)
    distances = np.round(distances)
    
    angles_normed = (angles + np.pi)/(2*np.pi)
    
    combined = angle_range*angles_normed + sgl_sat_range*distances
    combined[middle_indices] = -1
    cm = cmap_color_circle_saturation(n_sat, n_sin, rel_space=rel_space)
    
    plt.subplot(1, 2, 2)
    plt.imshow(combined, cmap=cm)
    plt.clim(0, sgl_sat_range*n_sat)
    plt.colorbar(shrink=0.5)     
    plt.show()    

def main():   
    my_cmap = cmap_color_circle()
    show_color_circle(my_cmap, nx=500, ny=500)
    

if __name__=="__main__":
    main()