import matplotlib.pyplot as plt
from generate_points import get_random_point
import numpy as np


def get_color(i):
    return plt.get_cmap('tab20')(i)


def calc_length(a, b):
    '''Calculate Euclidian distance between points'''
    assert len(a) == len(b)
    return np.square(np.asarray(b)-np.asarray(a)).sum()


def plot_data(data):
    lst_x, lst_y = zip(*data)
    lst_x = list(lst_x)
    lst_y = list(lst_y)
    plt.figure(1)
    ax = plt.axes()
    ax.scatter(lst_x, lst_y)
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    plt.grid(True)
    plt.show()


def plot_error_data(error_data, fname=None):
    fig, ax = plt.subplots()
    ax.set_xlabel('k')
    ax.set_ylabel('err')
    ax.set_xlim(2, 20)
    plt.title('Errors')
    plt.grid(True)

    lst_x, lst_y = zip(*error_data)
    lst_x = list(lst_x)
    lst_y = list(lst_y)
    ax.plot(lst_x, lst_y, 'ro-')

    if fname:
        plt.savefig(fname)
    else:
        plt.show()


def get_data1():
    data = []
    for _ in range(200):
        data.append(get_random_point((0, 0), 1))
    return data


def get_data2():
    data = []
    for i in range(2):
        for _ in range(100):
            data.append(get_random_point((3 * ((-1) ** i), 0), 0.5))
    return data