import matplotlib.pyplot as plt
# %matplotlib inline
import numpy as np
from tqdm import tqdm
from keras.utils import to_categorical

# On some implementations of matplotlib, you may need to change this value
IMAGE_SIZE = 72


def generate_a_drawing(figsize, U, V, noise=0.0):
    fig = plt.figure(figsize=(figsize, figsize))
    ax = plt.subplot(111)
    plt.axis('Off')
    ax.set_xlim(0, figsize)
    ax.set_ylim(0, figsize)
    ax.fill(U, V, "k")
    fig.canvas.draw()
    imdata = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)[::3].astype(np.float32)
    imdata = imdata + noise * np.random.random(imdata.size)
    plt.close(fig)
    return imdata


def generate_a_rectangle(noise=0.0, free_location=False):
    figsize = 1.0
    U = np.zeros(4)
    V = np.zeros(4)
    if free_location:
        corners = np.random.random(4)
        top = max(corners[0], corners[1])
        bottom = min(corners[0], corners[1])
        left = min(corners[2], corners[3])
        right = max(corners[2], corners[3])
    else:
        side = (0.3 + 0.7 * np.random.random()) * figsize
        top = figsize / 2 + side / 2
        bottom = figsize / 2 - side / 2
        left = bottom
        right = top
    U[0] = U[1] = top
    U[2] = U[3] = bottom
    V[0] = V[3] = left
    V[1] = V[2] = right
    return generate_a_drawing(figsize, U, V, noise)


def generate_a_disk(noise=0.0, free_location=False):
    figsize = 1.0
    if free_location:
        center = np.random.random(2)
    else:
        center = (figsize / 2, figsize / 2)
    radius = (0.3 + 0.7 * np.random.random()) * figsize / 2
    N = 50
    U = np.zeros(N)
    V = np.zeros(N)
    i = 0
    for t in np.linspace(0, 2 * np.pi, N):
        U[i] = center[0] + np.cos(t) * radius
        V[i] = center[1] + np.sin(t) * radius
        i = i + 1
    return generate_a_drawing(figsize, U, V, noise)


def generate_a_triangle(noise=0.0, free_location=False):
    figsize = 1.0
    if free_location:
        U = np.random.random(3)
        V = np.random.random(3)
    else:
        size = (0.3 + 0.7 * np.random.random()) * figsize / 2
        middle = figsize / 2
        U = (middle, middle + size, middle - size)
        V = (middle + size, middle - size, middle - size)
    imdata = generate_a_drawing(figsize, U, V, noise)
    return [imdata, [U[0], V[0], U[1], V[1], U[2], V[2]]]


# im = generate_a_rectangle(10, True)
# plt.imshow(im.reshape(IMAGE_SIZE, IMAGE_SIZE), cmap='gray')
#
# im = generate_a_disk(10)
# plt.imshow(im.reshape(IMAGE_SIZE, IMAGE_SIZE), cmap='gray')
#
# [im, v] = generate_a_triangle(20, False)
# plt.imshow(im.reshape(IMAGE_SIZE, IMAGE_SIZE), cmap='gray')


def generate_dataset_classification(nb_samples, noise=0.0, free_location=False):
    # Getting im_size:
    im_size = generate_a_rectangle().shape[0]
    X = np.zeros([nb_samples, im_size])
    Y = np.zeros(nb_samples)
    print('Creating data:')
    for i in tqdm(range(nb_samples)):
        category = np.random.randint(3)
        if category == 0:
            X[i] = generate_a_rectangle(noise, free_location)
        elif category == 1:
            X[i] = generate_a_disk(noise, free_location)
        else:
            [X[i], V] = generate_a_triangle(noise, free_location)
        Y[i] = category
    X = (X + noise) / (255 + 2 * noise)
    return [X, Y]


def generate_test_set_classification():
    np.random.seed(42)
    [X_test, Y_test] = generate_dataset_classification(300, 20, True)
    Y_test = to_categorical(Y_test, 3)
    return [X_test, Y_test]


def generate_dataset_regression(nb_samples, noise=0.0):
    # Getting im_size:
    im_size = generate_a_triangle()[0].shape[0]
    X = np.zeros([nb_samples, im_size])
    Y = np.zeros([nb_samples, 6])
    print('Creating data:')
    for i in tqdm(range(nb_samples)):
        [X[i], Y[i]] = generate_a_triangle(noise, True)
    X = (X + noise) / (255 + 2 * noise)
    return [X, Y]


import matplotlib.patches as patches


def visualize_prediction(x, y):
    fig, ax = plt.subplots(figsize=(5, 5))
    I_ = x.reshape((IMAGE_SIZE, IMAGE_SIZE))
    ax.imshow(I_, extent=[-0.15, 1.15, -0.15, 1.15], cmap='gray')
    ax.set_xlim([0, 1])
    ax.set_ylim([0, 1])

    xy = y.reshape(3, 2)
    tri = patches.Polygon(xy, closed=True, fill=False, edgecolor='r', linewidth=5, alpha=0.5)
    ax.add_patch(tri)

    plt.show()


def generate_test_set_regression():
    np.random.seed(42)
    [X_test, Y_test] = generate_dataset_regression(300, 20)
    return [X_test, Y_test]


def plotLossHistory(loss_history, metric):
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(20, 7))
    epochs = range(1, len(loss_history.history['loss']) + 1)
    color1 = 'tomato'
    color2 = 'purple'
    fontsize = 14
    ax1.set_xlabel('epochs', fontsize=fontsize)
    ax1.set_ylabel('loss', fontsize=fontsize)
    ax1.plot(epochs, loss_history.history['loss'], '--o', color=color1)
    ax1.plot(epochs, loss_history.history['val_loss'], '-o', color=color2)
    ax1.tick_params(axis='x', labelsize=fontsize)
    ax1.tick_params(axis='y', labelsize=fontsize)
    ax1.set_xticks(epochs)
    ax1.grid(alpha=0.3)
    ax1.set_title("Loss history")

    ax2.set_xlabel('epochs', fontsize=fontsize)
    ax2.set_ylabel('metric', fontsize=fontsize)
    ax2.plot(epochs, loss_history.history[metric], '--o', color=color1, label="training")
    ax2.plot(epochs, loss_history.history['val_' + metric], '-o', color=color2, label="validation")
    ax2.tick_params(axis='x', labelsize=fontsize)
    ax2.tick_params(axis='y', labelsize=fontsize)
    ax2.set_xticks(epochs)
    ax2.grid(alpha=0.3)
    ax2.set_title(metric)
    plt.legend(loc='best')
    plt.tight_layout()
    return fig
